""" Database abstraction layer supporting both SQLite (local) and BigQuery (cloud). Provides unified interface for CRUD operations across deployment modes. """ import sqlite3 import json import logging from abc import ABC, abstractmethod from datetime import datetime, timedelta from typing import List, Dict, Any, Optional, Union from contextlib import contextmanager from pathlib import Path from config import config from models import ProcessedEvent, SatelliteResult, Claim, INDIAN_STATES # Configure logging logger = logging.getLogger(__name__) class DatabaseInterface(ABC): """Abstract interface for database operations""" @abstractmethod async def initialize(self) -> bool: """Initialize database connection and schema""" pass @abstractmethod async def insert_event(self, event: ProcessedEvent) -> bool: """Insert a processed event into the database""" pass @abstractmethod async def get_event(self, event_id: str) -> Optional[ProcessedEvent]: """Retrieve a specific event by ID""" pass @abstractmethod async def get_events_by_region(self, region: str, limit: int = 100) -> List[ProcessedEvent]: """Get events for a specific Indian state/region""" pass @abstractmethod async def get_events_by_timerange(self, start_time: datetime, end_time: datetime) -> List[ProcessedEvent]: """Get events within a specific time range""" pass @abstractmethod async def get_heatmap_data(self, hours_back: int = 24) -> Dict[str, Dict[str, Any]]: """Get aggregated data for heatmap visualization""" pass @abstractmethod async def delete_old_events(self, days_old: int = 30) -> int: """Delete events older than specified days""" pass @abstractmethod async def get_stats(self) -> Dict[str, Any]: """Get database statistics""" pass class SQLiteDatabase(DatabaseInterface): """SQLite implementation for local development mode""" def __init__(self, db_path: str): self.db_path = Path(db_path) self.db_path.parent.mkdir(parents=True, exist_ok=True) logger.info(f"Initializing SQLite database at {self.db_path}") @contextmanager def get_connection(self): """Context manager for database connections""" conn = sqlite3.connect(self.db_path) conn.row_factory = sqlite3.Row # Enable dict-like access try: yield conn finally: conn.close() async def initialize(self) -> bool: """Initialize SQLite database with schema""" try: with self.get_connection() as conn: # Create events table conn.execute(""" CREATE TABLE IF NOT EXISTS events ( event_id TEXT PRIMARY KEY, source TEXT NOT NULL, original_text TEXT NOT NULL, timestamp DATETIME NOT NULL, lang TEXT, region_hint TEXT, lat REAL, lon REAL, entities TEXT, -- JSON array virality_score REAL, satellite_data TEXT, -- JSON object claims_data TEXT, -- JSON array processing_metadata TEXT, -- JSON object created_at DATETIME DEFAULT CURRENT_TIMESTAMP ) """) # Create indexes for performance conn.execute("CREATE INDEX IF NOT EXISTS idx_events_timestamp ON events(timestamp)") conn.execute("CREATE INDEX IF NOT EXISTS idx_events_region ON events(region_hint)") conn.execute("CREATE INDEX IF NOT EXISTS idx_events_source ON events(source)") conn.execute("CREATE INDEX IF NOT EXISTS idx_events_created_at ON events(created_at)") # Create aggregation table for faster heatmap queries conn.execute(""" CREATE TABLE IF NOT EXISTS region_stats ( region TEXT PRIMARY KEY, event_count INTEGER DEFAULT 0, avg_virality_score REAL DEFAULT 0.0, avg_reality_score REAL DEFAULT 0.0, last_updated DATETIME DEFAULT CURRENT_TIMESTAMP ) """) conn.commit() logger.info("SQLite database initialized successfully") return True except Exception as e: logger.error(f"Failed to initialize SQLite database: {e}") return False async def insert_event(self, event: ProcessedEvent) -> bool: """Insert a processed event into SQLite""" try: with self.get_connection() as conn: # Serialize complex fields to JSON entities_json = json.dumps(event.entities) satellite_json = json.dumps(event.satellite.to_dict() if event.satellite else {}) claims_json = json.dumps([claim.to_dict() for claim in event.claims]) metadata_json = json.dumps(event.processing_metadata) conn.execute(""" INSERT OR REPLACE INTO events ( event_id, source, original_text, timestamp, lang, region_hint, lat, lon, entities, virality_score, satellite_data, claims_data, processing_metadata, created_at ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) """, ( event.event_id, event.source.value, event.original_text, event.timestamp.isoformat(), event.lang.value, event.region_hint, event.lat, event.lon, entities_json, event.virality_score, satellite_json, claims_json, metadata_json, event.created_at.isoformat() )) # Update region statistics await self._update_region_stats(conn, event) conn.commit() logger.debug(f"Inserted event {event.event_id} into SQLite") return True except Exception as e: logger.error(f"Failed to insert event into SQLite: {e}") return False async def _update_region_stats(self, conn, event: ProcessedEvent): """Update aggregated region statistics""" if not event.region_hint: return reality_score = event.get_reality_score() # Get current stats cursor = conn.execute( "SELECT event_count, avg_virality_score, avg_reality_score FROM region_stats WHERE region = ?", (event.region_hint,) ) row = cursor.fetchone() if row: # Update existing stats count = row[0] + 1 new_virality = ((row[1] * row[0]) + event.virality_score) / count new_reality = ((row[2] * row[0]) + reality_score) / count conn.execute(""" UPDATE region_stats SET event_count = ?, avg_virality_score = ?, avg_reality_score = ?, last_updated = ? WHERE region = ? """, (count, new_virality, new_reality, datetime.utcnow().isoformat(), event.region_hint)) else: # Insert new stats conn.execute(""" INSERT INTO region_stats (region, event_count, avg_virality_score, avg_reality_score, last_updated) VALUES (?, ?, ?, ?, ?) """, (event.region_hint, 1, event.virality_score, reality_score, datetime.utcnow().isoformat())) async def get_event(self, event_id: str) -> Optional[ProcessedEvent]: """Retrieve a specific event by ID from SQLite""" try: with self.get_connection() as conn: cursor = conn.execute("SELECT * FROM events WHERE event_id = ?", (event_id,)) row = cursor.fetchone() if row: return self._row_to_event(row) return None except Exception as e: logger.error(f"Failed to get event from SQLite: {e}") return None async def get_events_by_region(self, region: str, limit: int = 100) -> List[ProcessedEvent]: """Get events for a specific region from SQLite""" try: with self.get_connection() as conn: cursor = conn.execute(""" SELECT * FROM events WHERE region_hint = ? ORDER BY timestamp DESC LIMIT ? """, (region, limit)) events = [] for row in cursor.fetchall(): event = self._row_to_event(row) if event: events.append(event) return events except Exception as e: logger.error(f"Failed to get events by region from SQLite: {e}") return [] async def get_events_by_timerange(self, start_time: datetime, end_time: datetime) -> List[ProcessedEvent]: """Get events within a time range from SQLite""" try: with self.get_connection() as conn: cursor = conn.execute(""" SELECT * FROM events WHERE timestamp BETWEEN ? AND ? ORDER BY timestamp DESC """, (start_time.isoformat(), end_time.isoformat())) events = [] for row in cursor.fetchall(): event = self._row_to_event(row) if event: events.append(event) return events except Exception as e: logger.error(f"Failed to get events by timerange from SQLite: {e}") return [] async def get_heatmap_data(self, hours_back: int = 24) -> Dict[str, Dict[str, Any]]: """Get aggregated heatmap data from SQLite""" try: cutoff_time = datetime.utcnow() - timedelta(hours=hours_back) with self.get_connection() as conn: cursor = conn.execute(""" SELECT region_hint, COUNT(*) as event_count, AVG(virality_score) as avg_virality, AVG(CASE WHEN satellite_data != '{}' AND satellite_data != '' THEN JSON_EXTRACT(satellite_data, '$.reality_score') ELSE 0.5 END) as avg_reality FROM events WHERE timestamp >= ? AND region_hint IS NOT NULL AND region_hint != '' GROUP BY region_hint """, (cutoff_time.isoformat(),)) heatmap_data = {} for row in cursor.fetchall(): region = row[0] if region and region.lower() in INDIAN_STATES: heatmap_data[region] = { "intensity": min(1.0, row[1] / 10.0), # Normalize to 0-1 "event_count": row[1], "avg_virality_score": round(row[2] or 0.0, 3), "avg_reality_score": round(row[3] or 0.5, 3), "misinformation_risk": round((row[2] or 0.0) * (1.0 - (row[3] or 0.5)), 3) } return heatmap_data except Exception as e: logger.error(f"Failed to get heatmap data from SQLite: {e}") return {} async def delete_old_events(self, days_old: int = 30) -> int: """Delete old events from SQLite""" try: cutoff_date = datetime.utcnow() - timedelta(days=days_old) with self.get_connection() as conn: cursor = conn.execute( "DELETE FROM events WHERE created_at < ?", (cutoff_date.isoformat(),) ) deleted_count = cursor.rowcount conn.commit() logger.info(f"Deleted {deleted_count} old events from SQLite") return deleted_count except Exception as e: logger.error(f"Failed to delete old events from SQLite: {e}") return 0 async def get_stats(self) -> Dict[str, Any]: """Get database statistics from SQLite""" try: with self.get_connection() as conn: # Total events cursor = conn.execute("SELECT COUNT(*) FROM events") total_events = cursor.fetchone()[0] # Events by source cursor = conn.execute("SELECT source, COUNT(*) FROM events GROUP BY source") events_by_source = dict(cursor.fetchall()) # Recent events (last 24 hours) cutoff_time = datetime.utcnow() - timedelta(hours=24) cursor = conn.execute( "SELECT COUNT(*) FROM events WHERE timestamp >= ?", (cutoff_time.isoformat(),) ) recent_events = cursor.fetchone()[0] # Database file size db_size = self.db_path.stat().st_size if self.db_path.exists() else 0 return { "total_events": total_events, "events_by_source": events_by_source, "recent_events_24h": recent_events, "database_size_bytes": db_size, "database_type": "sqlite" } except Exception as e: logger.error(f"Failed to get stats from SQLite: {e}") return {} def _row_to_event(self, row) -> Optional[ProcessedEvent]: """Convert SQLite row to ProcessedEvent object""" try: # Parse JSON fields entities = json.loads(row["entities"]) if row["entities"] else [] satellite_data = json.loads(row["satellite_data"]) if row["satellite_data"] else {} claims_data = json.loads(row["claims_data"]) if row["claims_data"] else [] metadata = json.loads(row["processing_metadata"]) if row["processing_metadata"] else {} # Create satellite result satellite = SatelliteResult.from_dict(satellite_data) if satellite_data else None # Create claims claims = [Claim.from_dict(claim_data) for claim_data in claims_data] # Parse timestamps timestamp = datetime.fromisoformat(row["timestamp"]) created_at = datetime.fromisoformat(row["created_at"]) return ProcessedEvent( event_id=row["event_id"], source=row["source"], original_text=row["original_text"], timestamp=timestamp, lang=row["lang"], region_hint=row["region_hint"], lat=row["lat"] or 0.0, lon=row["lon"] or 0.0, entities=entities, virality_score=row["virality_score"] or 0.0, satellite=satellite, claims=claims, processing_metadata=metadata, created_at=created_at ) except Exception as e: logger.error(f"Failed to convert row to event: {e}") return None class DatabaseManager: """Factory class for creating proper database implementation""" @staticmethod def create_database() -> DatabaseInterface: """Create database instance strictly for SQLite (cost-free)""" db_config = config.get_database_config() # Extract path from SQLite URL db_path = db_config.url.replace("sqlite:///", "") return SQLiteDatabase(db_path) # Global database instance database = DatabaseManager.create_database()