| """
|
| Supabase dataset ingestion service.
|
|
|
| Fetches document metadata and binary files (benign vs prompt-injection)
|
| from Supabase PostgreSQL (public.uploads) and Storage (team-files bucket)
|
| for AI model training and testing.
|
| """
|
|
|
| import os
|
| from typing import List, Dict, Optional
|
| from supabase import create_client, Client
|
|
|
| from app.core.config import settings
|
| from app.core.logging import get_logger
|
|
|
| logger = get_logger(__name__)
|
|
|
|
|
| class SupabaseDatasetService:
|
| """Service to interact with Supabase storage and database for training datasets."""
|
|
|
| def __init__(self):
|
| self._client: Optional[Client] = None
|
|
|
| @property
|
| def client(self) -> Client:
|
| """Lazy-initialize Supabase client."""
|
| if self._client is None:
|
| if not settings.SUPABASE_URL or not settings.SUPABASE_KEY:
|
| raise ValueError("SUPABASE_URL and SUPABASE_KEY must be configured in environment.")
|
| self._client = create_client(settings.SUPABASE_URL, settings.SUPABASE_KEY)
|
| return self._client
|
|
|
| @property
|
| def bucket_name(self) -> str:
|
| return settings.SUPABASE_STORAGE_BUCKET
|
|
|
| def list_dataset_records(self, category: Optional[str] = None) -> List[Dict]:
|
| """Fetch metadata records from `public.uploads` table."""
|
| try:
|
| query = self.client.from_("uploads").select("*")
|
| if category in ["benign", "injection"]:
|
| query = query.eq("category", category)
|
| response = query.order("created_at", desc=True).execute()
|
| return response.data or []
|
| except Exception as e:
|
| logger.error("Failed to list Supabase uploads: %s", str(e))
|
| raise
|
|
|
| def get_file_download_url(self, storage_path: str) -> str:
|
| """Get public download URL for a storage object."""
|
| res = self.client.storage.from_(self.bucket_name).get_public_url(storage_path)
|
| return res
|
|
|
| def download_file_bytes(self, storage_path: str) -> bytes:
|
| """Download raw binary content of a file from Supabase storage."""
|
| response = self.client.storage.from_(self.bucket_name).download(storage_path)
|
| return response
|
|
|
| def sync_dataset_to_disk(self, target_dir: str = settings.DATASET_BASE_DIR) -> Dict[str, int]:
|
| """
|
| Synchronize all clean (benign) and injected (injection) documents
|
| from Supabase storage to local disk under target_dir/benign and target_dir/injection.
|
| """
|
| records = self.list_dataset_records()
|
| stats = {"benign": 0, "injection": 0, "failed": 0, "skipped": 0}
|
|
|
| for item in records:
|
| cat = item.get("category")
|
| path = item.get("storage_path")
|
| file_name = item.get("file_name")
|
| record_id = item.get("id")
|
|
|
| if not cat or not path:
|
| continue
|
|
|
|
|
| cat_dir = os.path.join(target_dir, cat)
|
| os.makedirs(cat_dir, exist_ok=True)
|
|
|
| safe_filename = f"{record_id}_{file_name}" if record_id else file_name
|
| local_file_path = os.path.join(cat_dir, safe_filename)
|
|
|
|
|
| if os.path.exists(local_file_path):
|
| stats[cat] = stats.get(cat, 0) + 1
|
| stats["skipped"] += 1
|
| continue
|
|
|
| try:
|
| file_bytes = self.download_file_bytes(path)
|
| with open(local_file_path, "wb") as f:
|
| f.write(file_bytes)
|
| stats[cat] = stats.get(cat, 0) + 1
|
| logger.info("Downloaded dataset file: %s -> %s", path, local_file_path)
|
| except Exception as e:
|
| logger.error("Failed to download dataset file %s: %s", path, str(e))
|
| stats["failed"] += 1
|
|
|
| return stats
|
|
|
|
|
| dataset_service = SupabaseDatasetService()
|
|
|
