Spaces:
Running
Running
File size: 9,220 Bytes
7af9e82 | 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 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 | """
User Files Service - Manages file metadata in Supabase
This service:
1. Tracks all user files (plots, CSVs, reports, models) in Supabase
2. Provides file listing for the Assets panel
3. Handles file expiration and cleanup coordination
4. Works with R2StorageService for actual file storage
"""
import os
from datetime import datetime, timedelta
from typing import Optional, Dict, Any, List
from dataclasses import dataclass
from enum import Enum
# Supabase client import
try:
from supabase import create_client, Client
except ImportError:
print("Warning: supabase package not installed. Run: pip install supabase")
Client = None
SUPABASE_URL = os.getenv("SUPABASE_URL", "")
SUPABASE_SERVICE_KEY = os.getenv("SUPABASE_SERVICE_KEY", "") # Use service key for backend
class FileType(Enum):
PLOT = "plot"
CSV = "csv"
REPORT = "report"
MODEL = "model"
@dataclass
class UserFile:
"""Represents a user file record."""
id: str
user_id: str
session_id: Optional[str]
file_type: FileType
file_name: str
r2_key: str
size_bytes: int
mime_type: str
metadata: Dict[str, Any]
created_at: datetime
expires_at: datetime
download_url: Optional[str] = None
class UserFilesService:
"""Service for managing user file metadata in Supabase."""
def __init__(self):
"""Initialize Supabase client."""
if not SUPABASE_URL or not SUPABASE_SERVICE_KEY:
raise ValueError("SUPABASE_URL and SUPABASE_SERVICE_KEY must be set")
self.client: Client = create_client(SUPABASE_URL, SUPABASE_SERVICE_KEY)
self.table = "user_files"
# ==================== CREATE ====================
def create_file_record(
self,
user_id: str,
file_type: FileType,
file_name: str,
r2_key: str,
size_bytes: int,
session_id: Optional[str] = None,
mime_type: str = "application/octet-stream",
metadata: Optional[Dict[str, Any]] = None,
expires_in_days: int = 7
) -> UserFile:
"""
Create a file record in Supabase.
Args:
user_id: User ID
file_type: Type of file
file_name: Display name
r2_key: R2 storage key
size_bytes: File size
session_id: Optional chat session ID
mime_type: MIME type
metadata: Additional metadata (plot type, metrics, etc.)
expires_in_days: Days until file expires
Returns:
Created UserFile record
"""
expires_at = datetime.utcnow() + timedelta(days=expires_in_days)
data = {
"user_id": user_id,
"session_id": session_id,
"file_type": file_type.value,
"file_name": file_name,
"r2_key": r2_key,
"size_bytes": size_bytes,
"mime_type": mime_type,
"metadata": metadata or {},
"expires_at": expires_at.isoformat()
}
result = self.client.table(self.table).insert(data).execute()
if result.data:
return self._to_user_file(result.data[0])
raise Exception("Failed to create file record")
# ==================== READ ====================
def get_user_files(
self,
user_id: str,
file_type: Optional[FileType] = None,
session_id: Optional[str] = None,
include_expired: bool = False
) -> List[UserFile]:
"""
Get all files for a user.
Args:
user_id: User ID
file_type: Optional filter by type
session_id: Optional filter by session
include_expired: Include expired files
Returns:
List of UserFile records
"""
query = self.client.table(self.table)\
.select("*")\
.eq("user_id", user_id)\
.eq("is_deleted", False)
if file_type:
query = query.eq("file_type", file_type.value)
if session_id:
query = query.eq("session_id", session_id)
if not include_expired:
query = query.gt("expires_at", datetime.utcnow().isoformat())
query = query.order("created_at", desc=True)
result = query.execute()
return [self._to_user_file(row) for row in (result.data or [])]
def get_file_by_id(self, file_id: str) -> Optional[UserFile]:
"""Get a specific file by ID."""
result = self.client.table(self.table)\
.select("*")\
.eq("id", file_id)\
.single()\
.execute()
if result.data:
return self._to_user_file(result.data)
return None
def get_file_by_r2_key(self, r2_key: str) -> Optional[UserFile]:
"""Get a file by R2 key."""
result = self.client.table(self.table)\
.select("*")\
.eq("r2_key", r2_key)\
.single()\
.execute()
if result.data:
return self._to_user_file(result.data)
return None
def get_session_files(self, session_id: str) -> List[UserFile]:
"""Get all files for a chat session."""
result = self.client.table(self.table)\
.select("*")\
.eq("session_id", session_id)\
.eq("is_deleted", False)\
.order("created_at", desc=True)\
.execute()
return [self._to_user_file(row) for row in (result.data or [])]
# ==================== UPDATE ====================
def extend_expiration(self, file_id: str, additional_days: int = 7) -> bool:
"""Extend file expiration date."""
file = self.get_file_by_id(file_id)
if not file:
return False
new_expires = datetime.utcnow() + timedelta(days=additional_days)
result = self.client.table(self.table)\
.update({"expires_at": new_expires.isoformat()})\
.eq("id", file_id)\
.execute()
return bool(result.data)
# ==================== DELETE ====================
def soft_delete_file(self, file_id: str) -> bool:
"""Soft delete a file (mark as deleted)."""
result = self.client.table(self.table)\
.update({"is_deleted": True})\
.eq("id", file_id)\
.execute()
return bool(result.data)
def hard_delete_file(self, file_id: str) -> bool:
"""Permanently delete a file record."""
result = self.client.table(self.table)\
.delete()\
.eq("id", file_id)\
.execute()
return bool(result.data)
def get_expired_files(self) -> List[UserFile]:
"""Get all expired files for cleanup."""
result = self.client.table(self.table)\
.select("*")\
.lt("expires_at", datetime.utcnow().isoformat())\
.eq("is_deleted", False)\
.execute()
return [self._to_user_file(row) for row in (result.data or [])]
# ==================== STATS ====================
def get_user_storage_stats(self, user_id: str) -> Dict[str, Any]:
"""Get storage statistics for a user."""
files = self.get_user_files(user_id, include_expired=False)
stats = {
"total_files": len(files),
"total_size_bytes": sum(f.size_bytes for f in files),
"by_type": {}
}
for file_type in FileType:
type_files = [f for f in files if f.file_type == file_type]
stats["by_type"][file_type.value] = {
"count": len(type_files),
"size_bytes": sum(f.size_bytes for f in type_files)
}
stats["total_size_mb"] = round(stats["total_size_bytes"] / (1024 * 1024), 2)
return stats
# ==================== HELPERS ====================
def _to_user_file(self, row: Dict[str, Any]) -> UserFile:
"""Convert database row to UserFile object."""
return UserFile(
id=row["id"],
user_id=row["user_id"],
session_id=row.get("session_id"),
file_type=FileType(row["file_type"]),
file_name=row["file_name"],
r2_key=row["r2_key"],
size_bytes=row.get("size_bytes", 0),
mime_type=row.get("mime_type", "application/octet-stream"),
metadata=row.get("metadata", {}),
created_at=datetime.fromisoformat(row["created_at"].replace("Z", "+00:00")),
expires_at=datetime.fromisoformat(row["expires_at"].replace("Z", "+00:00"))
)
# ==================== SINGLETON ====================
_files_service: Optional[UserFilesService] = None
def get_files_service() -> UserFilesService:
"""Get or create UserFilesService singleton."""
global _files_service
if _files_service is None:
_files_service = UserFilesService()
return _files_service
|