database.py

import json
import mysql.connector
from mysql.connector import Error, pooling
from contextlib import contextmanager
from typing import Optional, List, Dict, Any

from config import DB_HOST, DB_PORT, DB_USER, DB_PASSWORD, DB_NAME

# ============================================================
# Connection Pool
# ============================================================
# A pool of 5 reusable connections avoids the overhead of
# opening a new TCP connection on every request.
# ============================================================

_pool: Optional[pooling.MySQLConnectionPool] = None


def _get_pool() -> pooling.MySQLConnectionPool:
    """Lazily create and return the shared connection pool."""
    global _pool
    if _pool is None:
        _pool = pooling.MySQLConnectionPool(
            pool_name="biosentinel_pool",
            pool_size=5,
            host=DB_HOST,
            port=DB_PORT,
            user=DB_USER,
            password=DB_PASSWORD,
            database=DB_NAME,
            autocommit=False,
            charset="utf8mb4",
            collation="utf8mb4_unicode_ci",
        )
    return _pool


@contextmanager
def get_db():
    """
    Context manager that yields a live MySQL connection from the pool.
    Commits on success, rolls back on exception, always returns the
    connection to the pool when the block exits.

    Usage:
        with get_db() as conn:
            cursor = conn.cursor(dictionary=True)
            cursor.execute(...)
    """
    conn = _get_pool().get_connection()
    try:
        yield conn
        conn.commit()
    except Exception:
        conn.rollback()
        raise
    finally:
        conn.close()


# ============================================================
# Helper Functions
# ============================================================

def db_execute(query: str, values: Optional[tuple] = None) -> bool:
    """Run INSERT / UPDATE / DELETE. Returns True on success."""
    with get_db() as conn:
        cursor = conn.cursor()
        cursor.execute(query, values)
    return True


def db_insert(query: str, values: Optional[tuple] = None) -> Optional[int]:
    """
    Run an INSERT and return the last-inserted row ID.
    Returns None on failure.
    """
    with get_db() as conn:
        cursor = conn.cursor()
        cursor.execute(query, values)
        return cursor.lastrowid


def db_fetch_one(query: str, values: Optional[tuple] = None) -> Optional[Dict[str, Any]]:
    """Return a single row as a dict, or None."""
    with get_db() as conn:
        cursor = conn.cursor(dictionary=True)
        cursor.execute(query, values)
        return cursor.fetchone()


def db_fetch_all(query: str, values: Optional[tuple] = None) -> List[Dict[str, Any]]:
    """Return all matching rows as a list of dicts."""
    with get_db() as conn:
        cursor = conn.cursor(dictionary=True)
        cursor.execute(query, values)
        return cursor.fetchall()


# ============================================================
# Database Initialization
# ============================================================

def init_db() -> None:
    """
    Run the full schema SQL from database.sql.
    Called once at startup (or from Flask CLI).
    Safe to call multiple times — uses IF NOT EXISTS guards.
    """
    import os
    sql_path = os.path.join(os.path.dirname(__file__), "database.sql")
    with open(sql_path, "r") as f:
        raw = f.read()

    # Split on semicolons to execute statement-by-statement
    statements = [s.strip() for s in raw.split(";") if s.strip()]

    # Use a direct connection (not the pool) for DDL
    conn = mysql.connector.connect(
        host=DB_HOST,
        port=DB_PORT,
        user=DB_USER,
        password=DB_PASSWORD,
        database=DB_NAME,
        charset="utf8mb4",
    )
    cursor = conn.cursor()
    for stmt in statements:
        try:
            cursor.execute(stmt)
        except Error as e:
            # Log but don't crash on duplicate-table warnings
            print(f"[DB] Warning: {e}")
    conn.commit()
    cursor.close()
    conn.close()
    print("[DB] Schema initialized.")


# ============================================================
# Legacy Class Wrapper (backwards-compatible with old code)
# ============================================================

class Database:
    """
    Thin wrapper kept for backward compatibility with any
    code that uses Database().connect() / .execute() etc.
    New code should use the module-level helper functions above.
    """

    def __init__(self):
        self.connection = None

    def connect(self):
        try:
            self.connection = mysql.connector.connect(
                host=DB_HOST, port=DB_PORT,
                user=DB_USER, password=DB_PASSWORD,
                database=DB_NAME
            )
        except Error as e:
            print(f"[DB] Error: {e}")

    def disconnect(self):
        if self.connection and self.connection.is_connected():
            self.connection.close()

    def execute(self, query, values=None):
        cursor = self.connection.cursor()
        try:
            cursor.execute(query, values)
            self.connection.commit()
            return True
        except Error as e:
            print(e)
            self.connection.rollback()
            return False
        finally:
            cursor.close()

    def execute_lastrowid(self, query, values=None):
        cursor = self.connection.cursor()
        try:
            cursor.execute(query, values)
            self.connection.commit()
            return cursor.lastrowid
        except Error as e:
            print(e)
            self.connection.rollback()
            return None
        finally:
            cursor.close()

    def fetch_one(self, query, values=None):
        cursor = self.connection.cursor(dictionary=True)
        try:
            cursor.execute(query, values)
            return cursor.fetchone()
        except Error as e:
            print(e)
            return None
        finally:
            cursor.close()

    def fetch_all(self, query, values=None):
        cursor = self.connection.cursor(dictionary=True)
        try:
            cursor.execute(query, values)
            return cursor.fetchall()
        except Error as e:
            print(e)
            return []
        finally:
            cursor.close()