Spaces:
Running on Zero
Running on Zero
| 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 | |
| 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() | |