Files
ss-tools/backend/src/core/database.py
2026-01-12 12:33:51 +03:00

79 lines
2.5 KiB
Python

# [DEF:backend.src.core.database:Module]
#
# @SEMANTICS: database, sqlite, sqlalchemy, session, persistence
# @PURPOSE: Configures the SQLite database connection and session management.
# @LAYER: Core
# @RELATION: DEPENDS_ON -> sqlalchemy
# @RELATION: USES -> backend.src.models.mapping
#
# @INVARIANT: A single engine instance is used for the entire application.
# [SECTION: IMPORTS]
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker, Session
from backend.src.models.mapping import Base
# Import models to ensure they're registered with Base
from backend.src.models.task import TaskRecord
from backend.src.models.connection import ConnectionConfig
import os
# [/SECTION]
# [DEF:DATABASE_URL:Constant]
DATABASE_URL = os.getenv("DATABASE_URL", "sqlite:///./mappings.db")
# [/DEF:DATABASE_URL:Constant]
# [DEF:TASKS_DATABASE_URL:Constant]
TASKS_DATABASE_URL = os.getenv("TASKS_DATABASE_URL", "sqlite:///./tasks.db")
# [/DEF:TASKS_DATABASE_URL:Constant]
# [DEF:engine:Variable]
engine = create_engine(DATABASE_URL, connect_args={"check_same_thread": False})
# [/DEF:engine:Variable]
# [DEF:tasks_engine:Variable]
tasks_engine = create_engine(TASKS_DATABASE_URL, connect_args={"check_same_thread": False})
# [/DEF:tasks_engine:Variable]
# [DEF:SessionLocal:Class]
# @PURPOSE: A session factory for the main mappings database.
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
# [/DEF:SessionLocal:Class]
# [DEF:TasksSessionLocal:Class]
# @PURPOSE: A session factory for the tasks execution database.
TasksSessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=tasks_engine)
# [/DEF:TasksSessionLocal:Class]
# [DEF:init_db:Function]
# @PURPOSE: Initializes the database by creating all tables.
def init_db():
Base.metadata.create_all(bind=engine)
Base.metadata.create_all(bind=tasks_engine)
# [/DEF:init_db:Function]
# [DEF:get_db:Function]
# @PURPOSE: Dependency for getting a database session.
# @POST: Session is closed after use.
# @RETURN: Generator[Session, None, None]
def get_db():
db = SessionLocal()
try:
yield db
finally:
db.close()
# [/DEF:get_db:Function]
# [DEF:get_tasks_db:Function]
# @PURPOSE: Dependency for getting a tasks database session.
# @POST: Session is closed after use.
# @RETURN: Generator[Session, None, None]
def get_tasks_db():
db = TasksSessionLocal()
try:
yield db
finally:
db.close()
# [/DEF:get_tasks_db:Function]
# [/DEF:backend.src.core.database:Module]