""" Database abstraction layer supporting SQLite, PostgreSQL, and MySQL. Backend selection is based on DATABASE_URL environment variable: - postgres://... or postgresql://... -> PostgreSQL - mysql://... -> MySQL - Not set -> SQLite (default) """ import os import json import time import asyncio from pathlib import Path from typing import Dict, List, Any, Optional, Tuple, Set from abc import ABC, abstractmethod import aiosqlite # Schema version for migrations SCHEMA_VERSION = 1 # Define all columns that should exist in the accounts table # Format: (column_name, column_type_sqlite, column_type_postgres, column_type_mysql, default_value) ACCOUNTS_COLUMNS = [ ("id", "TEXT PRIMARY KEY", "TEXT PRIMARY KEY", "VARCHAR(255) PRIMARY KEY", None), ("label", "TEXT", "TEXT", "TEXT", None), ("clientId", "TEXT", "TEXT", "TEXT", None), ("clientSecret", "TEXT", "TEXT", "TEXT", None), ("refreshToken", "TEXT", "TEXT", "TEXT", None), ("accessToken", "TEXT", "TEXT", "TEXT", None), ("other", "TEXT", "TEXT", "TEXT", None), ("last_refresh_time", "TEXT", "TEXT", "TEXT", None), ("last_refresh_status", "TEXT", "TEXT", "TEXT", None), ("created_at", "TEXT", "TEXT", "TEXT", None), ("updated_at", "TEXT", "TEXT", "TEXT", None), ("enabled", "INTEGER DEFAULT 1", "INTEGER DEFAULT 1", "INT DEFAULT 1", "1"), ("error_count", "INTEGER DEFAULT 0", "INTEGER DEFAULT 0", "INT DEFAULT 0", "0"), ("success_count", "INTEGER DEFAULT 0", "INTEGER DEFAULT 0", "INT DEFAULT 0", "0"), ("expires_at", "TEXT", "TEXT", "TEXT", None), ] # Optional imports for other backends try: import asyncpg HAS_ASYNCPG = True except ImportError: HAS_ASYNCPG = False try: import aiomysql HAS_AIOMYSQL = True except ImportError: HAS_AIOMYSQL = False class DatabaseBackend(ABC): """Abstract base class for database backends.""" @abstractmethod async def initialize(self) -> None: """Initialize connection and ensure schema exists.""" pass @abstractmethod async def close(self) -> None: """Close database connections.""" pass @abstractmethod async def execute(self, query: str, params: tuple = ()) -> int: """Execute a query and return affected row count.""" pass @abstractmethod async def fetchone(self, query: str, params: tuple = ()) -> Optional[Dict[str, Any]]: """Fetch a single row as dict.""" pass @abstractmethod async def fetchall(self, query: str, params: tuple = ()) -> List[Dict[str, Any]]: """Fetch all rows as list of dicts.""" pass class SQLiteBackend(DatabaseBackend): """SQLite database backend using aiosqlite.""" def __init__(self, db_path: Path): self._db_path = db_path self._initialized = False self._conn: Optional[aiosqlite.Connection] = None async def _get_existing_columns(self) -> Set[str]: """Get existing column names from accounts table.""" try: async with self._conn.execute("PRAGMA table_info(accounts)") as cursor: rows = await cursor.fetchall() return {row[1] for row in rows} except Exception: return set() async def _migrate_schema(self) -> None: """Add missing columns to accounts table.""" existing_cols = await self._get_existing_columns() if not existing_cols: return # Table doesn't exist yet, will be created fresh for col_name, col_type, _, _, _ in ACCOUNTS_COLUMNS: if col_name not in existing_cols and "PRIMARY KEY" not in col_type: # Extract just the type without DEFAULT clause for ALTER TABLE base_type = col_type.split(" DEFAULT")[0].strip() try: await self._conn.execute(f"ALTER TABLE accounts ADD COLUMN {col_name} {base_type}") print(f"[DB Migration] Added column: {col_name}") except Exception as e: print(f"[DB Migration] Failed to add column {col_name}: {e}") async def initialize(self) -> None: if self._initialized: return self._db_path.parent.mkdir(parents=True, exist_ok=True) self._conn = await aiosqlite.connect(self._db_path) # Performance tuning PRAGMAs await self._conn.execute("PRAGMA journal_mode=WAL;") await self._conn.execute("PRAGMA synchronous = NORMAL;") await self._conn.execute("PRAGMA cache_size = -65536; -- 64MB") await self._conn.execute("PRAGMA temp_store = MEMORY;") # Build CREATE TABLE statement from schema definition columns_sql = ", ".join([f"{col[0]} {col[1]}" for col in ACCOUNTS_COLUMNS]) await self._conn.execute(f""" CREATE TABLE IF NOT EXISTS accounts ({columns_sql}) """) # Run migrations for existing tables await self._migrate_schema() # Create indexes for performance await self._conn.execute("CREATE INDEX IF NOT EXISTS idx_accounts_enabled ON accounts (enabled);") await self._conn.execute("CREATE INDEX IF NOT EXISTS idx_accounts_created_at ON accounts (created_at);") await self._conn.execute("CREATE INDEX IF NOT EXISTS idx_accounts_success_count ON accounts (success_count);") await self._conn.commit() self._initialized = True async def close(self) -> None: if self._conn: await self._conn.close() self._conn = None self._initialized = False async def execute(self, query: str, params: tuple = ()) -> int: cursor = await self._conn.execute(query, params) await self._conn.commit() return cursor.rowcount async def fetchone(self, query: str, params: tuple = ()) -> Optional[Dict[str, Any]]: self._conn.row_factory = aiosqlite.Row async with self._conn.execute(query, params) as cursor: row = await cursor.fetchone() return dict(row) if row else None async def fetchall(self, query: str, params: tuple = ()) -> List[Dict[str, Any]]: self._conn.row_factory = aiosqlite.Row async with self._conn.execute(query, params) as cursor: rows = await cursor.fetchall() return [dict(row) for row in rows] class PostgresBackend(DatabaseBackend): """PostgreSQL database backend using asyncpg.""" def __init__(self, dsn: str): self._dsn = dsn self._pool: "Optional[asyncpg.pool.Pool]" = None self._initialized = False async def _get_existing_columns(self, conn) -> Set[str]: """Get existing column names from accounts table.""" try: rows = await conn.fetch(""" SELECT column_name FROM information_schema.columns WHERE table_name = 'accounts' """) return {row['column_name'] for row in rows} except Exception: return set() async def _migrate_schema(self, conn) -> None: """Add missing columns to accounts table.""" existing_cols = await self._get_existing_columns(conn) if not existing_cols: return # Table doesn't exist yet for col_name, _, col_type, _, _ in ACCOUNTS_COLUMNS: if col_name not in existing_cols and "PRIMARY KEY" not in col_type: base_type = col_type.split(" DEFAULT")[0].strip() try: await conn.execute(f"ALTER TABLE accounts ADD COLUMN IF NOT EXISTS {col_name} {base_type}") print(f"[DB Migration] Added column: {col_name}") except Exception as e: print(f"[DB Migration] Failed to add column {col_name}: {e}") async def initialize(self) -> None: if not HAS_ASYNCPG: raise ImportError("asyncpg is required for PostgreSQL support. Install with: pip install asyncpg") self._pool = await asyncpg.create_pool(dsn=self._dsn, min_size=1, max_size=20) async with self._pool.acquire() as conn: # Build CREATE TABLE statement from schema definition columns_sql = ", ".join([f"{col[0]} {col[2]}" for col in ACCOUNTS_COLUMNS]) await conn.execute(f""" CREATE TABLE IF NOT EXISTS accounts ({columns_sql}) """) # Run migrations await self._migrate_schema(conn) self._initialized = True async def close(self) -> None: if self._pool: await self._pool.close() self._pool = None self._initialized = False def _convert_placeholders(self, query: str) -> str: """Convert ? placeholders to $1, $2, etc.""" result = [] param_num = 0 i = 0 while i < len(query): if query[i] == '?': param_num += 1 result.append(f'${param_num}') else: result.append(query[i]) i += 1 return ''.join(result) async def execute(self, query: str, params: tuple = ()) -> int: pg_query = self._convert_placeholders(query) async with self._pool.acquire() as conn: result = await conn.execute(pg_query, *params) # asyncpg returns string like "UPDATE 1" try: return int(result.split()[-1]) except (ValueError, IndexError): return 0 async def fetchone(self, query: str, params: tuple = ()) -> Optional[Dict[str, Any]]: pg_query = self._convert_placeholders(query) async with self._pool.acquire() as conn: row = await conn.fetchrow(pg_query, *params) return dict(row) if row else None async def fetchall(self, query: str, params: tuple = ()) -> List[Dict[str, Any]]: pg_query = self._convert_placeholders(query) async with self._pool.acquire() as conn: rows = await conn.fetch(pg_query, *params) return [dict(row) for row in rows] class MySQLBackend(DatabaseBackend): """MySQL database backend using aiomysql.""" def __init__(self, dsn: str): self._dsn = dsn self._pool = None self._initialized = False self._config = self._parse_dsn(dsn) def _parse_dsn(self, dsn: str) -> Dict[str, Any]: """Parse MySQL DSN into connection parameters.""" # mysql://user:password@host:port/database from urllib.parse import urlparse, parse_qs parsed = urlparse(dsn) config = { 'host': parsed.hostname or 'localhost', 'port': parsed.port or 3306, 'user': parsed.username or 'root', 'password': parsed.password or '', 'db': parsed.path.lstrip('/') if parsed.path else 'test', } # Handle SSL query = parse_qs(parsed.query) if 'ssl' in query or 'sslmode' in query or 'ssl-mode' in query: config['ssl'] = True return config async def _get_existing_columns(self, cur) -> Set[str]: """Get existing column names from accounts table.""" try: await cur.execute(f"DESCRIBE accounts") rows = await cur.fetchall() return {row[0] if isinstance(row, tuple) else row['Field'] for row in rows} except Exception: return set() async def _migrate_schema(self, cur) -> None: """Add missing columns to accounts table.""" existing_cols = await self._get_existing_columns(cur) if not existing_cols: return # Table doesn't exist yet for col_name, _, _, col_type, _ in ACCOUNTS_COLUMNS: if col_name not in existing_cols and "PRIMARY KEY" not in col_type: base_type = col_type.split(" DEFAULT")[0].strip() try: await cur.execute(f"ALTER TABLE accounts ADD COLUMN {col_name} {base_type}") print(f"[DB Migration] Added column: {col_name}") except Exception as e: # Column might already exist if "Duplicate column" not in str(e): print(f"[DB Migration] Failed to add column {col_name}: {e}") async def initialize(self) -> None: if not HAS_AIOMYSQL: raise ImportError("aiomysql is required for MySQL support. Install with: pip install aiomysql") self._pool = await aiomysql.create_pool( host=self._config['host'], port=self._config['port'], user=self._config['user'], password=self._config['password'], db=self._config['db'], minsize=1, maxsize=20, autocommit=True ) async with self._pool.acquire() as conn: async with conn.cursor() as cur: # Build CREATE TABLE statement from schema definition columns_sql = ", ".join([f"{col[0]} {col[3]}" for col in ACCOUNTS_COLUMNS]) await cur.execute(f""" CREATE TABLE IF NOT EXISTS accounts ({columns_sql}) """) # Run migrations await self._migrate_schema(cur) self._initialized = True async def close(self) -> None: if self._pool: self._pool.close() await self._pool.wait_closed() self._pool = None self._initialized = False def _convert_placeholders(self, query: str) -> str: """Convert ? placeholders to %s for MySQL.""" return query.replace('?', '%s') async def execute(self, query: str, params: tuple = ()) -> int: mysql_query = self._convert_placeholders(query) async with self._pool.acquire() as conn: async with conn.cursor() as cur: await cur.execute(mysql_query, params) return cur.rowcount async def fetchone(self, query: str, params: tuple = ()) -> Optional[Dict[str, Any]]: mysql_query = self._convert_placeholders(query) async with self._pool.acquire() as conn: async with conn.cursor(aiomysql.DictCursor) as cur: await cur.execute(mysql_query, params) return await cur.fetchone() async def fetchall(self, query: str, params: tuple = ()) -> List[Dict[str, Any]]: mysql_query = self._convert_placeholders(query) async with self._pool.acquire() as conn: async with conn.cursor(aiomysql.DictCursor) as cur: await cur.execute(mysql_query, params) return await cur.fetchall() # Global database instance _db: Optional[DatabaseBackend] = None def get_database_backend() -> DatabaseBackend: """Get the configured database backend based on DATABASE_URL.""" global _db if _db is not None: return _db database_url = os.getenv('DATABASE_URL', '').strip() if database_url.startswith(('postgres://', 'postgresql://')): # Fix common postgres:// to postgresql:// for asyncpg dsn = database_url.replace('postgres://', 'postgresql://', 1) if database_url.startswith('postgres://') else database_url _db = PostgresBackend(dsn) print(f"[DB] Using PostgreSQL backend") elif database_url.startswith('mysql://'): _db = MySQLBackend(database_url) print(f"[DB] Using MySQL backend") else: # Default to SQLite base_dir = Path(__file__).resolve().parent db_path = base_dir / "data.sqlite3" _db = SQLiteBackend(db_path) print(f"[DB] Using SQLite backend: {db_path}") return _db async def init_db() -> DatabaseBackend: """Initialize and return the database backend.""" db = get_database_backend() await db.initialize() return db async def close_db() -> None: """Close the database backend.""" global _db if _db: await _db.close() _db = None # Helper functions for common operations def row_to_dict(row: Dict[str, Any]) -> Optional[Dict[str, Any]]: """Convert a database row to dict with JSON parsing for 'other' field.""" if row is None: return None d = dict(row) if d.get("other"): try: d["other"] = json.loads(d["other"]) except Exception: pass # normalize enabled to bool if "enabled" in d and d["enabled"] is not None: try: d["enabled"] = bool(int(d["enabled"])) except Exception: d["enabled"] = bool(d["enabled"]) return d