Spaces:
Running on Zero
Running on Zero
File size: 5,907 Bytes
8442ef1 e8553f3 8442ef1 8028058 c60d5a5 8442ef1 c60d5a5 8442ef1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 | import os
from contextlib import contextmanager
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker, declarative_base, Session
from utils.logger import logger
# Base class for model definitions
Base = declarative_base()
# Configure SQLite DB path (default: data/discord.db)
DB_PATH = os.getenv("DATABASE_URL", "sqlite:///data/discord.db")
# Ensure the parent directory for the database exists
if DB_PATH.startswith("sqlite:///"):
sqlite_file = DB_PATH.replace("sqlite:///", "")
parent_dir = os.path.dirname(sqlite_file)
if parent_dir:
try:
os.makedirs(parent_dir, exist_ok=True)
except PermissionError as e:
logger.warning(
f"Could not create database directory '{parent_dir}' due to PermissionError: {e}. "
"Falling back to local writable directory 'data/'."
)
local_dir = "data"
os.makedirs(local_dir, exist_ok=True)
DB_PATH = f"sqlite:///{local_dir}/{os.path.basename(sqlite_file)}"
logger.info(f"Connecting to database at {DB_PATH}")
# Create Engine
# Note: For SQLite, check_same_thread=False is needed since discord.py event loop
# runs async tasks on threads, and SQLite connections are bound to threads by default.
connect_args = {}
if DB_PATH.startswith("sqlite:"):
connect_args["check_same_thread"] = False
engine = create_engine(DB_PATH, connect_args=connect_args)
# Create SessionMaker
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
def init_db():
"""Create all database tables mapped by Base."""
try:
logger.info("Initializing database tables...")
Base.metadata.create_all(bind=engine)
logger.info("Database tables initialized successfully.")
# Check and clear message embeddings if dimension mismatches configuration
from sqlalchemy import text
from sqlalchemy.orm import Session
from config import settings
import json
with Session(engine) as session:
try:
result = session.execute(text("SELECT embedding FROM message_embeddings LIMIT 1")).fetchone()
if result:
emb_list = json.loads(result[0])
if len(emb_list) != settings.EMBEDDING_DIMENSION:
logger.warning(
f"Detected embedding dimension mismatch ({len(emb_list)} != {settings.EMBEDDING_DIMENSION}). "
"Clearing existing embeddings in the database to allow regeneration."
)
session.execute(text("DELETE FROM message_embeddings"))
session.commit()
except Exception as db_err:
logger.warning(f"Failed to check/clear embedding dimension mismatch: {db_err}")
try:
# Set up FTS5 virtual table and triggers for SQLite full-text search
if DB_PATH.startswith("sqlite"):
logger.info("Initializing SQLite FTS5 virtual table and synchronization triggers...")
# Create FTS5 virtual table
session.execute(text("""
CREATE VIRTUAL TABLE IF NOT EXISTS messages_fts USING fts5(
message_id UNINDEXED,
content
)
"""))
# Create synchronization triggers
session.execute(text("""
CREATE TRIGGER IF NOT EXISTS messages_ai AFTER INSERT ON messages BEGIN
INSERT INTO messages_fts(message_id, content) VALUES (new.message_id, new.content);
END;
"""))
session.execute(text("""
CREATE TRIGGER IF NOT EXISTS messages_ad AFTER DELETE ON messages BEGIN
DELETE FROM messages_fts WHERE message_id = old.message_id;
END;
"""))
session.execute(text("""
CREATE TRIGGER IF NOT EXISTS messages_au AFTER UPDATE OF content ON messages BEGIN
UPDATE messages_fts SET content = new.content WHERE message_id = old.message_id;
END;
"""))
# One-time migration: Populate FTS5 table if it is currently empty
fts_count = session.execute(text("SELECT count(*) FROM messages_fts")).fetchone()[0]
if fts_count == 0:
logger.info("Migrating existing messages into FTS5 index...")
session.execute(text("""
INSERT INTO messages_fts(message_id, content)
SELECT message_id, content
FROM messages
WHERE is_deleted = 0 AND content IS NOT NULL AND content != ''
"""))
session.commit()
logger.info("SQLite FTS5 virtual table and triggers initialized successfully.")
except Exception as fts_err:
logger.error(f"Failed to configure SQLite FTS5 index: {fts_err}", exc_info=True)
except Exception as e:
logger.error(f"Error initializing database tables: {e}", exc_info=True)
raise
@contextmanager
def get_db_session():
"""Provide a transactional scope around a series of operations."""
session: Session = SessionLocal()
try:
yield session
session.commit()
except Exception as e:
session.rollback()
logger.error(f"Database transaction error: {e}", exc_info=True)
raise
finally:
session.close()
|