File size: 4,033 Bytes
215f97f | 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 | """
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
# Target directory: e.g. ./data/raw/benign or ./data/raw/injection
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)
# Skip download if file already exists locally
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()
|