130 lines
3.6 KiB
Python
130 lines
3.6 KiB
Python
"""
|
||
Database utilities for RAG system
|
||
"""
|
||
import sqlite3
|
||
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 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()
|