Spaces:
Sleeping
Sleeping
File size: 3,153 Bytes
931223d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 | import os
import abc
from datetime import datetime
class StorageProvider(abc.ABC):
@abc.abstractmethod
def save_file(self, file_bytes: bytes, filename: str) -> str:
pass
@abc.abstractmethod
def get_file(self, filename: str) -> bytes:
pass
class LocalStorage(StorageProvider):
def __init__(self, base_dir="data_storage"):
self.base_dir = base_dir
os.makedirs(base_dir, exist_ok=True)
def save_file(self, file_bytes: bytes, filename: str) -> str:
# Prepend timestamp to avoid collision
ts = datetime.now().strftime("%Y%m%d_%H%M%S")
safe_name = f"{ts}_{filename}"
path = os.path.join(self.base_dir, safe_name)
with open(path, "wb") as f:
f.write(file_bytes)
return path
def get_file(self, filename: str) -> bytes:
path = os.path.join(self.base_dir, filename)
if not os.path.exists(path):
return None
with open(path, "rb") as f:
return f.read()
class SwiftStorage(StorageProvider):
"""
OpenStack Swift Storage Provider.
Requires python-swiftclient installed.
"""
def __init__(self, auth_url, username, password, project_name, container_name="elephmind_images"):
# Import here to avoid error on Windows if not installed
try:
from swiftclient import Connection
except ImportError:
raise ImportError("python-swiftclient not installed!")
self.container_name = container_name
self.conn = Connection(
authurl=auth_url,
user=username,
key=password,
tenant_name=project_name,
auth_version='3',
os_options={'user_domain_name': 'Default', 'project_domain_name': 'Default'}
)
# Ensure container exists
try:
self.conn.put_container(self.container_name)
except Exception as e:
print(f"Swift Connection Error: {e}")
def save_file(self, file_bytes: bytes, filename: str) -> str:
ts = datetime.now().strftime("%Y%m%d_%H%M%S")
safe_name = f"{ts}_{filename}"
self.conn.put_object(
self.container_name,
safe_name,
contents=file_bytes,
content_type='application/octet-stream'
)
return f"swift://{self.container_name}/{safe_name}"
def get_file(self, filename: str) -> bytes:
# filename could be safe_name
# logic to extract key if needed
try:
_, obj = self.conn.get_object(self.container_name, filename)
return obj
except Exception:
return None
# Factory
def get_storage_provider(config_mode="LOCAL"):
if config_mode == "OPENSTACK":
return SwiftStorage(
auth_url=os.getenv("OS_AUTH_URL"),
username=os.getenv("OS_USERNAME"),
password=os.getenv("OS_PASSWORD"),
project_name=os.getenv("OS_PROJECT_NAME")
)
else:
return LocalStorage()
|