bt-source/panel/class/flask_session/sqlite_pool.py

"""
SQLite专用连接池库
支持线程安全和gevent+greenlet异步环境
"""
import os
import re
import shutil
import sqlite3
import threading
import weakref
import time
from collections import deque
from contextlib import contextmanager
from typing import Optional, Dict, Any, Callable, Union, Generator

try:
    import gevent
    from gevent import local as gevent_local
    from gevent.lock import RLock as GeventRLock
    from gevent.event import Event as GeventEvent

    HAS_GEVENT = True
except ImportError:
    HAS_GEVENT = False
    gevent = None
    gevent_local = None
    GeventRLock = None
    GeventEvent = None


class SQLiteConnectionError(Exception):
    """SQLite连接池相关异常"""
    pass


class SQLiteTimeoutError(SQLiteConnectionError):
    """连接获取超时异常"""
    pass


def detect_async_environment():
    """检测当前运行环境"""
    if HAS_GEVENT and gevent is not None:
        try:
            # 检查是否在gevent上下文中
            gevent.getcurrent()
            return 'gevent'
        except:
            pass
    return 'thread'


class SQLiteConnection:
    """SQLite连接包装器"""

    def __init__(self, pool: 'SQLitePool', connection: sqlite3.Connection,
                 created_at: float):
        self.pool = pool
        self._connection = connection
        self.created_at = created_at
        self.last_used = time.time()
        self.in_use = False
        self._closed = False

    def __enter__(self):
        return self._connection

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()

    def close(self):
        """归还连接到池中"""
        if not self._closed and self.pool:
            self.pool._return_connection(self)

    def invalidate(self):
        """标记连接为无效"""
        self._closed = True
        if self._connection:
            try:
                self._connection.close()
            except:
                pass

    @property
    def connection(self) -> sqlite3.Connection:
        """获取底层sqlite3连接"""
        if self._closed:
            raise SQLiteConnectionError("Connection is closed")
        self.last_used = time.time()
        return self._connection

    def __getattr__(self, name):
        """代理sqlite3.Connection的方法"""
        return getattr(self.connection, name)


