Spaces:
Sleeping
Sleeping
| import os | |
| from fastapi import FastAPI, File, UploadFile, HTTPException, Form, Header | |
| from fastapi.middleware.cors import CORSMiddleware | |
| from fastapi.responses import JSONResponse, StreamingResponse | |
| from fastapi.encoders import jsonable_encoder | |
| from pydantic import BaseModel, Field, EmailStr | |
| from typing import Optional, List, Dict, Any | |
| import pandas as pd | |
| import io | |
| import json | |
| import time | |
| import asyncio | |
| from backend import RegexClassifier | |
| from email_service import send_welcome_email | |
| import evaluator_api | |
| import video_router | |
| import video_job_queue | |
| # db.schema_inspector provides unified SQLAlchemy-based schema reflection (TICKET-4) | |
| from db.schema_inspector import SchemaInspector | |
| _schema_inspector = SchemaInspector() # module-level singleton — stateless, safe to share | |
| # Initialize FastAPI app | |
| app = FastAPI( | |
| title="Segmento Sense API", | |
| description="AI-powered PII Detection and Data Classification Platform", | |
| version="1.0.0" | |
| ) | |
| # CORS Configuration | |
| # Keep all known origins: production domains + local dev ports | |
| app.add_middleware( | |
| CORSMiddleware, | |
| allow_origins=[ | |
| # ── Production ────────────────────────────────────── | |
| "https://segmento.in", | |
| "https://www.segmento.in", | |
| "https://segmento-sense.vercel.app", | |
| # ── Local development ──────────────────────────────── | |
| "http://localhost:3000", | |
| "http://localhost:3001", | |
| "http://localhost:3002", | |
| ], | |
| allow_credentials=True, | |
| allow_methods=["*"], | |
| allow_headers=["*"], | |
| ) | |
| # Initialize the classifier | |
| classifier = RegexClassifier() | |
| # Wire evaluator API — must happen after classifier is ready | |
| evaluator_api.setup(classifier) | |
| app.include_router(evaluator_api.router) | |
| # Wire video API — async job queue (startup worker launched below) | |
| video_job_queue.setup(classifier) | |
| app.include_router(video_router.router) | |
| async def _start_video_worker(): | |
| """Launch the background asyncio video processing worker.""" | |
| await video_job_queue.startup() | |
| # Maximum file size (1GB) | |
| MAX_FILE_SIZE = 1024 * 1024 * 1024 # 1GB in bytes | |
| # ==================== PYDANTIC MODELS ==================== | |
| class TextAnalysisRequest(BaseModel): | |
| text: str = Field(..., description="Text to analyze for PII") | |
| class PatternAddRequest(BaseModel): | |
| name: str = Field(..., description="Pattern name") | |
| regex: str = Field(..., description="Regex pattern") | |
| class DatabaseConnectionRequest(BaseModel): | |
| host: str | |
| port: str | |
| database: str | |
| user: str | |
| password: str | |
| table: str = Field(None, description="Table name (or collection for MongoDB)") | |
| scan_mode: str = Field("full", description="full, incremental, sampling, metadata_only, metadata_and_sampling") | |
| class AwsRdsConnectionRequest(DatabaseConnectionRequest): | |
| engine: str = Field(..., description="E.g., postgres, mysql, aurora_mysql") | |
| class DynamoDbConnectionRequest(BaseModel): | |
| access_key: str | |
| secret_key: str | |
| region: str | |
| table: str = Field(None, description="Table name") | |
| class MetadataScanRequest(DatabaseConnectionRequest): | |
| connector_type: str = Field(..., description="E.g., postgresql or mysql") | |
| engine: Optional[str] = Field(None, description="Only for aws-rds") | |
| access_key: Optional[str] = None | |
| secret_key: Optional[str] = None | |
| region: Optional[str] = None | |
| class S3ConnectionRequest(BaseModel): | |
| access_key: str | |
| secret_key: str | |
| region: str | |
| bucket: str = Field(None, description="Bucket name") | |
| file_key: str = Field(None, description="File key/path") | |
| class AzureConnectionRequest(BaseModel): | |
| connection_string: str | |
| container: str = Field(None, description="Container name") | |
| blob: str = Field(None, description="Blob name") | |
| class GCSConnectionRequest(BaseModel): | |
| credentials: Dict[str, Any] | |
| bucket: str = Field(None, description="Bucket name") | |
| file_name: str = Field(None, description="File name") | |
| class GoogleDriveRequest(BaseModel): | |
| credentials: Dict[str, Any] | |
| file_id: str = Field(None, description="Drive file ID") | |
| mime_type: str = Field(None, description="File MIME type") | |
| class SlackRequest(BaseModel): | |
| token: str | |
| channel_id: str | |
| class GmailRequest(BaseModel): | |
| credentials: Dict[str, Any] | |
| auth_type: str = Field("service_account", description="service_account or oauth2_token") | |
| class ZendeskRequest(BaseModel): | |
| subdomain: str | |
| email: str | |
| api_token: str | |
| class SalesforceRequest(BaseModel): | |
| instance_url: str | |
| access_token: str | |
| class GlueCredentials(BaseModel): | |
| access_key: str | |
| secret_key: str | |
| region: str | |
| class GlueListTablesRequest(GlueCredentials): | |
| database_name: str | |
| class GlueScanRequest(GlueCredentials): | |
| database_name: str | |
| table_name: str | |
| class ConfluenceRequest(BaseModel): | |
| url: str | |
| username: str | |
| token: str | |
| page_id: str | |
| class PDFPageRequest(BaseModel): | |
| page_number: int = 0 | |
| class WelcomeEmailRequest(BaseModel): | |
| name: str = Field(..., description="User's name") | |
| email: EmailStr = Field(..., description="User's email address") | |
| class RegisterRequest(BaseModel): | |
| name: str | |
| email: str | |
| password: str | |
| class LoginRequest(BaseModel): | |
| email: str | |
| password: str | |
| # ==================== HELPER FUNCTIONS ==================== | |
| def validate_file_size(file: UploadFile): | |
| """Validate uploaded file size""" | |
| file.file.seek(0, 2) # Seek to end | |
| size = file.file.tell() # Get position (file size) | |
| file.file.seek(0) # Reset to beginning | |
| if size > MAX_FILE_SIZE: | |
| raise HTTPException( | |
| status_code=413, | |
| detail=f"File size ({size} bytes) exceeds maximum allowed size (1GB)" | |
| ) | |
| return size | |
| def format_pii_response(df: pd.DataFrame, source_df: pd.DataFrame = None, text: str = None, | |
| selected_models: list = None) -> Dict: | |
| """Format PII analysis response""" | |
| count_df = classifier.get_pii_counts_dataframe(df, selected_models) if source_df is not None else classifier.get_pii_counts(text, selected_models) | |
| response = { | |
| "pii_counts": count_df.fillna("").to_dict(orient="records") if not count_df.empty else [], | |
| "total_pii_found": int(count_df["Count"].sum()) if not count_df.empty else 0 | |
| } | |
| # Add schema if source dataframe provided | |
| if source_df is not None and not source_df.empty: | |
| schema_df = classifier.get_data_schema(source_df) | |
| response["schema"] = schema_df.fillna("").to_dict(orient="records") | |
| # Add inspector results if text provided | |
| if text: | |
| inspector_df = classifier.run_full_inspection(text, selected_models) | |
| if not inspector_df.empty: | |
| response["inspector"] = inspector_df.fillna("").to_dict(orient="records") | |
| return response | |
| # ==================== MODEL CATALOGUE ENDPOINT ==================== | |
| async def get_available_models(): | |
| """ | |
| Returns the full list of available PII detection models, | |
| separated into always-on and lazy-loaded categories. | |
| """ | |
| return JSONResponse(content={ | |
| "always_on": [ | |
| {"key": "regex", "label": "🛠️ Regex", "description": "Fast rule-based pattern matching (emails, phones, SSNs)"}, | |
| {"key": "nltk", "label": "🧠 NLTK", "description": "Statistical NLP chunker for names and locations"}, | |
| {"key": "spacy", "label": "🤖 SpaCy", "description": "Industrial-strength NER (en_core_web_lg)"}, | |
| {"key": "presidio", "label": "🛡️ Presidio", "description": "Microsoft Presidio — enterprise PII analyser"}, | |
| {"key": "gliner", "label": "🦅 GLiNER", "description": "Zero-shot entity extraction (urchade/gliner_small-v2.1)"}, | |
| {"key": "deberta", "label": "🚀 DeBERTa", "description": "Kaggle-winning DeBERTa V3 fine-tuned for PII"}, | |
| ], | |
| "lazy_loaded": [ | |
| {"key": "pasteproof", "label": "📋 Pasteproof", "description": "joneauxedgar/pasteproof-pii-detector-v2 — broad PII detection"}, | |
| {"key": "piiranha", "label": "🐟 Piiranha", "description": "iiiorg/piiranha-v1 — personal information specialist"}, | |
| {"key": "nvidia_gliner", "label": "⚡ NVIDIA-GLiNER", "description": "nvidia/gliner-PII — enterprise-grade zero-shot NER"}, | |
| {"key": "mmbert", "label": "🌐 mmbert32k", "description": "llm-semantic-router/mmbert32k — 32k-context document scanner"}, | |
| ] | |
| }) | |
| # ==================== FILE UPLOAD ENDPOINTS ==================== | |
| async def upload_csv(file: UploadFile = File(...), mask: bool = Form(False), | |
| selected_models: str = Form("")): | |
| """Upload and analyze CSV file""" | |
| try: | |
| validate_file_size(file) | |
| content = await file.read() | |
| df = pd.read_csv(io.BytesIO(content)) | |
| models = [m.strip() for m in selected_models.split(",") if m.strip()] or None | |
| text_sample = df.head(10).to_string() | |
| response = format_pii_response(df, df, text_sample, models) | |
| if mask: | |
| masked_df = classifier.mask_dataframe(df.head(50), models) | |
| response["data"] = masked_df.fillna("").to_dict(orient="records") | |
| else: | |
| highlighted_df = classifier.scan_dataframe_with_html(df.head(50), models) | |
| response["data"] = highlighted_df.fillna("").to_dict(orient="records") | |
| return JSONResponse(content=response) | |
| except Exception as e: | |
| raise HTTPException(status_code=500, detail=str(e)) | |
| async def upload_txt(file: UploadFile = File(...), mask: bool = Form(False), | |
| selected_models: str = Form("")): | |
| """Upload and analyze TXT file""" | |
| try: | |
| validate_file_size(file) | |
| content = await file.read() | |
| text = content.decode("utf-8") | |
| df = pd.DataFrame({"Content": [text]}) | |
| models = [m.strip() for m in selected_models.split(",") if m.strip()] or None | |
| text_sample = text[:2000] | |
| response = format_pii_response(df, df, text_sample, models) | |
| if mask: | |
| masked_df = classifier.mask_dataframe(df, models) | |
| response["data"] = masked_df.fillna("").to_dict(orient="records") | |
| else: | |
| highlighted_df = classifier.scan_dataframe_with_html(df, models) | |
| response["data"] = highlighted_df.fillna("").to_dict(orient="records") | |
| return JSONResponse(content=response) | |
| except Exception as e: | |
| raise HTTPException(status_code=500, detail=str(e)) | |
| async def upload_json(file: UploadFile = File(...), mask: bool = Form(False), | |
| selected_models: str = Form("")): | |
| """Upload and analyze JSON file""" | |
| try: | |
| validate_file_size(file) | |
| df = classifier.get_json_data(file.file) | |
| models = [m.strip() for m in selected_models.split(",") if m.strip()] or None | |
| text_sample = df.head(10).to_string() | |
| response = format_pii_response(df, df, text_sample, models) | |
| if mask: | |
| masked_df = classifier.mask_dataframe(df.head(50), models) | |
| response["data"] = masked_df.fillna("").to_dict(orient="records") | |
| else: | |
| highlighted_df = classifier.scan_dataframe_with_html(df.head(50), models) | |
| response["data"] = highlighted_df.fillna("").to_dict(orient="records") | |
| return JSONResponse(content=response) | |
| except Exception as e: | |
| raise HTTPException(status_code=500, detail=str(e)) | |
| async def upload_parquet(file: UploadFile = File(...), mask: bool = Form(False), | |
| selected_models: str = Form("")): | |
| """Upload and analyze Parquet file""" | |
| try: | |
| validate_file_size(file) | |
| content = await file.read() | |
| df = classifier.get_parquet_data(content) | |
| models = [m.strip() for m in selected_models.split(",") if m.strip()] or None | |
| text_sample = df.head(10).to_string() | |
| response = format_pii_response(df, df, text_sample, models) | |
| if mask: | |
| masked_df = classifier.mask_dataframe(df.head(50), models) | |
| response["data"] = masked_df.fillna("").to_dict(orient="records") | |
| else: | |
| highlighted_df = classifier.scan_dataframe_with_html(df.head(50), models) | |
| response["data"] = highlighted_df.fillna("").to_dict(orient="records") | |
| return JSONResponse(content=response) | |
| except Exception as e: | |
| raise HTTPException(status_code=500, detail=str(e)) | |
| async def upload_avro(file: UploadFile = File(...), mask: bool = Form(False), | |
| selected_models: str = Form("")): | |
| """Upload and analyze Apache Avro file""" | |
| try: | |
| validate_file_size(file) | |
| content = await file.read() | |
| df = classifier.get_avro_data(content) | |
| models = [m.strip() for m in selected_models.split(",") if m.strip()] or None | |
| text_sample = df.head(10).to_string() | |
| response = format_pii_response(df, df, text_sample, models) | |
| if mask: | |
| masked_df = classifier.mask_dataframe(df.head(50), models) | |
| response["data"] = masked_df.fillna("").to_dict(orient="records") | |
| else: | |
| highlighted_df = classifier.scan_dataframe_with_html(df.head(50), models) | |
| response["data"] = highlighted_df.fillna("").to_dict(orient="records") | |
| return JSONResponse(content=response) | |
| except Exception as e: | |
| raise HTTPException(status_code=500, detail=str(e)) | |
| async def upload_pdf(file: UploadFile = File(...), page_number: int = Form(0)): | |
| """Upload and analyze PDF file (with pagination)""" | |
| try: | |
| validate_file_size(file) | |
| content = await file.read() | |
| # Get total pages and extract text from specific page | |
| total_pages = classifier.get_pdf_total_pages(content) | |
| text = classifier.get_pdf_page_text(content, page_number) | |
| # Format PII response | |
| response = format_pii_response(None, None, text) | |
| response["total_pages"] = total_pages | |
| response["current_page"] = page_number | |
| # Get labeled PDF image | |
| img = classifier.get_labeled_pdf_image(content, page_number) | |
| if img: | |
| import base64 | |
| from PIL import Image | |
| # Check if img is already bytes or a PIL Image | |
| if isinstance(img, bytes): | |
| # Already bytes, just encode | |
| img_str = base64.b64encode(img).decode() | |
| elif isinstance(img, Image.Image): | |
| # PIL Image, need to convert to bytes | |
| buffered = io.BytesIO() | |
| img.save(buffered, format="PNG") | |
| img_str = base64.b64encode(buffered.getvalue()).decode() | |
| else: | |
| # Unknown type, skip image | |
| img_str = None | |
| if img_str: | |
| response["image"] = f"data:image/png;base64,{img_str}" | |
| return JSONResponse(content=response) | |
| except Exception as e: | |
| raise HTTPException(status_code=500, detail=str(e)) | |
| async def upload_image(file: UploadFile = File(...), mask: bool = Form(False)): | |
| """Upload and analyze image with OCR""" | |
| try: | |
| validate_file_size(file) | |
| content = await file.read() | |
| # Extract text via OCR | |
| text = classifier.get_ocr_text_from_image(content) | |
| if not text: | |
| raise HTTPException(status_code=400, detail="No text could be extracted from the image") | |
| df = pd.DataFrame({"Content": [text]}) | |
| response = format_pii_response(df, df, text) | |
| if mask: | |
| masked_df = classifier.mask_dataframe(df) | |
| response["data"] = masked_df.fillna("").to_dict(orient="records") | |
| else: | |
| highlighted_df = classifier.scan_dataframe_with_html(df) | |
| response["data"] = highlighted_df.fillna("").to_dict(orient="records") | |
| # Return original image as base64 | |
| import base64 | |
| img_str = base64.b64encode(content).decode() | |
| response["original_image"] = f"data:image/png;base64,{img_str}" | |
| return JSONResponse(content=response) | |
| except Exception as e: | |
| raise HTTPException(status_code=500, detail=str(e)) | |
| # ==================== ANALYSIS ENDPOINTS ==================== | |
| async def analyze_text(request: TextAnalysisRequest): | |
| """Analyze plain text for PII""" | |
| try: | |
| matches = classifier.analyze_text_hybrid(request.text) | |
| count_df = classifier.get_pii_counts(request.text) | |
| return JSONResponse(content={ | |
| "matches": matches, | |
| "pii_counts": count_df.fillna("").to_dict(orient="records") if not count_df.empty else [], | |
| "total_pii_found": len(matches) | |
| }) | |
| except Exception as e: | |
| raise HTTPException(status_code=500, detail=str(e)) | |
| async def inspect_text(request: TextAnalysisRequest): | |
| """Run full model inspection on text""" | |
| try: | |
| inspector_df = classifier.run_full_inspection(request.text) | |
| if inspector_df.empty: | |
| return JSONResponse(content={ | |
| "inspector": [], | |
| "message": "No PII detected by any model" | |
| }) | |
| return JSONResponse(content={ | |
| "inspector": inspector_df.fillna("").to_dict(orient="records") | |
| }) | |
| except Exception as e: | |
| raise HTTPException(status_code=500, detail=str(e)) | |
| async def mask_text(request: TextAnalysisRequest): | |
| """Mask PII in text""" | |
| try: | |
| df = pd.DataFrame({"Content": [request.text]}) | |
| masked_df = classifier.mask_dataframe(df) | |
| return JSONResponse(content={ | |
| "original": request.text, | |
| "masked": masked_df.iloc[0]["Content"] | |
| }) | |
| except Exception as e: | |
| raise HTTPException(status_code=500, detail=str(e)) | |
| # ==================== PATTERN MANAGEMENT ==================== | |
| async def get_patterns(): | |
| """Get all regex patterns""" | |
| try: | |
| patterns = classifier.list_patterns() | |
| return JSONResponse(content={ | |
| "patterns": [{"name": k, "regex": v} for k, v in patterns.items()] | |
| }) | |
| except Exception as e: | |
| raise HTTPException(status_code=500, detail=str(e)) | |
| async def add_pattern(request: PatternAddRequest): | |
| """Add a new regex pattern""" | |
| try: | |
| classifier.add_pattern(request.name, request.regex) | |
| return JSONResponse(content={ | |
| "message": f"Pattern '{request.name}' added successfully", | |
| "pattern": {"name": request.name, "regex": request.regex} | |
| }) | |
| except Exception as e: | |
| raise HTTPException(status_code=500, detail=str(e)) | |
| async def delete_pattern(pattern_name: str): | |
| """Remove a regex pattern""" | |
| try: | |
| classifier.remove_pattern(pattern_name) | |
| return JSONResponse(content={ | |
| "message": f"Pattern '{pattern_name}' removed successfully" | |
| }) | |
| except Exception as e: | |
| raise HTTPException(status_code=500, detail=str(e)) | |
| # ==================== DATABASE CONNECTORS ==================== | |
| async def connect_postgresql(request: DatabaseConnectionRequest, authorization: str = Header(None)): | |
| """Connect to PostgreSQL, scan table, persist to Supabase, return analysis.""" | |
| uid = _uid_from_auth(authorization) | |
| try: | |
| limit = 5 if request.scan_mode in ("sampling", "metadata_and_sampling") else 100 | |
| df = classifier.get_postgres_data( | |
| request.host, request.port, request.database, | |
| request.user, request.password, request.table, limit=limit | |
| ) | |
| text_sample = df.head(10).to_string() | |
| response = format_pii_response(df, df, text_sample) | |
| rows_scanned = len(df) | |
| response["rows_scanned"] = rows_scanned | |
| masked_df = classifier.mask_dataframe(df.head(50)) | |
| response["data"] = masked_df.fillna("").to_dict(orient="records") | |
| # If hybrid mode, fetch metadata and append to response | |
| if request.scan_mode == "metadata_and_sampling": | |
| from classifier_manager.metadata_scanner import MetadataScanner | |
| scanner = MetadataScanner() | |
| cols_meta = classifier.pg_handler.get_schema_metadata( | |
| request.host, request.port, request.database, | |
| request.user, request.password, request.table | |
| ) | |
| cat_resp = scanner.scan_schema(cols_meta) | |
| response["metadata"] = cat_resp.get("metadata", {}) | |
| # ── Persist to Supabase universal schema ────────────────────────── | |
| from db.supabase_client import persist_scan_results | |
| persist_scan_results( | |
| uid=uid, | |
| connector_type="postgresql", | |
| db_name=request.database, | |
| table_name=request.table, | |
| response=response, | |
| column_count=len(df.columns), | |
| rows_scanned=rows_scanned, | |
| scan_mode=request.scan_mode | |
| ) | |
| # ───────────────────────────────────────────────────────────────── | |
| return JSONResponse(content=response) | |
| except Exception as e: | |
| raise HTTPException(status_code=500, detail=f"PostgreSQL connection failed: {str(e)}") | |
| async def connect_mysql(request: DatabaseConnectionRequest, authorization: str = Header(None)): | |
| """Connect to MySQL, scan table, persist to Supabase, return analysis.""" | |
| uid = _uid_from_auth(authorization) | |
| try: | |
| limit = 5 if request.scan_mode in ("sampling", "metadata_and_sampling") else 100 | |
| df = classifier.get_mysql_data( | |
| request.host, request.port, request.database, | |
| request.user, request.password, request.table, limit=limit | |
| ) | |
| text_sample = df.head(10).to_string() | |
| response = format_pii_response(df, df, text_sample) | |
| rows_scanned = len(df) | |
| response["rows_scanned"] = rows_scanned # additive: row count after 100-row cap | |
| masked_df = classifier.mask_dataframe(df.head(50)) | |
| response["data"] = masked_df.fillna("").to_dict(orient="records") | |
| # If hybrid mode, fetch metadata and append to response | |
| if request.scan_mode == "metadata_and_sampling": | |
| from classifier_manager.metadata_scanner import MetadataScanner | |
| scanner = MetadataScanner() | |
| cols_meta = classifier.mysql_handler.get_schema_metadata( | |
| request.host, request.port, request.database, | |
| request.user, request.password, request.table | |
| ) | |
| cat_resp = scanner.scan_schema(cols_meta) | |
| response["metadata"] = cat_resp.get("metadata", {}) | |
| # ── Persist to Supabase universal schema ────────────────────────── | |
| from db.supabase_client import persist_scan_results | |
| persist_scan_results( | |
| uid=uid, | |
| connector_type="mysql", | |
| db_name=request.database, | |
| table_name=request.table, | |
| response=response, | |
| column_count=len(df.columns), | |
| rows_scanned=rows_scanned, | |
| scan_mode=request.scan_mode | |
| ) | |
| # ───────────────────────────────────────────────────────────────── | |
| return JSONResponse(content=response) | |
| except Exception as e: | |
| raise HTTPException(status_code=500, detail=f"MySQL connection failed: {str(e)}") | |
| async def list_postgresql_tables(request: DatabaseConnectionRequest): | |
| """List all tables in a PostgreSQL database (via SchemaInspector).""" | |
| try: | |
| engine = _schema_inspector.build_engine( | |
| "postgresql", request.host, request.port, | |
| request.database, request.user, request.password | |
| ) | |
| tables = _schema_inspector.get_tables(engine) | |
| return JSONResponse(content={"tables": tables}) | |
| except Exception as e: | |
| raise HTTPException(status_code=500, detail=f"PostgreSQL list-tables failed: {str(e)}") | |
| async def list_mysql_tables_endpoint(request: DatabaseConnectionRequest): | |
| """List all tables in a MySQL database (via SchemaInspector).""" | |
| try: | |
| engine = _schema_inspector.build_engine( | |
| "mysql", request.host, request.port, | |
| request.database, request.user, request.password | |
| ) | |
| tables = _schema_inspector.get_tables(engine) | |
| return JSONResponse(content={"tables": tables}) | |
| except Exception as e: | |
| raise HTTPException(status_code=500, detail=f"MySQL list-tables failed: {str(e)}") | |
| async def connect_mongodb(request: DatabaseConnectionRequest, authorization: str = Header(None)): | |
| """Connect to MongoDB, scan collection, persist to Supabase, return analysis.""" | |
| uid = _uid_from_auth(authorization) | |
| try: | |
| limit = 5 if request.scan_mode in ("sampling", "metadata_and_sampling") else 100 | |
| df = classifier.get_mongodb_data( | |
| request.host, request.port, request.database, | |
| request.user, request.password, request.table, limit=limit | |
| ) | |
| text_sample = df.head(10).to_string() | |
| response = format_pii_response(df, df, text_sample) | |
| rows_scanned = len(df) | |
| response["rows_scanned"] = rows_scanned | |
| masked_df = classifier.mask_dataframe(df.head(50)) | |
| response["data"] = masked_df.fillna("").to_dict(orient="records") | |
| # ── Persist to Supabase universal schema ────────────────────────── | |
| from db.supabase_client import persist_scan_results | |
| persist_scan_results( | |
| uid=uid, | |
| connector_type="mongodb", | |
| db_name=request.database, | |
| table_name=request.table, | |
| response=response, | |
| column_count=len(df.columns), | |
| rows_scanned=rows_scanned, | |
| scan_mode=request.scan_mode | |
| ) | |
| # ───────────────────────────────────────────────────────────────── | |
| return JSONResponse(content=response) | |
| except Exception as e: | |
| raise HTTPException(status_code=500, detail=f"MongoDB connection failed: {str(e)}") | |
| async def list_mongodb_collections_endpoint(request: DatabaseConnectionRequest): | |
| """List all collections in a MongoDB database.""" | |
| try: | |
| collections = classifier.list_mongodb_collections( | |
| request.host, request.port, request.database, | |
| request.user, request.password | |
| ) | |
| return JSONResponse(content={"tables": collections}) | |
| except Exception as e: | |
| raise HTTPException(status_code=500, detail=f"MongoDB list-collections failed: {str(e)}") | |
| async def connect_mariadb(request: DatabaseConnectionRequest, authorization: str = Header(None)): | |
| """Connect to MariaDB, scan table, persist to Supabase, return analysis.""" | |
| uid = _uid_from_auth(authorization) | |
| try: | |
| limit = 5 if request.scan_mode in ("sampling", "metadata_and_sampling") else 100 | |
| df = classifier.get_mariadb_data( | |
| request.host, request.port, request.database, | |
| request.user, request.password, request.table, limit=limit | |
| ) | |
| text_sample = df.head(10).to_string() | |
| response = format_pii_response(df, df, text_sample) | |
| rows_scanned = len(df) | |
| response["rows_scanned"] = rows_scanned | |
| masked_df = classifier.mask_dataframe(df.head(50)) | |
| response["data"] = masked_df.fillna("").to_dict(orient="records") | |
| if request.scan_mode == "metadata_and_sampling": | |
| from classifier_manager.metadata_scanner import MetadataScanner | |
| scanner = MetadataScanner() | |
| cols_meta = classifier.mariadb_handler.get_schema_metadata( | |
| request.host, request.port, request.database, | |
| request.user, request.password, request.table | |
| ) | |
| cat_resp = scanner.scan_schema(cols_meta) | |
| response["metadata"] = cat_resp.get("metadata", {}) | |
| # ── Persist to Supabase universal schema ────────────────────────── | |
| from db.supabase_client import persist_scan_results | |
| persist_scan_results( | |
| uid=uid, | |
| connector_type="mariadb", | |
| db_name=request.database, | |
| table_name=request.table, | |
| response=response, | |
| column_count=len(df.columns), | |
| rows_scanned=rows_scanned, | |
| scan_mode=request.scan_mode | |
| ) | |
| # ───────────────────────────────────────────────────────────────── | |
| return JSONResponse(content=response) | |
| except Exception as e: | |
| raise HTTPException(status_code=500, detail=f"MariaDB connection failed: {str(e)}") | |
| async def list_mariadb_tables_endpoint(request: DatabaseConnectionRequest): | |
| """List all tables in a MariaDB database (via SchemaInspector).""" | |
| try: | |
| engine = _schema_inspector.build_engine( | |
| "mariadb", request.host, request.port, | |
| request.database, request.user, request.password | |
| ) | |
| tables = _schema_inspector.get_tables(engine) | |
| return JSONResponse(content={"tables": tables}) | |
| except Exception as e: | |
| raise HTTPException(status_code=500, detail=f"MariaDB list-tables failed: {str(e)}") | |
| async def connect_slack(request: SlackRequest, authorization: str = Header(None)): | |
| """Connect to Slack, fetch channel messages, scan, persist to Supabase.""" | |
| uid = _uid_from_auth(authorization) | |
| try: | |
| df = classifier.get_slack_messages(request.token, request.channel_id) | |
| if df.empty: | |
| raise ValueError("No messages found or invalid channel/token.") | |
| text_sample = df.head(10).to_string() | |
| response = format_pii_response(df, df, text_sample) | |
| rows_scanned = len(df) | |
| response["rows_scanned"] = rows_scanned | |
| masked_df = classifier.mask_dataframe(df.head(50)) | |
| response["data"] = masked_df.fillna("").to_dict(orient="records") | |
| # Supabase Persistence | |
| from db.supabase_client import persist_scan_results | |
| persist_scan_results( | |
| uid=uid, | |
| connector_type="slack", | |
| db_name=request.channel_id, | |
| table_name="Messages", | |
| response=response, | |
| column_count=len(df.columns), | |
| rows_scanned=rows_scanned | |
| ) | |
| return JSONResponse(content=response) | |
| except Exception as e: | |
| raise HTTPException(status_code=500, detail=f"Slack connection failed: {str(e)}") | |
| async def connect_gmail(request: GmailRequest, authorization: str = Header(None)): | |
| """Connect to Gmail, fetch messages, scan, persist to Supabase.""" | |
| from api import _uid_from_auth | |
| uid = _uid_from_auth(authorization) | |
| if not uid: | |
| raise HTTPException(status_code=401, detail="Authentication required to scan") | |
| try: | |
| df = classifier.get_gmail_data(request.credentials, request.auth_type, num_emails=10) | |
| if df.empty: | |
| raise ValueError("No emails found or invalid credentials.") | |
| text_sample = df.head(10).to_string() | |
| response = format_pii_response(df, df, text_sample) | |
| rows_scanned = len(df) | |
| response["rows_scanned"] = rows_scanned | |
| masked_df = classifier.mask_dataframe(df.head(50)) | |
| response["data"] = masked_df.fillna("").to_dict(orient="records") | |
| # Supabase Persistence | |
| from db.supabase_client import persist_scan_results | |
| persist_scan_results( | |
| uid=uid, | |
| connector_type="gmail", | |
| db_name="Gmail Inbox", | |
| table_name="Emails", | |
| response=response, | |
| column_count=len(df.columns), | |
| rows_scanned=rows_scanned | |
| ) | |
| return JSONResponse(content=response) | |
| except Exception as e: | |
| raise HTTPException(status_code=500, detail=f"Gmail connection failed: {str(e)}") | |
| async def connect_zendesk(request: ZendeskRequest, authorization: str = Header(None)): | |
| """Connect to Zendesk, fetch tickets, scan, persist to Supabase.""" | |
| from api import _uid_from_auth | |
| uid = _uid_from_auth(authorization) | |
| if not uid: | |
| raise HTTPException(status_code=401, detail="Authentication required to scan") | |
| try: | |
| df = classifier.get_zendesk_tickets(request.subdomain, request.email, request.api_token) | |
| if df.empty: | |
| raise ValueError("No tickets found or invalid credentials.") | |
| text_sample = df.head(10).to_string() | |
| response = format_pii_response(df, df, text_sample) | |
| rows_scanned = len(df) | |
| response["rows_scanned"] = rows_scanned | |
| masked_df = classifier.mask_dataframe(df.head(50)) | |
| response["data"] = masked_df.fillna("").to_dict(orient="records") | |
| # Supabase Persistence | |
| from db.supabase_client import persist_scan_results | |
| persist_scan_results( | |
| uid=uid, | |
| connector_type="zendesk", | |
| db_name=f"{request.subdomain}.zendesk.com", | |
| table_name="Tickets", | |
| response=response, | |
| column_count=len(df.columns), | |
| rows_scanned=rows_scanned | |
| ) | |
| return JSONResponse(content=response) | |
| except Exception as e: | |
| raise HTTPException(status_code=500, detail=f"Zendesk connection failed: {str(e)}") | |
| async def connect_salesforce(request: SalesforceRequest, authorization: str = Header(None)): | |
| """Connect to Salesforce, fetch records, scan, persist to Supabase.""" | |
| from api import _uid_from_auth | |
| uid = _uid_from_auth(authorization) | |
| if not uid: | |
| raise HTTPException(status_code=401, detail="Authentication required to scan") | |
| try: | |
| df = classifier.get_salesforce_records(request.instance_url, request.access_token) | |
| if df.empty: | |
| raise ValueError("No records found or invalid credentials.") | |
| text_sample = df.head(10).to_string() | |
| response = format_pii_response(df, df, text_sample) | |
| rows_scanned = len(df) | |
| response["rows_scanned"] = rows_scanned | |
| masked_df = classifier.mask_dataframe(df.head(50)) | |
| response["data"] = masked_df.fillna("").to_dict(orient="records") | |
| # Supabase Persistence | |
| from db.supabase_client import persist_scan_results | |
| persist_scan_results( | |
| uid=uid, | |
| connector_type="salesforce", | |
| db_name=request.instance_url.replace("https://", ""), | |
| table_name="Records", | |
| response=response, | |
| column_count=len(df.columns), | |
| rows_scanned=rows_scanned | |
| ) | |
| return JSONResponse(content=response) | |
| except Exception as e: | |
| raise HTTPException(status_code=500, detail=f"Salesforce connection failed: {str(e)}") | |
| async def list_glue_databases(request: GlueCredentials): | |
| try: | |
| databases = classifier.list_glue_databases(request.access_key, request.secret_key, request.region) | |
| return {"databases": databases} | |
| except Exception as e: | |
| raise HTTPException(status_code=500, detail=str(e)) | |
| async def list_glue_tables(request: GlueListTablesRequest): | |
| try: | |
| tables = classifier.list_glue_tables(request.access_key, request.secret_key, request.region, request.database_name) | |
| return {"tables": tables} | |
| except Exception as e: | |
| raise HTTPException(status_code=500, detail=str(e)) | |
| async def connect_glue(request: GlueScanRequest, authorization: str = Header(None)): | |
| """Connect to AWS Glue, fetch table schema, scan for PII, and persist metadata.""" | |
| from api import _uid_from_auth | |
| uid = _uid_from_auth(authorization) | |
| if not uid: | |
| raise HTTPException(status_code=401, detail="Authentication required to scan") | |
| try: | |
| df = classifier.get_glue_table_schema(request.access_key, request.secret_key, request.region, request.database_name, request.table_name) | |
| if df.empty: | |
| raise ValueError("No schema found or invalid credentials.") | |
| text_sample = df.head(10).to_string() | |
| response = format_pii_response(df, df, text_sample) | |
| rows_scanned = len(df) | |
| response["rows_scanned"] = rows_scanned | |
| masked_df = classifier.mask_dataframe(df.head(50)) | |
| response["data"] = masked_df.fillna("").to_dict(orient="records") | |
| # Supabase Persistence | |
| from db.supabase_client import persist_scan_results | |
| persist_scan_results( | |
| uid=uid, | |
| connector_type="glue", | |
| db_name=request.database_name, | |
| table_name=request.table_name, | |
| response=response, | |
| column_count=len(df.columns), | |
| rows_scanned=rows_scanned | |
| ) | |
| return JSONResponse(content=response) | |
| except Exception as e: | |
| raise HTTPException(status_code=500, detail=f"AWS Glue connection failed: {str(e)}") | |
| async def connect_aws_rds(request: AwsRdsConnectionRequest, authorization: str = Header(None)): | |
| """Connect to AWS RDS, scan table, persist to Supabase, return analysis.""" | |
| from api import _uid_from_auth | |
| uid = _uid_from_auth(authorization) | |
| if not uid: | |
| raise HTTPException(status_code=401, detail="Authentication required to scan") | |
| try: | |
| df = classifier.get_aws_rds_data( | |
| request.engine, request.host, request.port, request.database, | |
| request.user, request.password, request.table | |
| ) | |
| text_sample = df.head(10).to_string() | |
| response = format_pii_response(df, df, text_sample) | |
| rows_scanned = len(df) | |
| response["rows_scanned"] = rows_scanned | |
| masked_df = classifier.mask_dataframe(df.head(50)) | |
| response["data"] = masked_df.fillna("").to_dict(orient="records") | |
| # ── Persist to Supabase universal schema ────────────────────────── | |
| from db.supabase_client import persist_scan_results | |
| persist_scan_results( | |
| uid=uid, | |
| connector_type="aws-rds", | |
| db_name=request.database, | |
| table_name=request.table, | |
| response=response, | |
| column_count=len(df.columns), | |
| rows_scanned=rows_scanned | |
| ) | |
| # ───────────────────────────────────────────────────────────────── | |
| return JSONResponse(content=response) | |
| except Exception as e: | |
| raise HTTPException(status_code=500, detail=f"AWS RDS connection failed: {str(e)}") | |
| async def connect_dynamodb(request: DynamoDbConnectionRequest, authorization: str = Header(None)): | |
| """Connect to DynamoDB, scan table, persist to Supabase, return analysis.""" | |
| uid = _uid_from_auth(authorization) | |
| try: | |
| df = classifier.get_dynamodb_data( | |
| request.region, request.access_key, request.secret_key, request.table | |
| ) | |
| text_sample = df.head(10).to_string() | |
| response = format_pii_response(df, df, text_sample) | |
| # Bug #8 fix: rows_scanned was undefined — define it before persist call | |
| rows_scanned = len(df) | |
| response["rows_scanned"] = rows_scanned | |
| masked_df = classifier.mask_dataframe(df.head(50)) | |
| response["data"] = masked_df.fillna("").to_dict(orient="records") | |
| # ── Persist to Supabase universal schema ────────────────────────── | |
| from db.supabase_client import persist_scan_results | |
| persist_scan_results( | |
| uid=uid, | |
| connector_type="dynamodb", | |
| db_name=request.region, | |
| table_name=request.table, | |
| response=response, | |
| column_count=len(df.columns), | |
| rows_scanned=rows_scanned | |
| ) | |
| # ───────────────────────────────────────────────────────────────── | |
| return JSONResponse(content=response) | |
| except Exception as e: | |
| raise HTTPException(status_code=500, detail=f"DynamoDB connection failed: {str(e)}") | |
| async def connect_aws_rds_list_tables(request: AwsRdsConnectionRequest): | |
| try: | |
| tables = classifier.list_aws_rds_tables( | |
| request.engine, request.host, request.port, request.database, | |
| request.user, request.password | |
| ) | |
| return JSONResponse(content={"tables": tables}) | |
| except Exception as e: | |
| raise HTTPException(status_code=500, detail=f"Failed to list tables: {str(e)}") | |
| async def connect_dynamodb_list_tables(request: DynamoDbConnectionRequest): | |
| try: | |
| tables = classifier.list_dynamodb_tables( | |
| request.region, request.access_key, request.secret_key | |
| ) | |
| return JSONResponse(content={"tables": tables}) | |
| except Exception as e: | |
| raise HTTPException(status_code=500, detail=f"Failed to list tables: {str(e)}") | |
| async def connect_metadata_scan(request: MetadataScanRequest, authorization: str = Header(None)): | |
| """ | |
| Connect to a DB (Postgres/MySQL), run a metadata-only scan via MetadataScanner, | |
| and persist results to Supabase returning the CatalogResponse shape. | |
| """ | |
| uid = _uid_from_auth(authorization) | |
| try: | |
| from classifier_manager.metadata_scanner import MetadataScanner | |
| from db.supabase_client import ( | |
| create_scan_session, finish_scan_session, | |
| upsert_db_registry_entry, upsert_db_scan_state, | |
| get_file_catalog, get_scan_sessions | |
| ) | |
| scanner = MetadataScanner() | |
| db_name = request.database | |
| table_name = request.table | |
| connector_type = request.connector_type.lower() | |
| # Postgres schema defaults to public for deterministic file_ids | |
| schema_name = "public" if connector_type in ["postgresql", "aws-rds"] else "" | |
| # 1. Fetch schema metadata via unified SchemaInspector (TICKET-5) | |
| # DynamoDB has no SQL schema — falls back to native handler | |
| try: | |
| if connector_type in ("postgresql", "mysql", "mariadb", "aws-rds"): | |
| rds_engine = getattr(request, "engine", None) | |
| engine = _schema_inspector.build_engine( | |
| connector_type, request.host, request.port, db_name, | |
| request.user, request.password, | |
| **({} if not rds_engine else {"rds_engine": rds_engine}) | |
| ) | |
| raw_cols = _schema_inspector.get_columns(engine, table_name) | |
| # MetadataScanner expects list[dict] with column_name / data_type keys | |
| columns_metadata = raw_cols | |
| elif connector_type == "dynamodb": | |
| db_name = request.region | |
| columns_metadata = classifier.dynamodb_handler.get_schema_metadata( | |
| request.region, request.access_key, request.secret_key, table_name | |
| ) | |
| else: | |
| raise HTTPException( | |
| status_code=400, | |
| detail=f"Unsupported connector_type for metadata scan: {connector_type}" | |
| ) | |
| except HTTPException: | |
| raise | |
| except Exception as e: | |
| err_str = str(e) | |
| if "missing read permissions" in err_str.lower(): | |
| raise HTTPException(status_code=403, detail=err_str) | |
| raise HTTPException(status_code=500, detail=err_str) | |
| # 2. Run MetadataScanner | |
| scan_results = scanner.scan_schema(columns_metadata) | |
| flagged_columns = scan_results.get("flagged_columns", []) | |
| # 3. Generate LLM Explanations for metadata scan | |
| if flagged_columns and classifier: | |
| rows = [{"PII Type": c["matched_rule"], "matched_rule": c["matched_rule"], "contributing_models": ["MetadataScanner"]} for c in flagged_columns] | |
| classifier.generate_explanations(rows) | |
| for c, r in zip(flagged_columns, rows): | |
| c["llm_explanation"] = r.get("llm_explanation", "") | |
| # 4. Persist to Supabase | |
| try: | |
| session_id = create_scan_session(uid, connector_type, "METADATA_ONLY") | |
| # Upsert Folder row | |
| folder_metadata = {"scan_mode": "metadata_only"} | |
| upsert_db_registry_entry(uid, connector_type, db_name, "", None, True, folder_metadata) | |
| # Upsert Table row | |
| if table_name: | |
| table_classification = "SENSITIVE" if len(flagged_columns) > 0 else "NON-SENSITIVE" | |
| upsert_db_registry_entry(uid, connector_type, db_name, schema_name, table_name, False, scan_results) | |
| upsert_db_scan_state(uid, connector_type, db_name, schema_name, table_name, | |
| table_classification, "METADATA_ONLY", session_id) | |
| # Update folder scan state | |
| upsert_db_scan_state(uid, connector_type, db_name, "", None, | |
| table_classification, "METADATA_ONLY", session_id) | |
| total_pii = len(flagged_columns) | |
| finish_scan_session(session_id, 1, total_pii) | |
| except Exception as persist_err: | |
| print(f"[WARN] Metadata scan persistence failed: {persist_err}") | |
| # 4. Return standard CatalogResponse format | |
| files = get_file_catalog(uid, connector_type) | |
| sessions = get_scan_sessions(uid, connector_type) | |
| for f in files: | |
| f["connector_type"] = connector_type | |
| files.sort(key=lambda f: (not f.get("is_folder", False), f.get("file_name", "").lower())) | |
| sessions.sort(key=lambda s: s.get("triggered_at", ""), reverse=True) | |
| completed_sessions = [s for s in sessions if s.get("status") == "completed"] | |
| last_session = completed_sessions[0] if completed_sessions else None | |
| response_content = jsonable_encoder({"files": files, "last_session": last_session, "sessions": sessions}) | |
| return JSONResponse(content=response_content) | |
| except HTTPException: | |
| raise | |
| except Exception as e: | |
| raise HTTPException(status_code=500, detail=f"Metadata scan failed: {str(e)}") | |
| # ==================== CLOUD STORAGE - AWS S3 ==================== | |
| async def list_s3_buckets(request: S3ConnectionRequest): | |
| """List S3 buckets""" | |
| try: | |
| buckets = classifier.get_s3_buckets(request.access_key, request.secret_key, request.region) | |
| return JSONResponse(content={"buckets": buckets}) | |
| except Exception as e: | |
| raise HTTPException(status_code=500, detail=f"S3 connection failed: {str(e)}") | |
| async def list_s3_files(request: S3ConnectionRequest): | |
| """List files in S3 bucket""" | |
| try: | |
| if not request.bucket: | |
| raise HTTPException(status_code=400, detail="Bucket name is required") | |
| files = classifier.get_s3_files( | |
| request.access_key, request.secret_key, request.region, request.bucket | |
| ) | |
| return JSONResponse(content={"files": files}) | |
| except Exception as e: | |
| raise HTTPException(status_code=500, detail=f"Failed to list S3 files: {str(e)}") | |
| async def scan_s3_file(request: S3ConnectionRequest): | |
| """Download and scan S3 file""" | |
| try: | |
| if not request.bucket or not request.file_key: | |
| raise HTTPException(status_code=400, detail="Bucket and file_key are required") | |
| content = classifier.download_s3_file( | |
| request.access_key, request.secret_key, request.region, | |
| request.bucket, request.file_key | |
| ) | |
| # Assume CSV for now | |
| df = pd.read_csv(io.BytesIO(content)) | |
| text_sample = df.head(10).to_string() | |
| response = format_pii_response(df, df, text_sample) | |
| masked_df = classifier.mask_dataframe(df.head(50)) | |
| response["data"] = masked_df.fillna("").to_dict(orient="records") | |
| return JSONResponse(content=response) | |
| except Exception as e: | |
| raise HTTPException(status_code=500, detail=f"S3 scan failed: {str(e)}") | |
| # ==================== CLOUD STORAGE - AZURE ==================== | |
| async def list_azure_containers(request: AzureConnectionRequest): | |
| """List Azure containers""" | |
| try: | |
| containers = classifier.get_azure_containers(request.connection_string) | |
| return JSONResponse(content={"containers": containers}) | |
| except Exception as e: | |
| raise HTTPException(status_code=500, detail=f"Azure connection failed: {str(e)}") | |
| async def list_azure_blobs(request: AzureConnectionRequest): | |
| """List blobs in Azure container""" | |
| try: | |
| if not request.container: | |
| raise HTTPException(status_code=400, detail="Container name is required") | |
| blobs = classifier.get_azure_blobs(request.connection_string, request.container) | |
| return JSONResponse(content={"blobs": blobs}) | |
| except Exception as e: | |
| raise HTTPException(status_code=500, detail=f"Failed to list blobs: {str(e)}") | |
| async def scan_azure_blob(request: AzureConnectionRequest): | |
| """Download and scan Azure blob""" | |
| try: | |
| if not request.container or not request.blob: | |
| raise HTTPException(status_code=400, detail="Container and blob are required") | |
| content = classifier.download_azure_blob( | |
| request.connection_string, request.container, request.blob | |
| ) | |
| df = pd.read_csv(io.BytesIO(content)) | |
| text_sample = df.head(10).to_string() | |
| response = format_pii_response(df, df, text_sample) | |
| masked_df = classifier.mask_dataframe(df.head(50)) | |
| response["data"] = masked_df.fillna("").to_dict(orient="records") | |
| return JSONResponse(content=response) | |
| except Exception as e: | |
| raise HTTPException(status_code=500, detail=f"Azure scan failed: {str(e)}") | |
| # ==================== CLOUD STORAGE - GCS ==================== | |
| async def list_gcs_buckets(request: GCSConnectionRequest): | |
| """List GCS buckets""" | |
| try: | |
| buckets = classifier.get_gcs_buckets(request.credentials) | |
| return JSONResponse(content={"buckets": buckets}) | |
| except Exception as e: | |
| raise HTTPException(status_code=500, detail=f"GCS connection failed: {str(e)}") | |
| async def list_gcs_files(request: GCSConnectionRequest): | |
| """List files in GCS bucket""" | |
| try: | |
| if not request.bucket: | |
| raise HTTPException(status_code=400, detail="Bucket name is required") | |
| files = classifier.get_gcs_files(request.credentials, request.bucket) | |
| return JSONResponse(content={"files": files}) | |
| except Exception as e: | |
| raise HTTPException(status_code=500, detail=f"Failed to list GCS files: {str(e)}") | |
| async def scan_gcs_file(request: GCSConnectionRequest): | |
| """Download and scan GCS file""" | |
| try: | |
| if not request.bucket or not request.file_name: | |
| raise HTTPException(status_code=400, detail="Bucket and file_name are required") | |
| content = classifier.download_gcs_file( | |
| request.credentials, request.bucket, request.file_name | |
| ) | |
| df = pd.read_csv(io.BytesIO(content)) | |
| text_sample = df.head(10).to_string() | |
| response = format_pii_response(df, df, text_sample) | |
| masked_df = classifier.mask_dataframe(df.head(50)) | |
| response["data"] = masked_df.fillna("").to_dict(orient="records") | |
| return JSONResponse(content=response) | |
| except Exception as e: | |
| raise HTTPException(status_code=500, detail=f"GCS scan failed: {str(e)}") | |
| # ==================== CLOUD STORAGE - GOOGLE DRIVE ==================== | |
| async def list_drive_files(request: GoogleDriveRequest): | |
| """List Google Drive files""" | |
| try: | |
| files = classifier.get_google_drive_files(request.credentials) | |
| return JSONResponse(content={"files": files}) | |
| except Exception as e: | |
| raise HTTPException(status_code=500, detail=f"Google Drive connection failed: {str(e)}") | |
| async def scan_drive_file(request: GoogleDriveRequest): | |
| """Download and scan Google Drive file""" | |
| try: | |
| if not request.file_id or not request.mime_type: | |
| raise HTTPException(status_code=400, detail="file_id and mime_type are required") | |
| content = classifier.download_drive_file( | |
| request.file_id, request.mime_type, request.credentials | |
| ) | |
| if isinstance(content, bytes): | |
| try: | |
| text = content.decode('utf-8') | |
| df = pd.DataFrame({"Content": [text]}) | |
| response = format_pii_response(df, df, text) | |
| highlighted_df = classifier.scan_dataframe_with_html(df) | |
| response["data"] = highlighted_df.fillna("").to_dict(orient="records") | |
| return JSONResponse(content=response) | |
| except: | |
| raise HTTPException(status_code=400, detail="Binary file cannot be processed") | |
| except Exception as e: | |
| raise HTTPException(status_code=500, detail=f"Google Drive scan failed: {str(e)}") | |
| # ==================== ENTERPRISE CONNECTORS ==================== | |
| async def scan_gmail(file: UploadFile = File(...), num_emails: int = Form(10)): | |
| """Scan Gmail messages""" | |
| try: | |
| df = classifier.get_gmail_data(file.file, num_emails) | |
| if df.empty: | |
| raise HTTPException(status_code=400, detail="No emails fetched") | |
| text_sample = df.iloc[0]['Content'] | |
| response = format_pii_response(df, df, text_sample) | |
| masked_df = classifier.mask_dataframe(df) | |
| response["data"] = masked_df.fillna("").to_dict(orient="records") | |
| return JSONResponse(content=response) | |
| except Exception as e: | |
| raise HTTPException(status_code=500, detail=f"Gmail scan failed: {str(e)}") | |
| async def scan_slack(request: SlackRequest): | |
| """Scan Slack messages""" | |
| try: | |
| df = classifier.get_slack_messages(request.token, request.channel_id) | |
| if df.empty: | |
| raise HTTPException(status_code=400, detail="No messages found or authentication failed") | |
| text_sample = df.iloc[0]['Content'] | |
| response = format_pii_response(df, df, text_sample) | |
| masked_df = classifier.mask_dataframe(df) | |
| response["data"] = masked_df.fillna("").to_dict(orient="records") | |
| return JSONResponse(content=response) | |
| except Exception as e: | |
| raise HTTPException(status_code=500, detail=f"Slack scan failed: {str(e)}") | |
| async def scan_confluence(request: ConfluenceRequest): | |
| """Scan Confluence page""" | |
| try: | |
| df = classifier.get_confluence_page( | |
| request.url, request.username, request.token, request.page_id | |
| ) | |
| if df.empty: | |
| raise HTTPException(status_code=400, detail="Failed to fetch page") | |
| text_sample = df.iloc[0]['Content'] | |
| response = format_pii_response(df, df, text_sample) | |
| highlighted_df = classifier.scan_dataframe_with_html(df) | |
| response["data"] = highlighted_df.fillna("").to_dict(orient="records") | |
| return JSONResponse(content=response) | |
| except Exception as e: | |
| raise HTTPException(status_code=500, detail=f"Confluence scan failed: {str(e)}") | |
| # ==================== EMAIL FUNCTIONALITY ==================== | |
| async def send_welcome(request: WelcomeEmailRequest): | |
| """ | |
| Send a welcome email to a new user. | |
| This endpoint is called by the frontend after a user submits the contact form. | |
| """ | |
| try: | |
| # Send the welcome email | |
| success = send_welcome_email(request.name, request.email) | |
| if success: | |
| return JSONResponse(content={ | |
| "success": True, | |
| "message": f"Welcome email sent to {request.email}" | |
| }) | |
| else: | |
| raise HTTPException( | |
| status_code=500, | |
| detail="Failed to send welcome email. SMTP configuration may be missing." | |
| ) | |
| except Exception as e: | |
| raise HTTPException( | |
| status_code=500, | |
| detail=f"Email sending failed: {str(e)}" | |
| ) | |
| # ==================== EVALUATOR ENDPOINTS ==================== | |
| # Full model catalogue served to the MODEL LAB frontend | |
| EVALUATOR_MODEL_CATALOGUE = [ | |
| {"key": "regex", "label": "Regex Engine", "hf_id": "deterministic", "type": "Rule-based", "params": "—", "f1_benchmark": 1.0, "lazy": False, "description": "Deterministic regex patterns for EMAIL, PHONE, SSN, CC, IP, URL, MAC."}, | |
| {"key": "nltk", "label": "NLTK Chunker", "hf_id": "nltk", "type": "Statistical", "params": "—", "f1_benchmark": 0.0, "lazy": False, "description": "NLTK ne_chunk: PERSON → FIRST_NAME, GPE → LOCATION."}, | |
| {"key": "spacy", "label": "SpaCy LG", "hf_id": "en_core_web_lg", "type": "Statistical", "params": "685M", "f1_benchmark": 0.0, "lazy": False, "description": "SpaCy en_core_web_lg NER model."}, | |
| {"key": "presidio", "label": "MS Presidio", "hf_id": "microsoft/presidio-analyzer", "type": "Rule+ML", "params": "—", "f1_benchmark": 0.0, "lazy": False, "description": "Microsoft Presidio enterprise PII analyzer."}, | |
| {"key": "gliner", "label": "GLiNER Small", "hf_id": "urchade/gliner_small-v2.1", "type": "GLiNER", "params": "small", "f1_benchmark": 0.850, "lazy": False, "description": "Zero-shot GLiNER small model."}, | |
| {"key": "deberta", "label": "DeBERTa PII", "hf_id": "lakshyakh93/deberta-large-finetuned-pii", "type": "NER", "params": "86M", "f1_benchmark": 0.920, "lazy": False, "description": "Kaggle-winning DeBERTa V3 fine-tuned for PII."}, | |
| {"key": "pasteproof", "label": "Pasteproof v2", "hf_id": "joneauxedgar/pasteproof-pii-detector-v2", "type": "NER", "params": "149M", "f1_benchmark": 0.970, "lazy": True, "description": "ModernBERT 149M fine-tuned PII detector."}, | |
| {"key": "piiranha", "label": "Piiranha v1", "hf_id": "iiiorg/piiranha-v1-detect-personal-information", "type": "NER", "params": "86M", "f1_benchmark": 0.931, "lazy": True, "description": "DeBERTa-based PII detector."}, | |
| {"key": "nvidia_gliner","label": "NVIDIA GLiNER", "hf_id": "nvidia/gliner-PII-0.1", "type": "GLiNER", "params": "570M", "f1_benchmark": 0.870, "lazy": True, "description": "NVIDIA GLiNER with 37-label PII vocabulary."}, | |
| {"key": "mmbert", "label": "mmbert32k", "hf_id": "llm-semantic-router/mmbert32k-pii-detector-merged", "type": "NER", "params": "307M", "f1_benchmark": 0.969, "lazy": True, "description": "ModernBERT 32k-context PII detector."}, | |
| {"key": "nerguard", "label": "NerGuard-0.3B", "hf_id": "exdsgift/NerGuard-0.3B", "type": "NER", "params": "300M", "f1_benchmark": 0.996, "lazy": True, "description": "mDeBERTa 300M — highest F1 in the registry."}, | |
| {"key": "gliner_large", "label": "GLiNER PII Large", "hf_id": "knowledgator/gliner-pii-large-v1.0", "type": "GLiNER", "params": "large","f1_benchmark": 0.833, "lazy": True, "description": "GLiNER large architecture fine-tuned for PII."}, | |
| ] | |
| class EvaluatorScanRequest(BaseModel): | |
| text: str = Field(..., description="Document text to scan") | |
| gt_spans: List[Dict[str, Any]] = Field(default=[], description="Ground truth spans from /api/evaluator/parse") | |
| model_keys: List[str] = Field(default=["regex", "spacy", "deberta"], description="Model keys to run") | |
| conf_threshold: float = Field(default=0.5, ge=0.0, le=1.0) | |
| entropy_threshold: float = Field(default=4.5, ge=3.0, le=6.0) | |
| async def evaluator_models(): | |
| """Return the full model catalogue for the MODEL LAB page.""" | |
| return JSONResponse(content={"models": EVALUATOR_MODEL_CATALOGUE}) | |
| async def evaluator_parse( | |
| file: UploadFile = File(...), | |
| format: str = Form("auto"), | |
| doc_index: int = Form(0), | |
| schema: str = Form(""), | |
| ): | |
| """ | |
| Parse a labeled (or unlabeled) dataset file. | |
| Returns the document text, ground-truth spans, format detected, and doc count. | |
| format: 'auto' | 'bigcode' | 'nemotron' | 'csv_spans' | 'json_spans' | 'unlabeled' | |
| schema: JSON string e.g. '{"text_col":"text","spans_col":"spans"}' | |
| """ | |
| try: | |
| content = await file.read() | |
| schema_dict = json.loads(schema) if schema.strip() else None | |
| text, gt_spans, has_gt, doc_count, fmt_detected = detect_and_parse( | |
| content, file.filename or "", doc_index, schema_dict, format | |
| ) | |
| return JSONResponse(content={ | |
| "text": text, | |
| "gt_spans": gt_spans, | |
| "has_gt": has_gt, | |
| "format_detected": fmt_detected, | |
| "doc_count": doc_count, | |
| "char_count": len(text), | |
| }) | |
| except Exception as e: | |
| raise HTTPException(status_code=500, detail=f"Parse error: {str(e)}") | |
| async def evaluator_scan(request: EvaluatorScanRequest): | |
| """ | |
| Run selected models on text in parallel, compare against GT spans. | |
| Returns partial results — if a model times out or errors, other models | |
| still return results. Each model entry includes timed_out + error fields. | |
| """ | |
| try: | |
| t0 = time.time() | |
| has_gt = bool(request.gt_spans) | |
| # Parallel scan — new shape: {key: {detections, error, timed_out}} | |
| raw_results = await asyncio.get_event_loop().run_in_executor( | |
| None, | |
| lambda: classifier.scan_with_models(request.text, request.model_keys), | |
| ) | |
| per_model: Dict[str, Any] = {} | |
| for model_key, payload in raw_results.items(): | |
| predictions = payload["detections"] | |
| timed_out = payload["timed_out"] | |
| model_error = payload["error"] | |
| # Normalise canonical label | |
| for p in predictions: | |
| p["canonical"] = norm_model_out(p.get("label", "")) | |
| if has_gt and not timed_out: | |
| comparison = compare_spans(predictions, request.gt_spans, model_key) | |
| metrics = compute_metrics(comparison) | |
| coverage = get_label_coverage(request.gt_spans, model_key) | |
| failures = analyse_failures(comparison, request.text) | |
| else: | |
| comparison = {"TP": [], "FP": [], "FN": []} | |
| metrics = [] | |
| coverage = {"in_scope": [], "out_of_scope": []} | |
| failures = {"missed": [], "false_positives": []} | |
| per_model[model_key] = { | |
| "predictions": predictions, | |
| "comparison": comparison, | |
| "metrics": metrics, | |
| "coverage": coverage, | |
| "failures": failures, | |
| "timed_out": timed_out, | |
| "error": model_error, | |
| } | |
| elapsed = round(time.time() - t0, 2) | |
| any_timeout = any(v["timed_out"] for v in per_model.values()) | |
| return JSONResponse(content={ | |
| "per_model": per_model, | |
| "has_gt": has_gt, | |
| "elapsed": elapsed, | |
| "any_timeout": any_timeout, | |
| }) | |
| except Exception as e: | |
| raise HTTPException(status_code=500, detail=f"Evaluator scan error: {str(e)}") | |
| async def evaluator_scan_stream(request: EvaluatorScanRequest): | |
| """ | |
| SSE streaming version of /api/evaluator/scan. | |
| Emits one Server-Sent Event per model as results arrive, so the frontend | |
| can show live per-model progress rather than waiting for all models. | |
| Event format (JSON): | |
| data: {"model_key": ..., "done": false, "status": "running"} | |
| data: {"model_key": ..., "done": true, "payload": {...}, "elapsed": ...} | |
| data: {"done": true, "all_complete": true} | |
| """ | |
| import concurrent.futures | |
| has_gt = bool(request.gt_spans) | |
| loop = asyncio.get_event_loop() | |
| async def event_generator(): | |
| t0 = time.time() | |
| always_on = {"regex", "nltk", "spacy", "presidio", "gliner", "deberta"} | |
| def _run_one_key(key: str): | |
| payload = classifier.scan_with_models(request.text, [key]) | |
| return key, payload.get(key, {"detections": [], "error": "No result", "timed_out": False}) | |
| with concurrent.futures.ThreadPoolExecutor(max_workers=4) as pool: | |
| future_map = {pool.submit(_run_one_key, key): key for key in request.model_keys} | |
| # Yield start event for each model | |
| for key in request.model_keys: | |
| yield f'data: {{"model_key": "{key}", "done": false, "status": "running"}}\n\n' | |
| for f in concurrent.futures.as_completed(future_map, timeout=110): | |
| try: | |
| key, raw = await loop.run_in_executor(None, f.result) | |
| predictions = raw["detections"] | |
| for p in predictions: | |
| p["canonical"] = norm_model_out(p.get("label", "")) | |
| if has_gt and not raw["timed_out"]: | |
| comparison = compare_spans(predictions, request.gt_spans, key) | |
| metrics = compute_metrics(comparison) | |
| coverage = get_label_coverage(request.gt_spans, key) | |
| failures = analyse_failures(comparison, request.text) | |
| else: | |
| comparison = {"TP": [], "FP": [], "FN": []} | |
| metrics = []; coverage = {"in_scope": [], "out_of_scope": []}; failures = {"missed": [], "false_positives": []} | |
| result_payload = { | |
| "predictions": predictions, "comparison": comparison, | |
| "metrics": metrics, "coverage": coverage, "failures": failures, | |
| "timed_out": raw["timed_out"], "error": raw["error"], | |
| } | |
| event = json.dumps({"model_key": key, "done": True, "payload": result_payload, "elapsed": round(time.time() - t0, 2)}) | |
| yield f'data: {event}\n\n' | |
| except Exception as e: | |
| err_event = json.dumps({"model_key": future_map[f], "done": True, "error": str(e), "timed_out": True}) | |
| yield f'data: {err_event}\n\n' | |
| yield f'data: {{"done": true, "all_complete": true, "total_elapsed": {round(time.time() - t0, 2)}}}\n\n' | |
| return StreamingResponse( | |
| event_generator(), | |
| media_type="text/event-stream", | |
| headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"}, | |
| ) | |
| async def evaluator_batch( | |
| file: UploadFile = File(...), | |
| format: str = Form("nemotron"), | |
| n_docs: int = Form(50), | |
| model_keys: str = Form("deberta"), | |
| conf_threshold: float = Form(0.5), | |
| entropy_threshold: float = Form(4.5), | |
| ): | |
| """ | |
| Batch evaluation over N documents from a labeled file. | |
| Returns aggregate metrics + per-doc breakdown for each selected model. | |
| model_keys: comma-separated string e.g. 'deberta,piiranha' | |
| """ | |
| try: | |
| content = await file.read() | |
| keys = [k.strip() for k in model_keys.split(",") if k.strip()] | |
| n_docs = max(1, min(n_docs, 1000)) | |
| agg: Dict[str, Dict[str, int]] = {k: {"TP": 0, "FP": 0, "FN": 0} for k in keys} | |
| per_doc: List[Dict[str, Any]] = [] | |
| for i in range(n_docs): | |
| try: | |
| text, gt_spans, has_gt, doc_count, _ = detect_and_parse( | |
| content, file.filename or "", i, None, format | |
| ) | |
| if i >= doc_count: | |
| break | |
| if not has_gt: | |
| continue | |
| model_results = classifier.scan_with_models(text, keys) | |
| doc_entry: Dict[str, Any] = {"doc_index": i} | |
| for key, predictions in model_results.items(): | |
| for p in predictions: | |
| p["canonical"] = norm_model_out(p.get("label", "")) | |
| comp = compare_spans(predictions, gt_spans, key) | |
| tp, fp, fn = len(comp["TP"]), len(comp["FP"]), len(comp["FN"]) | |
| agg[key]["TP"] += tp | |
| agg[key]["FP"] += fp | |
| agg[key]["FN"] += fn | |
| prec = tp / (tp + fp) if (tp + fp) else 0.0 | |
| rec = tp / (tp + fn) if (tp + fn) else 0.0 | |
| f1 = 2 * prec * rec / (prec + rec) if (prec + rec) else 0.0 | |
| doc_entry[key] = { | |
| "f1": round(f1, 4), "precision": round(prec, 4), | |
| "recall": round(rec, 4), "tp": tp, "fp": fp, "fn": fn, | |
| } | |
| per_doc.append(doc_entry) | |
| except Exception as doc_err: | |
| per_doc.append({"doc_index": i, "error": str(doc_err)}) | |
| # Aggregate metrics per model | |
| aggregate: Dict[str, Any] = {} | |
| for key, c in agg.items(): | |
| tp, fp, fn = c["TP"], c["FP"], c["FN"] | |
| p = tp / (tp + fp) if (tp + fp) else 0.0 | |
| r = tp / (tp + fn) if (tp + fn) else 0.0 | |
| f = 2 * p * r / (p + r) if (p + r) else 0.0 | |
| aggregate[key] = { | |
| "f1": round(f, 4), "precision": round(p, 4), | |
| "recall": round(r, 4), "tp": tp, "fp": fp, "fn": fn, | |
| } | |
| return JSONResponse(content={ | |
| "aggregate": aggregate, | |
| "per_doc": per_doc, | |
| "n_docs_evaluated": len(per_doc), | |
| }) | |
| except Exception as e: | |
| raise HTTPException(status_code=500, detail=f"Batch eval error: {str(e)}") | |
| # ==================== AUTH ENDPOINTS ==================== | |
| from db.supabase_auth import auth_sign_up, auth_sign_in, auth_get_user, auth_sign_out | |
| def _uid_from_auth(authorization: Optional[str]) -> str: | |
| """Extract and verify the Supabase GoTrue user_id from a Bearer token. | |
| Raises HTTP 401 if the token is missing or invalid.""" | |
| if not authorization or not authorization.startswith("Bearer "): | |
| raise HTTPException(status_code=401, detail="Missing or invalid Authorization header") | |
| token = authorization.split(" ", 1)[1] | |
| user_data = auth_get_user(access_token=token) | |
| # auth_get_user returns { "id": "...", "email": "...", ... } or { "error": "..." } | |
| if "error" in user_data or "id" not in user_data: | |
| raise HTTPException(status_code=401, detail=user_data.get("error", "Invalid token")) | |
| return user_data["id"] | |
| async def register(req: RegisterRequest): | |
| """Register a new Sense user via Supabase GoTrue.""" | |
| # Validate minimum password length before hitting Supabase | |
| if len(req.password) < 6: | |
| raise HTTPException(status_code=422, detail="Password must be at least 6 characters") | |
| # Delegate to GoTrue | |
| result = auth_sign_up(email=req.email, password=req.password, full_name=req.name) | |
| # Seed credits row — non-fatal if it fails | |
| try: | |
| if "id" in result: | |
| from db.credits_manager import seed_user_credits | |
| seed_user_credits(result["id"]) | |
| except Exception as e: | |
| print(f"[WARN] Credit seeding failed: {e}") | |
| # Send welcome email — failure must NOT block registration | |
| try: | |
| send_welcome_email(name=req.name, email=req.email) | |
| except Exception: | |
| pass | |
| return JSONResponse(content=result) | |
| async def login(req: LoginRequest): | |
| """Authenticate a Sense user via Supabase GoTrue.""" | |
| result = auth_sign_in(email=req.email, password=req.password) | |
| # Restore credits if it's Sunday (idempotent) — non-fatal | |
| try: | |
| if "id" in result: | |
| from db.credits_manager import maybe_restore_credits | |
| maybe_restore_credits(result["id"]) | |
| except Exception as e: | |
| print(f"[WARN] Credit restore check failed: {e}") | |
| return JSONResponse(content=result) | |
| async def get_me(authorization: str = Header(None)): | |
| """Return the authenticated user's profile. Requires Bearer token.""" | |
| if not authorization or not authorization.startswith("Bearer "): | |
| raise HTTPException(status_code=401, detail="Missing or invalid Authorization header") | |
| token = authorization.split(" ", 1)[1] | |
| user = auth_get_user(access_token=token) | |
| return JSONResponse(content=user) | |
| async def logout(authorization: str = Header(None)): | |
| """Sign out the user from Supabase GoTrue (server-side token invalidation).""" | |
| if authorization and authorization.startswith("Bearer "): | |
| token = authorization.split(" ", 1)[1] | |
| auth_sign_out(access_token=token) | |
| return JSONResponse(content={"message": "Signed out"}) | |
| # ==================== CREDIT ENDPOINTS ==================== | |
| async def get_credits_endpoint(authorization: str = Header(None)): | |
| """Return the authenticated user's credit balance.""" | |
| user_id = _uid_from_auth(authorization) | |
| from db.credits_manager import get_credits | |
| data = get_credits(user_id) | |
| return JSONResponse(content=data) | |
| async def deduct_credits_endpoint(authorization: str = Header(None), amount: int = 1): | |
| """Atomically deduct `amount` credits. Returns 402 when out of credits.""" | |
| user_id = _uid_from_auth(authorization) | |
| from db.credits_manager import deduct_credit | |
| result = deduct_credit(user_id, amount) | |
| if not result["success"]: | |
| return JSONResponse(status_code=402, content=result) | |
| return JSONResponse(content=result) | |
| async def get_profile_stats_endpoint(authorization: str = Header(None)): | |
| """Return aggregated scan statistics + credits for the profile dashboard.""" | |
| user_id = _uid_from_auth(authorization) | |
| from db.credits_manager import get_profile_stats | |
| stats = get_profile_stats(user_id) | |
| return JSONResponse(content=stats) | |
| # ==================== HEALTH CHECK ==================== | |
| async def debug_db(): | |
| """Debug: show recent scan sessions and registry rows (uses unified schema).""" | |
| try: | |
| from db.supabase_client import _get_conn | |
| conn = _get_conn() | |
| cur = conn.cursor() | |
| cur.execute( | |
| "SELECT id, uid, connector_type, scan_type, scan_mode, status, " | |
| "triggered_at, files_scanned, total_pii_found " | |
| "FROM sense_scan_sessions ORDER BY triggered_at DESC LIMIT 10" | |
| ) | |
| sessions = [dict(r) for r in cur.fetchall()] | |
| cur.execute( | |
| "SELECT file_id, file_name, connector_type, is_folder " | |
| "FROM sense_file_registry ORDER BY first_seen_at DESC LIMIT 20" | |
| ) | |
| registry = [dict(r) for r in cur.fetchall()] | |
| cur.close() | |
| conn.close() | |
| # Stringify datetimes | |
| for row in sessions + registry: | |
| for k, v in row.items(): | |
| if hasattr(v, 'isoformat'): | |
| row[k] = v.isoformat() | |
| return {"sessions": sessions, "registry": registry} | |
| except Exception as e: | |
| return {"error": str(e)} | |
| async def root(): | |
| """API health check""" | |
| return { | |
| "message": "Segmento Sense API", | |
| "status": "operational", | |
| "version": "1.0.0" | |
| } | |
| async def health_check(): | |
| """Detailed health check""" | |
| return { | |
| "status": "healthy", | |
| "classifiers": { | |
| "regex": True, | |
| "nltk": True, | |
| "spacy": True, | |
| "presidio": True, | |
| "gliner": True, | |
| "deberta": True | |
| } | |
| } | |
| # ==================== CONNECTOR CATALOG ENDPOINTS ==================== | |
| async def get_file_catalog_endpoint(connector_type: str, authorization: str = Header(None)): | |
| uid = _uid_from_auth(authorization) | |
| try: | |
| from db.supabase_client import get_file_catalog, get_scan_sessions | |
| files = get_file_catalog(uid, connector_type) | |
| sessions = get_scan_sessions(uid, connector_type) | |
| completed_sessions = [s for s in sessions if s.get("status") == "completed"] | |
| last_session = completed_sessions[0] if completed_sessions else None | |
| return {"files": files, "last_session": last_session, "sessions": sessions} | |
| except Exception as e: | |
| raise HTTPException(status_code=500, detail=str(e)) | |
| async def get_db_catalog( | |
| connector_type: Optional[str] = None, | |
| db_name: Optional[str] = None, | |
| authorization: str = Header(None) | |
| ): | |
| """ | |
| Unified catalog read endpoint (TICKET-5). | |
| Returns sense_file_registry + sense_file_scan_state rows for given user. | |
| Optional filters: connector_type, db_name. | |
| All scan modes (FULL_LOAD, SAMPLING, METADATA_ONLY, etc.) appear here. | |
| """ | |
| uid = _uid_from_auth(authorization) | |
| try: | |
| from db.supabase_client import get_file_catalog, get_scan_sessions | |
| DB_CONNECTOR_TYPES = ["postgresql", "mysql", "mongodb", "mariadb"] | |
| types_to_query = [connector_type] if connector_type else DB_CONNECTOR_TYPES | |
| all_files: list = [] | |
| all_sessions: list = [] | |
| for ct in types_to_query: | |
| # Pass db_name filter to get_file_catalog (TICKET-5) | |
| files = get_file_catalog(uid, ct, db_name=db_name) | |
| sessions = get_scan_sessions(uid, ct) | |
| for f in files: | |
| f["connector_type"] = ct | |
| all_files.extend(files) | |
| all_sessions.extend(sessions) | |
| all_files.sort(key=lambda f: (not f.get("is_folder", False), f.get("file_name", "").lower())) | |
| all_sessions.sort(key=lambda s: s.get("triggered_at", ""), reverse=True) | |
| completed_sessions = [s for s in all_sessions if s.get("status") == "completed"] | |
| last_session = completed_sessions[0] if completed_sessions else None | |
| return {"files": all_files, "last_session": last_session, "sessions": all_sessions} | |
| except Exception as e: | |
| raise HTTPException(status_code=500, detail=str(e)) | |
| async def get_scan_sessions_endpoint(connector_type: str, authorization: str = Header(None)): | |
| uid = _uid_from_auth(authorization) | |
| try: | |
| from db.supabase_client import get_scan_sessions | |
| sessions = get_scan_sessions(uid, connector_type) | |
| return {"sessions": sessions} | |
| except Exception as e: | |
| raise HTTPException(status_code=500, detail=str(e)) | |
| # NOTE: /api/scan/async, /api/scan/status/{job_id}, /api/catalog/{db_name} | |
| # have been removed in TICKET-5. | |
| # Use /api/connectors/metadata-scan (sync) + GET /db/catalog?db_name=... instead. | |
| # ==================== INCREMENTAL SCAN ENDPOINTS (TICKET-8) ==================== | |
| async def watermark_check( | |
| connector_type: str, | |
| db_name: str, | |
| table_name: str, | |
| host: str, | |
| port: str, | |
| user: str, | |
| password: str, | |
| authorization: str = Header(None) | |
| ): | |
| """ | |
| Check whether incremental scan is possible for a given table. | |
| Returns: | |
| { supported: true, watermark_column: str, last_value: str | null } | |
| { supported: false, reason: str } | |
| """ | |
| uid = _uid_from_auth(authorization) | |
| try: | |
| ct = connector_type.lower() | |
| if ct not in ("postgresql", "mysql", "mariadb", "aws-rds"): | |
| return {"supported": False, "reason": f"{ct} does not support watermark-based incremental scan."} | |
| engine = _schema_inspector.build_engine(ct, host, port, db_name, user, password) | |
| wm_col = _schema_inspector.detect_watermark_column(engine, table_name) | |
| if not wm_col: | |
| return { | |
| "supported": False, | |
| "reason": "No suitable watermark column found (updated_at, created_at, or auto-increment PK required)." | |
| } | |
| # Build the canonical file_id to look up existing watermark | |
| schema_name = "public" if ct in ("postgresql", "aws-rds") else "" | |
| from db.supabase_client import _build_db_file_id, check_watermark | |
| file_id = _build_db_file_id(db_name, schema_name, table_name) | |
| _, last_value = check_watermark(uid, ct, file_id) | |
| return { | |
| "supported": True, | |
| "watermark_column": wm_col, | |
| "last_value": last_value, # None means first-time (no prior watermark) | |
| } | |
| except Exception as e: | |
| raise HTTPException(status_code=500, detail=str(e)) | |
| async def incremental_scan( | |
| db_type: str, | |
| request: DatabaseConnectionRequest, | |
| authorization: str = Header(None) | |
| ): | |
| """ | |
| Watermark-based incremental scan (TICKET-8). | |
| 1. Detect watermark column. | |
| 2. Read last stored watermark from sense_file_scan_state. | |
| 3. Fetch only rows WHERE watermark_col > last_value. | |
| 4. Run PII scan on new rows. | |
| 5. Persist results + advance watermark. | |
| """ | |
| uid = _uid_from_auth(authorization) | |
| ct = db_type.lower() | |
| if ct not in ("postgresql", "mysql", "mariadb"): | |
| raise HTTPException(status_code=400, detail=f"Incremental scan not supported for: {ct}") | |
| try: | |
| from sqlalchemy import text | |
| from db.supabase_client import ( | |
| persist_scan_results, check_watermark, update_watermark, | |
| _build_db_file_id | |
| ) | |
| schema_name = "public" if ct == "postgresql" else "" | |
| file_id = _build_db_file_id(request.database, schema_name, request.table) | |
| # 1. Detect watermark column | |
| engine = _schema_inspector.build_engine( | |
| ct, request.host, request.port, request.database, request.user, request.password | |
| ) | |
| wm_col = _schema_inspector.detect_watermark_column(engine, request.table) | |
| if not wm_col: | |
| raise HTTPException( | |
| status_code=400, | |
| detail=( | |
| f"Table '{request.table}' has no suitable watermark column " | |
| "(needs updated_at, created_at, or auto-increment PK). " | |
| "Use Full Load or Sampling scan instead." | |
| ) | |
| ) | |
| # 2. Read last watermark | |
| _, last_value = check_watermark(uid, ct, file_id) | |
| # 3. Fetch incremental rows | |
| with engine.connect() as conn: | |
| if last_value is not None: | |
| # Quoted identifiers prevent SQL injection on col/table names | |
| query = text( | |
| f'SELECT * FROM "{request.table}" ' | |
| f'WHERE "{wm_col}" > :last_val ORDER BY "{wm_col}" ASC LIMIT 500' | |
| ) | |
| rows = conn.execute(query, {"last_val": last_value}).mappings().fetchall() | |
| else: | |
| # First incremental scan = full load (no prior watermark) | |
| query = text( | |
| f'SELECT * FROM "{request.table}" ORDER BY "{wm_col}" ASC LIMIT 500' | |
| ) | |
| rows = conn.execute(query).mappings().fetchall() | |
| if not rows: | |
| return JSONResponse(content={ | |
| "message": "No new rows since last scan.", | |
| "rows_scanned": 0, | |
| "total_pii_found": 0, | |
| "watermark_column": wm_col, | |
| "last_value": last_value, | |
| }) | |
| import pandas as pd | |
| df = pd.DataFrame([dict(r) for r in rows]) | |
| text_sample = df.head(10).to_string() | |
| response = format_pii_response(df, df, text_sample) | |
| rows_scanned = len(df) | |
| response["rows_scanned"] = rows_scanned | |
| masked_df = classifier.mask_dataframe(df.head(50)) | |
| response["data"] = masked_df.fillna("").to_dict(orient="records") | |
| # 4. Persist results | |
| persist_scan_results( | |
| uid=uid, | |
| connector_type=ct, | |
| db_name=request.database, | |
| table_name=request.table, | |
| response=response, | |
| column_count=len(df.columns), | |
| rows_scanned=rows_scanned, | |
| scan_mode="incremental", | |
| ) | |
| # 5. Advance watermark | |
| new_wm = _schema_inspector.get_max_watermark_value(engine, request.table, wm_col) | |
| if new_wm: | |
| update_watermark(uid, ct, file_id, wm_col, new_wm) | |
| response["watermark_column"] = wm_col | |
| response["last_value"] = new_wm | |
| return JSONResponse(content=response) | |
| except HTTPException: | |
| raise | |
| except Exception as e: | |
| raise HTTPException(status_code=500, detail=f"Incremental scan failed: {str(e)}") | |
| if __name__ == "__main__": | |
| import uvicorn | |
| uvicorn.run(app, host="0.0.0.0", port=7860) # HuggingFace Spaces default port |