Restructure backend into microservices, add envelope encryption, DM files, and message editing

backend/services/main/db.py

@@ -0,0 +1,154 @@
+import os
+from typing import Generator, Optional
+
+import time
+import logging
+from sqlalchemy import create_engine, event, text
+from sqlalchemy.engine import Engine
+from sqlalchemy.orm import sessionmaker, Session
+from sqlalchemy.pool import StaticPool
+from sqlalchemy.exc import OperationalError
+
+from .constants import DATABASE_URL
+
+logger = logging.getLogger(__name__)
+
+"""
+Universal database interface that provides identical behavior for PostgreSQL and SQLite.
+
+Features:
+- Auto-creates parent directory for SQLite files.
+- Applies SQLite pragmas (foreign_keys=ON, journal_mode=WAL) for improved compatibility.
+- Uses StaticPool for in-memory or file-based SQLite when appropriate.
+- Exposes `engine`, `SessionLocal`, `get_db` dependency, and `POOL_CONFIG`.
+"""
+
+# Ensure parent directory exists for SQLite file DBs
+def _ensure_sqlite_parent_dir(url: str) -> None:
+    if not url or not url.startswith("sqlite"):
+        return
+    # strip sqlite:/// prefix
+    path = url.replace("sqlite:///", "", 1)
+    parent = os.path.dirname(path)
+    if parent:
+        os.makedirs(parent, exist_ok=True)
+
+
+# Pool and engine configuration (tunable via env)
+POOL_SIZE = int(os.getenv("DB_POOL_SIZE", "20"))
+MAX_OVERFLOW = int(os.getenv("DB_MAX_OVERFLOW", "40"))
+POOL_RECYCLE = int(os.getenv("DB_POOL_RECYCLE", "1800"))
+POOL_TIMEOUT = int(os.getenv("DB_POOL_TIMEOUT", "30"))
+
+POOL_CONFIG = {
+    "pool_size": POOL_SIZE,
+    "max_overflow": MAX_OVERFLOW,
+    "pool_recycle": POOL_RECYCLE,
+    "pool_timeout": POOL_TIMEOUT,
+    "pool_pre_ping": True,
+}
+
+
+def get_engine(database_url: Optional[str] = None) -> Engine:
+    """
+    Create and return a SQLAlchemy Engine configured for the given database URL.
+    This function ensures SQLite-specific pragmas and connection args are applied.
+    Includes retry logic for database connection failures during startup.
+    """
+    url = database_url or DATABASE_URL
+    _ensure_sqlite_parent_dir(url)
+
+    # Retry database connection during startup (helps with Docker initialization timing)
+    if url.startswith("postgresql"):
+        max_retries = 15
+        retry_delay = 2
+
+        for attempt in range(max_retries):
+            try:
+                logger.info(f"Attempting database connection (attempt {attempt + 1}/{max_retries})...")
+                # Test the connection by creating engine and trying to connect
+                test_engine = create_engine(url, pool_size=1, max_overflow=0, pool_timeout=5, future=True)
+                with test_engine.connect() as conn:
+                    conn.execute(text("SELECT 1"))
+                test_engine.dispose()
+                logger.info("Database connection successful")
+                break
+            except OperationalError as e:
+                if attempt < max_retries - 1:
+                    logger.warning(f"Database connection failed (attempt {attempt + 1}): {e}")
+                    time.sleep(retry_delay)
+                else:
+                    logger.error(f"Database connection failed after {max_retries} attempts: {e}")
+                    raise
+            except Exception as e:
+                logger.error(f"Unexpected error during database connection: {e}")
+                raise
+
+    if url.startswith("sqlite"):
+        # For SQLite file-based DBs, use standard pooling but set connection timeout and pragmas.
+        # Use StaticPool only for in-memory SQLite.
+        in_memory = url in ("sqlite:///:memory:", "sqlite://")
+        connect_args = {"check_same_thread": False, "timeout": int(os.getenv("SQLITE_BUSY_TIMEOUT", "5"))}
+
+        if in_memory:
+            engine = create_engine(url, connect_args=connect_args, poolclass=StaticPool, future=True)
+        else:
+            engine = create_engine(url, connect_args=connect_args, future=True)
+
+        # Apply pragmas on connect for SQLite (foreign keys, WAL, busy_timeout)
+        @event.listens_for(engine, "connect")
+        def _sqlite_on_connect(dbapi_conn, connection_record):
+            try:
+                cursor = dbapi_conn.cursor()
+                cursor.execute("PRAGMA foreign_keys = ON")
+                cursor.execute("PRAGMA journal_mode = WAL")
+                # busy_timeout in milliseconds
+                busy_ms = int(os.getenv("SQLITE_BUSY_TIMEOUT_MS", "5000"))
+                cursor.execute(f"PRAGMA busy_timeout = {busy_ms}")
+                cursor.close()
+            except Exception:
+                # Best-effort; do not fail engine creation if pragmas cannot be set
+                pass
+
+        return engine
+
+    # Default for Postgres / MySQL etc. - use pool sizing from env
+    engine_kwargs = {
+        "pool_size": POOL_SIZE,
+        "max_overflow": MAX_OVERFLOW,
+        "pool_recycle": POOL_RECYCLE,
+        "pool_timeout": POOL_TIMEOUT,
+        "future": True,
+    }
+    return create_engine(url, **engine_kwargs)
+
+
+# Create global engine and session factory for convenient imports
+engine = get_engine()
+# Keep loaded attributes available after commit/close to avoid DetachedInstanceError
+SessionLocal = sessionmaker(class_=Session, autocommit=False, autoflush=False, bind=engine, expire_on_commit=False)
+
+
+def init_db(create_tables: bool = False, base_metadata=None) -> None:
+    """
+    Initialize the database. If `create_tables` is True and `base_metadata` is provided,
+    create all tables using the provided SQLAlchemy metadata.
+    """
+    if create_tables:
+        if base_metadata is None:
+            raise ValueError("base_metadata is required to create tables")
+        base_metadata.create_all(bind=engine)
+
+
+def get_db() -> Generator[Session, None, None]:
+    """
+    FastAPI dependency that yields a SQLAlchemy Session and ensures proper close().
+    """
+    db = SessionLocal()
+    try:
+        yield db
+    finally:
+        try:
+            db.close()
+        except Exception:
+            pass