class SQLitePool:
    """SQLite连接池基类"""

    def __init__(self,
                 database: str,
                 pool_size: int = 5,
                 max_overflow: int = 10,
                 timeout: float = 30.0,
                 recycle: int = 3600,
                 **connect_args):
        """
        初始化SQLite连接池

        Args:
            database: SQLite数据库文件路径
            pool_size: 池大小
            max_overflow: 最大溢出连接数
            timeout: 获取连接超时时间(秒)
            recycle: 连接回收时间(秒)
            **connect_args: sqlite3.connect的参数
        """
        if database.startswith("sqlite"):
            database = re.sub(r"sqlite[^/]*//", "", database)
        self.database = database
        self.pool_size = pool_size
        self.max_overflow = max_overflow
        self.timeout = timeout
        self.recycle = recycle
        self.connect_args = connect_args

        # 设置SQLite文件数据库的默认参数
        if database != ':memory:':
            self.connect_args.setdefault('check_same_thread', False)
            self.connect_args.setdefault('timeout', 20.0)
        else:
            # 内存数据库需要共享连接
            self.connect_args.setdefault('check_same_thread', False)

        self._pool = deque()
        self._overflow_count = 0
        self._created_connections = 0

        # 环境适配
        self._env = detect_async_environment()
        self._setup_environment()

    def _setup_environment(self):
        """根据环境设置锁和本地存储"""
        if self._env == 'gevent' and GeventRLock is not None and gevent_local is not None:
            self._lock = GeventRLock()
            self._local = gevent_local.local()
        else:
            self._lock = threading.RLock()
            self._local = threading.local()

    def _create_connection(self) -> sqlite3.Connection:
        """创建新的SQLite连接"""
        try:
            conn = sqlite3.connect(self.database, **self.connect_args)

            # 设置SQLite优化参数
            if self.database != ':memory:':
                conn.execute("PRAGMA journal_mode=WAL")
                conn.execute("PRAGMA synchronous=NORMAL")
                conn.execute("PRAGMA cache_size=-64000")  # 64MB缓存
                conn.execute("PRAGMA temp_store=MEMORY")

            return conn
        except Exception as e:
            if self.database == ':memory:':
                raise SQLiteConnectionError(f"Failed to create connection: {e}")
            print("链接数据库失败，尝试重建数据库")
            os.makedirs(os.path.dirname(self.database), exist_ok=True)
            if os.path.exists(self.database):
                if os.path.exists(self.database + '.bak'):
                    os.remove(self.database + '.bak')
                shutil.move(self.database, self.database + '.bak')
            open(self.database, 'w').close()
            try:
                conn = self._create_connection()
                # 设置SQLite优化参数
                if self.database != ':memory:':
                    conn.execute("PRAGMA journal_mode=WAL")
                    conn.execute("PRAGMA synchronous=NORMAL")
                    conn.execute("PRAGMA cache_size=-64000")  # 64MB缓存
                    conn.execute("PRAGMA temp_store=MEMORY")

                return conn
            except:
                raise SQLiteConnectionError(f"Failed to create connection: {e}")

    def _is_connection_valid(self, conn: sqlite3.Connection,
                             created_at: float) -> bool:
        """检查连接是否有效"""
        # 检查回收时间
        if self.recycle > 0 and time.time() - created_at > self.recycle:
            return False

        # 简单的ping测试
        try:
            conn.execute("SELECT 1").fetchone()
            return True
        except:
            return False

    def _get_connection_from_pool(self) -> Optional[SQLiteConnection]:
        """从池中获取连接"""
        with self._lock:
            while self._pool:
                wrapped_conn = self._pool.popleft()
                if self._is_connection_valid(wrapped_conn._connection,
                                             wrapped_conn.created_at):
                    wrapped_conn.in_use = True
                    return wrapped_conn
                else:
                    # 连接已过期，关闭它
                    wrapped_conn.invalidate()
                    self._created_connections -= 1
        return None

    def _create_new_connection(self) -> SQLiteConnection:
        """创建新连接"""
        conn = self._create_connection()
        wrapped_conn = SQLiteConnection(self, conn, time.time())
        wrapped_conn.in_use = True

        with self._lock:
            self._created_connections += 1
            if self._created_connections > self.pool_size:
                self._overflow_count += 1

        return wrapped_conn

    def _cleanup_expired(self):
        """清理过期的连接"""
        current_time = time.time()
        expired_connections = []

        with self._lock:
            # 创建一个新的队列来存储有效连接
            valid_connections = deque()

            # 检查池中的每个连接
            while self._pool:
                wrapped_conn = self._pool.popleft()

                # 如果连接未在使用且已过期
                if (not wrapped_conn.in_use and
                        (self.recycle > 0 and
                         current_time - wrapped_conn.created_at > self.recycle)):
                    expired_connections.append(wrapped_conn)
                elif (not wrapped_conn.in_use and
                      not self._is_connection_valid(wrapped_conn._connection,
                                                    wrapped_conn.created_at)):
                    # 连接无效
                    expired_connections.append(wrapped_conn)
                else:
                    # 连接仍然有效
                    valid_connections.append(wrapped_conn)

            # 更新池
            self._pool = valid_connections

            # 更新计数器
            for _ in expired_connections:
                self._created_connections -= 1
                if self._overflow_count > 0:
                    self._overflow_count -= 1

        # 在锁外关闭过期连接
        for wrapped_conn in expired_connections:
            wrapped_conn.invalidate()

        return len(expired_connections)

    def get_connection(self) -> SQLiteConnection:
        """获取连接"""
        raise NotImplementedError

    def _return_connection(self, wrapped_conn: SQLiteConnection):
        """归还连接到池中"""
        if wrapped_conn._closed:
            return

        wrapped_conn.in_use = False

        with self._lock:
            # 检查是否应该保留在池中
            if (len(self._pool) < self.pool_size and
                    self._is_connection_valid(wrapped_conn._connection,
                                              wrapped_conn.created_at)):
                self._pool.append(wrapped_conn)
            else:
                # 关闭溢出连接或过期连接
                wrapped_conn.invalidate()
                self._created_connections -= 1
                if self._overflow_count > 0:
                    self._overflow_count -= 1

    def dispose(self):
        """清理所有连接"""
        with self._lock:
            while self._pool:
                conn = self._pool.popleft()
                conn.invalidate()

            self._created_connections = 0
            self._overflow_count = 0

    def status(self) -> Dict[str, Any]:
        """获取池状态"""
        with self._lock:
            return {
                'pool_size': self.pool_size,
                'connections_in_pool': len(self._pool),
                'total_created': self._created_connections,
                'overflow_count': self._overflow_count,
                'environment': self._env
            }


class ThreadSafePool(SQLitePool):
    """线程安全的SQLite连接池"""

    def get_connection(self) -> SQLiteConnection:
        """获取连接（线程安全版本）"""
        start_time = time.time()

        while True:
            # 尝试从池中获取连接
            conn = self._get_connection_from_pool()
            if conn:
                return conn

            # 检查是否可以创建新连接
            with self._lock:
                can_create = (self._created_connections <
                              self.pool_size + self.max_overflow)

            if can_create:
                return self._create_new_connection()

            # 检查超时
            if time.time() - start_time > self.timeout:
                raise SQLiteTimeoutError(
                    f"Could not get connection within {self.timeout} seconds"
                )

            # 等待一段时间后重试
            if self._env == 'gevent' and gevent is not None:
                gevent.sleep(0.01)
            else:
                time.sleep(0.01)


