RAG/db_utils.py

"""
Database utilities for RAG system
"""
import sqlite3
import json
from pathlib import Path
from datetime import datetime
from typing import Tuple, Optional
from loguru import logger


def get_db_connection() -> Tuple[sqlite3.Connection, sqlite3.Cursor]:
    """
    Get a SQLite database connection with data_sources table initialized
    
    Returns:
        tuple: (connection, cursor)
    """
    DATA_DIR = Path(__file__).parent / "data"
    DATA_DIR.mkdir(parents=True, exist_ok=True)
    DB_PATH = DATA_DIR / "sessions.db"
    
    conn = sqlite3.connect(DB_PATH)
    cursor = conn.cursor()
    
    # 检查并创建data_sources表（如果不存在）
    try:
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS data_sources (
                name TEXT PRIMARY KEY,
                config TEXT NOT NULL,
                update_at TEXT NULL
            )
        ''')
        conn.commit()
    except Exception as e:
        logger.error(f"Error creating data_sources table: {e}")
        conn.close()
        raise
    
    return conn, cursor


def get_data_source_update_at(source_name: str) -> Optional[datetime]:
    """
    Get update_at for a data source from data_sources table
    
    Args:
        source_name: Name of the data source
        
    Returns:
        datetime: Update time if found, None otherwise
    """
    try:
        conn, cursor = get_db_connection()
        try:
            cursor.execute('SELECT update_at FROM data_sources WHERE name = ?', (source_name,))
            result = cursor.fetchone()
            if result and result[0]:
                return datetime.fromisoformat(result[0])
            return None
        finally:
            conn.close()
    except Exception as e:
        logger.warning(f"Error reading update_at from data_sources: {e}")
        return None


def update_data_source_update_at(source_name: str, update_at: datetime) -> bool:
    """
    Update update_at for a data source in data_sources table
    
    Args:
        source_name: Name of the data source
        update_at: New update time
        
    Returns:
        bool: True if update succeeded, False otherwise
    """
    try:
        conn, cursor = get_db_connection()
        try:
            cursor.execute('UPDATE data_sources SET update_at = ? WHERE name = ?', (update_at.isoformat(), source_name))
            conn.commit()
            logger.info(f"Updated update_at in data_sources for {source_name}: {update_at}")
            return True
        finally:
            conn.close()
    except Exception as e:
        logger.warning(f"Error updating update_at in data_sources: {e}")
        return False


def add_git_datasource(user_id: str, repo_config: dict):
    """
    新增Git仓库配置
    
    Args:
        user_id: 用户ID
        repo_config: 仓库配置
    """
    conn, cursor = get_db_connection()
    try:
        cursor.execute("""
            INSERT INTO datasource (user_id, name, type, git_config, create_time)
            VALUES (?, ?, 'git', ?, datetime('now'))
        """, (user_id, repo_config["name"], json.dumps(repo_config)))
        conn.commit()
        logger.info(f"新增Git数据源: {repo_config['name']}")
    finally:
        conn.close()


def update_git_sync_status(user_id: str, repo_id: str, sync_status: dict):
    """
    更新Git仓库同步状态
    
    Args:
        user_id: 用户ID
        repo_id: 仓库ID
        sync_status: 同步状态
    """
    conn, cursor = get_db_connection()
    try:
        cursor.execute("""
            UPDATE datasource SET git_config = json_set(git_config, '$.latest_commit_id', ?, '$.last_sync_time', ?)
            WHERE user_id = ? AND id = ?
        """, (sync_status["latest_commit_id"], sync_status["last_sync_time"], user_id, repo_id))
        conn.commit()
        logger.info(f"更新Git同步状态: {repo_id}")
    finally:
        conn.close()


def update_git_repo_config(user_id: str, repo_id: str, config: dict):
    """
    更新Git仓库配置
    
    Args:
        user_id: 用户ID
        repo_id: 仓库ID
        config: 配置信息
    """
    conn, cursor = get_db_connection()
    try:
        # 获取当前配置
        cursor.execute("SELECT git_config FROM datasource WHERE user_id = ? AND id = ?", (user_id, repo_id))
        result = cursor.fetchone()
        if result:
            current_config = json.loads(result[0])
            # 更新配置
            current_config.update(config)
            cursor.execute("""
                UPDATE datasource SET git_config = ?
                WHERE user_id = ? AND id = ?
            """, (json.dumps(current_config), user_id, repo_id))
            conn.commit()
            logger.info(f"更新Git仓库配置: {repo_id}")
    finally:
        conn.close()


def init_session_db():
    """
    Initialize session database with users and sessions tables
    """
    DATA_DIR = Path(__file__).parent / "data"
    DATA_DIR.mkdir(parents=True, exist_ok=True)
    DB_PATH = DATA_DIR / "sessions.db"
    
    conn = sqlite3.connect(DB_PATH)
    try:
        conn.execute(
            """
            CREATE TABLE IF NOT EXISTS users (
                id TEXT PRIMARY KEY,
                username TEXT UNIQUE,
                password TEXT,
                create_time TEXT
            )
            """
        )
        conn.execute(
            """
            CREATE TABLE IF NOT EXISTS sessions (
                id TEXT PRIMARY KEY,
                user_login TEXT,
                title TEXT,
                data TEXT,
                update_time TEXT
            )
            """
        )
        # Improve concurrency for writes
        conn.execute("PRAGMA journal_mode=WAL;")
        conn.execute("PRAGMA synchronous=NORMAL;")
        conn.commit()
    finally:
        conn.close()