StarPunk/starpunk/database/pool.py

"""
Database connection pool for StarPunk

Per ADR-053 and developer Q&A Q2, CQ1:
- Provides connection pooling for improved performance
- Integrates with Flask's g object for request-scoped connections
- Maintains same interface as get_db() for transparency
- Pool statistics available for metrics
- Wraps connections with MonitoredConnection for timing (v1.1.2 Phase 1)

Note: Migrations use direct connections (not pooled) for isolation
"""

import sqlite3
from pathlib import Path
from threading import Lock
from collections import deque
from flask import g
from typing import Optional


class ConnectionPool:
    """
    Simple connection pool for SQLite

    SQLite doesn't benefit from traditional connection pooling like PostgreSQL,
    but this provides connection reuse and request-scoped connection management.
    """

    def __init__(self, db_path, pool_size=5, timeout=10.0, slow_query_threshold=1.0, metrics_enabled=True):
        """
        Initialize connection pool

        Args:
            db_path: Path to SQLite database file
            pool_size: Maximum number of connections in pool
            timeout: Timeout for getting connection (seconds)
            slow_query_threshold: Threshold in seconds for slow query detection (v1.1.2)
            metrics_enabled: Whether to enable metrics collection (v1.1.2)
        """
        self.db_path = Path(db_path)
        self.pool_size = pool_size
        self.timeout = timeout
        self.slow_query_threshold = slow_query_threshold
        self.metrics_enabled = metrics_enabled
        self._pool = deque(maxlen=pool_size)
        self._lock = Lock()
        self._stats = {
            'connections_created': 0,
            'connections_reused': 0,
            'connections_closed': 0,
            'pool_hits': 0,
            'pool_misses': 0,
        }

    def _create_connection(self):
        """
        Create a new database connection

        Per CQ1: Wraps connection with MonitoredConnection if metrics enabled
        """
        conn = sqlite3.connect(
            self.db_path,
            timeout=self.timeout,
            check_same_thread=False  # Allow connection reuse across threads
        )
        conn.row_factory = sqlite3.Row  # Return rows as dictionaries

        # Enable WAL mode for better concurrency
        conn.execute("PRAGMA journal_mode=WAL")

        self._stats['connections_created'] += 1

        # Wrap with monitoring if enabled (v1.1.2 Phase 1)
        if self.metrics_enabled:
            from starpunk.monitoring import MonitoredConnection
            return MonitoredConnection(conn, self.slow_query_threshold)

        return conn

    def get_connection(self):
        """
        Get a connection from the pool

        Returns:
            sqlite3.Connection: Database connection
        """
        with self._lock:
            if self._pool:
                # Reuse existing connection
                conn = self._pool.pop()
                self._stats['pool_hits'] += 1
                self._stats['connections_reused'] += 1
                return conn
            else:
                # Create new connection
                self._stats['pool_misses'] += 1
                return self._create_connection()

    def return_connection(self, conn):
        """
        Return a connection to the pool

        Args:
            conn: Database connection to return
        """
        if not conn:
            return

        with self._lock:
            if len(self._pool) < self.pool_size:
                # Return to pool
                self._pool.append(conn)
            else:
                # Pool is full, close connection
                conn.close()
                self._stats['connections_closed'] += 1

    def close_connection(self, conn):
        """
        Close a connection without returning to pool

        Args:
            conn: Database connection to close
        """
        if conn:
            conn.close()
            self._stats['connections_closed'] += 1

    def get_stats(self):
        """
        Get pool statistics

        Returns:
            dict: Pool statistics for monitoring
        """
        with self._lock:
            return {
                **self._stats,
                'pool_size': len(self._pool),
                'max_pool_size': self.pool_size,
            }

    def close_all(self):
        """Close all connections in the pool"""
        with self._lock:
            while self._pool:
                conn = self._pool.pop()
                conn.close()
                self._stats['connections_closed'] += 1


# Global pool instance (initialized by app factory)
_pool = None


def init_pool(app):
    """
    Initialize the connection pool

    Per CQ2: Passes metrics configuration from app config

    Args:
        app: Flask application instance
    """
    global _pool

    db_path = app.config['DATABASE_PATH']
    pool_size = app.config.get('DB_POOL_SIZE', 5)
    timeout = app.config.get('DB_TIMEOUT', 10.0)
    slow_query_threshold = app.config.get('METRICS_SLOW_QUERY_THRESHOLD', 1.0)
    metrics_enabled = app.config.get('METRICS_ENABLED', True)

    _pool = ConnectionPool(
        db_path,
        pool_size,
        timeout,
        slow_query_threshold,
        metrics_enabled
    )
    app.logger.info(
        f"Database connection pool initialized "
        f"(size={pool_size}, metrics={'enabled' if metrics_enabled else 'disabled'})"
    )

    # Register teardown handler
    @app.teardown_appcontext
    def close_connection(error):
        """Return connection to pool when request context ends"""
        conn = g.pop('db', None)
        if conn:
            _pool.return_connection(conn)


def get_db(app=None):
    """
    Get database connection for current request

    Uses Flask's g object for request-scoped connection management.
    Connection is automatically returned to pool at end of request.

    Args:
        app: Flask application (optional, for backward compatibility with tests)
             When provided, this parameter is ignored as we use the pool

    Returns:
        sqlite3.Connection: Database connection
    """
    # Note: app parameter is kept for backward compatibility but ignored
    # The pool is request-scoped via Flask's g object
    if 'db' not in g:
        g.db = _pool.get_connection()
    return g.db


def get_pool_stats():
    """
    Get connection pool statistics

    Returns:
        dict: Pool statistics for monitoring
    """
    if _pool:
        return _pool.get_stats()
    return {}