class StaticPool(SQLitePool):
    """静态连接池 - 适用于内存数据库"""

    def __init__(self, database: str, **connect_args):
        super().__init__(database, pool_size=1, max_overflow=0, **connect_args)
        self._static_connection: Optional[SQLiteConnection] = None
        if self._env == 'gevent' and GeventRLock is not None:
            self._static_lock = GeventRLock()
        else:
            self._static_lock = threading.RLock()

    def get_connection(self) -> SQLiteConnection:
        """获取静态连接"""
        with self._static_lock:
            if (self._static_connection is None or
                    not self._is_connection_valid(self._static_connection._connection, self._static_connection.created_at)):
                if self._static_connection:
                    self._static_connection.invalidate()

                conn = self._create_connection()
                self._static_connection = SQLiteConnection(self, conn, time.time())

            return self._static_connection

    def _return_connection(self, wrapped_conn: SQLiteConnection):
        """静态池不需要归还连接"""
        wrapped_conn.in_use = False

    def _cleanup_expired(self):
        """静态池的清理逻辑"""
        with self._static_lock:
            if (self._static_connection and
                    not self._static_connection.in_use and
                    (0 < self.recycle < time.time() - self._static_connection.created_at)):
                self._static_connection.invalidate()
                self._static_connection = None
                return 1
        return 0


def create_pool(database: str,
                pool_type: str = 'auto',
                **kwargs) -> SQLitePool:
    """
    创建SQLite连接池

    Args:
        database: 数据库文件路径
        pool_type: 池类型 ('auto', 'thread', 'static')
        **kwargs: 其他参数

    Returns:
        SQLitePool实例
    """
    if pool_type == 'auto':
        if database == ':memory:':
            pool_type = 'static'
        else:
            pool_type = 'thread'

    if pool_type == 'static':
        return StaticPool(database, **kwargs)
    elif pool_type == 'thread':
        return ThreadSafePool(database, **kwargs)
    else:
        raise ValueError(f"Unknown pool type: {pool_type}")


@contextmanager
def get_connection(pool: SQLitePool) -> Generator[sqlite3.Connection, None, None]:
    """
    上下文管理器方式获取连接

    Usage:
        pool = create_pool('test.db')
        with get_connection(pool) as conn:
            cursor = conn.cursor()
            cursor.execute("SELECT 1")
    """
    wrapped_conn = pool.get_connection()
    try:
        yield wrapped_conn.connection
    finally:
        wrapped_conn.close()


# Flask扩展集成示例
class FlaskSQLitePool:
    """Flask扩展集成"""

    def __init__(self, app=None):
        self.pool = None
        if app:
            self.init_app(app)

    def init_app(self, app):
        """初始化Flask应用"""
        database = app.config.get('SQLITE_DATABASE', ':memory:')
        pool_config = app.config.get('SQLITE_POOL_CONFIG', {})

        self.pool = create_pool(database, **pool_config)

        # 注册清理函数
        @app.teardown_appcontext
        def close_db(error):
            # 在请求结束时清理过期连接
            if self.pool and hasattr(self.pool, '_cleanup_expired'):
                cleaned = self.pool._cleanup_expired()
                if cleaned > 0:
                    app.logger.debug(f"Cleaned up {cleaned} expired connections")

        app.extensions['sqlite_pool'] = self

    def get_connection(self) -> SQLiteConnection:
        """获取连接"""
        if not self.pool:
            raise RuntimeError("Pool not initialized")
        return self.pool.get_connection()


# 使用示例
if __name__ == "__main__":
    # 文件数据库示例
    print("=== 文件数据库示例 ===")
    file_pool = create_pool('test.db', pool_size=3, max_overflow=2)

    with get_connection(file_pool) as conn:
        conn.execute("""
                     CREATE TABLE IF NOT EXISTS users
                     (
                         id   INTEGER PRIMARY KEY,
                         name TEXT NOT NULL
                     )
                     """)
        conn.execute("INSERT INTO users (name) VALUES (?)", ("Alice",))
        conn.commit()

    print("文件数据库池状态:", file_pool.status())

    # 内存数据库示例
    print("\n=== 内存数据库示例 ===")
    memory_pool = create_pool(':memory:', pool_type='static')

    with get_connection(memory_pool) as conn:
        conn.execute("CREATE TABLE test (id INTEGER, value TEXT)")
        conn.execute("INSERT INTO test VALUES (1, 'test')")
        result = conn.execute("SELECT * FROM test").fetchall()
        print("查询结果:", result)

    print("内存数据库池状态:", memory_pool.status())

    # 多线程测试
    print("\n=== 多线程测试 ===")
    import concurrent.futures


    def worker(worker_id):
        with get_connection(file_pool) as conn:
            conn.execute("INSERT INTO users (name) VALUES (?)",
                         (f"Worker-{worker_id}",))
            conn.commit()
            result = conn.execute("SELECT COUNT(*) FROM users").fetchone()
            return f"Worker-{worker_id}: {result[0]} users"


    with concurrent.futures.ThreadPoolExecutor(max_workers=5) as executor:
        futures = [executor.submit(worker, i) for i in range(5)]
        for future in concurrent.futures.as_completed(futures):
            print(future.result())

    print("最终池状态:", file_pool.status())

    # 清理
    file_pool.dispose()
    memory_pool.dispose()