306 lines
12 KiB
Python
306 lines
12 KiB
Python
import logging
|
|
import io
|
|
import time
|
|
from typing import Any, Optional
|
|
from app.core.settings import settings
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
# get_storage_client() was being called fresh on every single file operation
|
|
# (every upload, download, stream, copy, ...), each paying for a DB round-trip
|
|
# (TenantStorageConfig, sometimes SystemConfiguration too) plus a credential
|
|
# decrypt and a new boto3 client construction — real, measured cost on a
|
|
# remote DB, and a direct contributor to "opening any document takes a long
|
|
# time." The (client, quarantine_bucket, clean_bucket) tuple is safe to reuse
|
|
# across calls for the same tenant, so it's cached in-process with a short
|
|
# TTL — long enough to eliminate the redundant work within a single document
|
|
# open (which can call this several times), short enough that an admin
|
|
# changing storage credentials in System Settings takes effect within a few
|
|
# minutes without needing a server restart.
|
|
_STORAGE_CLIENT_CACHE_TTL = 300 # seconds
|
|
_storage_client_cache: dict = {}
|
|
|
|
|
|
def invalidate_storage_client_cache(tenant_id=None) -> None:
|
|
"""Drop cached storage client(s) so the next call rebuilds from the
|
|
database — call this right after a tenant's storage config is
|
|
created/updated/cleared, so admin credential changes take effect
|
|
immediately instead of waiting out the TTL. Matches by string form
|
|
since callers pass tenant_id as either a UUID or a str.
|
|
"""
|
|
if tenant_id is None:
|
|
_storage_client_cache.clear()
|
|
return
|
|
for key in [k for k in _storage_client_cache if str(k[0]) == str(tenant_id)]:
|
|
_storage_client_cache.pop(key, None)
|
|
|
|
class B2StorageClient:
|
|
"""
|
|
Backblaze B2 Storage Client using Boto3 (S3-compatible API).
|
|
"""
|
|
|
|
def __init__(self, key_id: str, application_key: str, endpoint_url: str):
|
|
import boto3
|
|
from botocore.config import Config
|
|
|
|
region = "us-west-004"
|
|
if "s3." in endpoint_url:
|
|
parts = endpoint_url.split("s3.")
|
|
if len(parts) > 1:
|
|
region = parts[1].split(".")[0]
|
|
|
|
try:
|
|
self.client = boto3.client(
|
|
's3',
|
|
endpoint_url=endpoint_url,
|
|
aws_access_key_id=key_id,
|
|
aws_secret_access_key=application_key,
|
|
region_name=region,
|
|
config=Config(signature_version='s3v4')
|
|
)
|
|
except Exception as e:
|
|
logger.error(f"Failed to initialize B2 client: {e}")
|
|
raise
|
|
|
|
def put_object(self, Bucket: str, Key: str, Body: Any, ContentType: str = "application/octet-stream", Metadata: Optional[dict] = None, length: Optional[int] = None):
|
|
try:
|
|
if isinstance(Body, bytes):
|
|
data = Body
|
|
elif hasattr(Body, "read"):
|
|
data = Body
|
|
else:
|
|
data = Body
|
|
|
|
if Metadata:
|
|
Metadata = {str(k): str(v) for k, v in Metadata.items()}
|
|
else:
|
|
Metadata = {}
|
|
|
|
self.client.put_object(
|
|
Bucket=Bucket,
|
|
Key=Key,
|
|
Body=data,
|
|
ContentType=ContentType,
|
|
Metadata=Metadata
|
|
)
|
|
return {"status": "success", "key": Key}
|
|
except Exception as e:
|
|
logger.error(f"B2 Put Error: {e}")
|
|
raise
|
|
|
|
def get_object(self, Bucket: str, Key: str, Range: Optional[str] = None):
|
|
try:
|
|
kwargs = {"Bucket": Bucket, "Key": Key}
|
|
if Range:
|
|
kwargs["Range"] = Range
|
|
response = self.client.get_object(**kwargs)
|
|
return {
|
|
"Body": response['Body'],
|
|
"ContentType": response.get('ContentType'),
|
|
"ContentLength": response.get('ContentLength'),
|
|
"ContentRange": response.get('ContentRange'),
|
|
"StatusCode": 206 if Range else 200,
|
|
}
|
|
except Exception as e:
|
|
logger.error(f"B2 Get Error: {e}")
|
|
raise
|
|
|
|
def delete_object(self, Bucket: str, Key: str):
|
|
try:
|
|
self.client.delete_object(Bucket=Bucket, Key=Key)
|
|
return {"status": "success"}
|
|
except Exception as e:
|
|
logger.error(f"B2 Delete Error: {e}")
|
|
raise
|
|
|
|
def copy_object(self, Bucket: str, Key: str, CopySource: dict):
|
|
try:
|
|
self.client.copy_object(
|
|
Bucket=Bucket,
|
|
Key=Key,
|
|
CopySource=CopySource
|
|
)
|
|
return {"status": "success"}
|
|
except Exception as e:
|
|
logger.error(f"B2 Copy Error: {e}")
|
|
raise
|
|
|
|
def generate_presigned_url(self, ClientMethod: str, Params: Optional[dict] = None, ExpiresIn: int = 3600, user_id: Optional[int] = None, tenant_id: Optional[str] = None):
|
|
try:
|
|
params = Params or {}
|
|
key = params.get("Key")
|
|
|
|
if user_id is not None and key:
|
|
allowed_prefixes = [
|
|
f"user_{user_id}/", f"users/{user_id}/",
|
|
f"chatbot_documents/{user_id}/", f"doc_conversion/{user_id}/",
|
|
]
|
|
if tenant_id:
|
|
allowed_prefixes.extend([
|
|
f"tenants/{tenant_id}/user_{user_id}/", f"tenants/{tenant_id}/users/{user_id}/",
|
|
f"tenants/{tenant_id}/chatbot_documents/{user_id}/", f"tenants/{tenant_id}/doc_conversion/{user_id}/",
|
|
])
|
|
if not any(key.startswith(p) for p in allowed_prefixes) and not any(f"/{p}" in f"/{key}" for p in allowed_prefixes):
|
|
logger.warning(f"🚨 Security Alert: User {user_id} attempted to generate presigned URL for unauthorized key: {key}")
|
|
raise PermissionError(f"Access denied: User {user_id} does not own resource '{key}'")
|
|
|
|
if ClientMethod in ["get_object", "put_object"]:
|
|
return self.client.generate_presigned_url(
|
|
ClientMethod=ClientMethod,
|
|
Params=Params,
|
|
ExpiresIn=ExpiresIn
|
|
)
|
|
return None
|
|
except Exception as e:
|
|
if not isinstance(e, PermissionError):
|
|
logger.error(f"B2 Presigned URL Error: {e}")
|
|
raise
|
|
|
|
def list_objects(self, Bucket: str, Prefix: str = ""):
|
|
try:
|
|
response = self.client.list_objects_v2(Bucket=Bucket, Prefix=Prefix)
|
|
objects = response.get('Contents', [])
|
|
return [{"Key": obj['Key'], "Size": obj['Size']} for obj in objects]
|
|
except Exception as e:
|
|
logger.error(f"B2 List Error: {e}")
|
|
raise
|
|
|
|
|
|
def format_bytes(size: float) -> str:
|
|
"""Format bytes to human readable string."""
|
|
for unit in ["B", "KB", "MB", "GB", "TB"]:
|
|
if size < 1024.0:
|
|
return f"{size:.2f} {unit}"
|
|
size /= 1024.0
|
|
return f"{size:.2f} PB"
|
|
|
|
|
|
def generate_storage_key(tenant_id: str, user_id: int, filename: str) -> str:
|
|
"""Generate a unique original storage key for an initial upload."""
|
|
import uuid
|
|
import os
|
|
|
|
tenant_part = f"tenants/{tenant_id}/" if tenant_id and tenant_id != 'None' else ""
|
|
return f"{tenant_part}users/{user_id}/uploads/{uuid.uuid4()}_{filename}"
|
|
|
|
|
|
def _derive_b2_endpoint(key_id: str) -> str:
|
|
"""
|
|
Derive the B2 S3 endpoint URL from the key ID.
|
|
B2 key IDs start with a account cluster prefix.
|
|
"""
|
|
region_map = {
|
|
"000": "us-west-000",
|
|
"001": "us-west-001",
|
|
"002": "us-west-002",
|
|
"003": "us-west-004",
|
|
"004": "us-west-004",
|
|
"005": "us-east-005",
|
|
}
|
|
prefix = key_id[:3] if key_id and len(key_id) >= 3 else ""
|
|
region = region_map.get(prefix, "us-west-004")
|
|
|
|
if key_id.startswith("003"):
|
|
region = "us-west-004"
|
|
elif key_id.startswith("005"):
|
|
region = "us-east-005"
|
|
|
|
return f"https://s3.{region}.backblazeb2.com"
|
|
|
|
|
|
def get_storage_client(db=None, tenant_id=None, provider: Optional[str] = None):
|
|
"""
|
|
Get storage client dynamically.
|
|
1. Checks for tenant-specific configuration if tenant_id is provided.
|
|
2. Falls back to global system configuration.
|
|
Mandatory Backblaze B2.
|
|
"""
|
|
cache_key = (tenant_id, provider)
|
|
cached = _storage_client_cache.get(cache_key)
|
|
if cached is not None and time.time() - cached[3] < _STORAGE_CLIENT_CACHE_TTL:
|
|
return cached[0], cached[1], cached[2]
|
|
|
|
if db is None:
|
|
logger.error("get_storage_client called without DB session. Cannot fetch B2 config.")
|
|
raise RuntimeError("Storage configuration requires a database session.")
|
|
|
|
from app.modules.configuration.models.system_configuration_model import SystemConfiguration
|
|
from app.modules.tenant.models.tenant_storage_config_model import TenantStorageConfig
|
|
from app.core.crypto import decrypt_data
|
|
|
|
if tenant_id:
|
|
tenant_cfg = db.query(TenantStorageConfig).filter(
|
|
TenantStorageConfig.tenant_id == tenant_id,
|
|
TenantStorageConfig.is_active.is_(True)
|
|
).first()
|
|
|
|
if tenant_cfg:
|
|
try:
|
|
key_id = (tenant_cfg.b2_key_id or "").strip()
|
|
application_key = decrypt_data(tenant_cfg.encrypted_b2_application_key) if tenant_cfg.encrypted_b2_application_key else None
|
|
|
|
if not key_id or not application_key:
|
|
logger.warning(f"[Storage] Tenant {tenant_id} configuration is incomplete (missing key or secret). Falling back to global.")
|
|
else:
|
|
endpoint = tenant_cfg.b2_endpoint
|
|
if endpoint and endpoint.strip():
|
|
endpoint = endpoint.strip()
|
|
if not endpoint.startswith("http"):
|
|
endpoint = f"https://{endpoint}"
|
|
else:
|
|
endpoint = _derive_b2_endpoint(key_id)
|
|
|
|
quarantine_bucket = tenant_cfg.b2_quarantine_bucket or "quarantine-bucket"
|
|
clean_bucket = tenant_cfg.b2_clean_bucket or "clean-document-bucket"
|
|
|
|
logger.info(f"[Storage] Initializing B2 client for tenant {tenant_id} (Key ID: ...{key_id[-4:]})")
|
|
|
|
client = B2StorageClient(
|
|
key_id=key_id,
|
|
application_key=application_key,
|
|
endpoint_url=endpoint
|
|
)
|
|
_storage_client_cache[cache_key] = (client, quarantine_bucket, clean_bucket, time.time())
|
|
return client, quarantine_bucket, clean_bucket
|
|
except Exception as e:
|
|
logger.error(f"[Storage] Failed to initialize tenant storage for {tenant_id}: {e}. Falling back to global.")
|
|
else:
|
|
logger.debug(f"[Storage] No active tenant config for {tenant_id}, using global settings.")
|
|
|
|
config_items = db.query(SystemConfiguration).filter(SystemConfiguration.config_key.like("storage.%")).all()
|
|
configs = {cfg.config_key: cfg for cfg in config_items}
|
|
|
|
def get_cfg_val(key, default=None):
|
|
cfg = configs.get(key)
|
|
return cfg.text_value if cfg and cfg.text_value is not None else default
|
|
|
|
key_id = get_cfg_val("storage.b2_key_id")
|
|
enc_app_key = get_cfg_val("storage.b2_application_key")
|
|
|
|
if not key_id or not enc_app_key:
|
|
logger.error("Backblaze B2 is not configured in system settings.")
|
|
raise RuntimeError("Storage is not configured. Please set up Backblaze B2 in System Settings.")
|
|
|
|
try:
|
|
key_id = key_id.strip()
|
|
app_key = decrypt_data(enc_app_key).strip()
|
|
|
|
custom_endpoint = get_cfg_val("storage.b2_endpoint")
|
|
if custom_endpoint and custom_endpoint.strip():
|
|
endpoint = custom_endpoint.strip()
|
|
if not endpoint.startswith("http"):
|
|
endpoint = f"https://{endpoint}"
|
|
else:
|
|
endpoint = _derive_b2_endpoint(key_id)
|
|
|
|
quarantine = get_cfg_val("storage.b2_quarantine_bucket", "quarantine-bucket").strip()
|
|
clean = get_cfg_val("storage.b2_clean_bucket", "clean-document-bucket").strip()
|
|
|
|
b2_client = B2StorageClient(key_id, app_key, endpoint_url=endpoint)
|
|
logger.info(f"[Storage] Initializing global B2 client (Key: ...{key_id[-4:]})")
|
|
_storage_client_cache[cache_key] = (b2_client, quarantine, clean, time.time())
|
|
return b2_client, quarantine, clean
|
|
except Exception as e:
|
|
logger.error(f"Failed to initialize Global B2 client: {e}")
|
|
raise RuntimeError(f"Failed to initialize storage client: {e}") |