busya/ss-tools
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
Files
d1695fe536060ed5bb0f2009fcaf7c86a14d59b8
ss-tools/backend/src/core/database.py
busya a3c7c402b7 fix(llm): add fetch-models endpoint, fix SQL Lab INSERT (client_id truncation, sync mode, target_column, timestamp normalization) 
- Add POST /api/llm/providers/fetch-models route with LLMClient.fetch_models()
- Add target_column to TranslationJob model/schema/service/orchestrator
- Fix SQL Lab execute: truncate client_id to 11 chars (varchar(11))
- Switch SQL Lab to sync mode (runAsync: false) — no Celery workers
- Fix polling: unwrap nested result from Superset query API
- Fix ClickHouse timestamp: normalize float timestamps to YYYY-MM-DD
2026-05-13 20:06:15 +03:00

661 lines
25 KiB
Python
Raw Blame History

# #region DatabaseModule [C:3] [TYPE Module] [SEMANTICS sqlalchemy, connection, session]
#
# @BRIEF Configures database connection and session management (PostgreSQL-first).
# @LAYER: Core
# @RELATION DEPENDS_ON -> [MappingModels]
# @RELATION DEPENDS_ON -> [auth_config]
# @RELATION DEPENDS_ON -> [ConnectionConfig]
#
# @INVARIANT: A single engine instance is used for the entire application.
from sqlalchemy import create_engine, inspect, text
from sqlalchemy.orm import sessionmaker
from ..models.mapping import Base
from ..models.connection import ConnectionConfig
# Import models to ensure they're registered with Base
from ..models import task as _task_models # noqa: F401
from ..models import auth as _auth_models # noqa: F401
from ..models import config as _config_models # noqa: F401
from ..models import llm as _llm_models # noqa: F401
from ..models import assistant as _assistant_models # noqa: F401
from ..models import profile as _profile_models # noqa: F401
from ..models import clean_release as _clean_release_models # noqa: F401
from ..models import connection as _connection_models # noqa: F401
from ..models import dataset_review as _dataset_review_models # noqa: F401
from .logger import belief_scope, logger
from .auth.config import auth_config
import os
from pathlib import Path
# #region BASE_DIR [C:1] [TYPE Variable]
# @BRIEF Base directory for the backend.
BASE_DIR = Path(__file__).resolve().parent.parent.parent
# #endregion BASE_DIR
# #region DATABASE_URL [C:1] [TYPE Constant]
# @BRIEF URL for the main application database.
DEFAULT_POSTGRES_URL = os.getenv(
"POSTGRES_URL",
"postgresql+psycopg2://postgres:postgres@localhost:5432/ss_tools",
)
DATABASE_URL = os.getenv("DATABASE_URL", DEFAULT_POSTGRES_URL)
# #endregion DATABASE_URL
# #region TASKS_DATABASE_URL [C:1] [TYPE Constant]
# @BRIEF URL for the tasks execution database.
# Defaults to DATABASE_URL to keep task logs in the same PostgreSQL instance.
TASKS_DATABASE_URL = os.getenv("TASKS_DATABASE_URL", DATABASE_URL)
# #endregion TASKS_DATABASE_URL
# #region AUTH_DATABASE_URL [C:1] [TYPE Constant]
# @BRIEF URL for the authentication database.
AUTH_DATABASE_URL = os.getenv("AUTH_DATABASE_URL", auth_config.AUTH_DATABASE_URL)
# #endregion AUTH_DATABASE_URL
# #region engine [C:1] [TYPE Variable]
# @BRIEF SQLAlchemy engine for mappings database.
# @SIDE_EFFECT: Creates database engine and manages connection pool.
def _build_engine(db_url: str):
with belief_scope("_build_engine"):
if db_url.startswith("sqlite"):
return create_engine(db_url, connect_args={"check_same_thread": False})
return create_engine(db_url, pool_pre_ping=True)
engine = _build_engine(DATABASE_URL)
# #endregion engine
# #region tasks_engine [C:1] [TYPE Variable]
# @BRIEF SQLAlchemy engine for tasks database.
tasks_engine = _build_engine(TASKS_DATABASE_URL)
# #endregion tasks_engine
# #region auth_engine [C:1] [TYPE Variable]
# @BRIEF SQLAlchemy engine for authentication database.
auth_engine = _build_engine(AUTH_DATABASE_URL)
# #endregion auth_engine
# #region SessionLocal [C:1] [TYPE Class]
# @BRIEF A session factory for the main mappings database.
# @PRE: engine is initialized.
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
# #endregion SessionLocal
# #region TasksSessionLocal [C:1] [TYPE Class]
# @BRIEF A session factory for the tasks execution database.
# @PRE: tasks_engine is initialized.
TasksSessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=tasks_engine)
# #endregion TasksSessionLocal
# #region AuthSessionLocal [C:1] [TYPE Class]
# @BRIEF A session factory for the authentication database.
# @PRE: auth_engine is initialized.
AuthSessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=auth_engine)
# #endregion AuthSessionLocal
# #region _ensure_user_dashboard_preferences_columns [C:3] [TYPE Function]
# @BRIEF Applies additive schema upgrades for user_dashboard_preferences table.
# @PRE: bind_engine points to application database where profile table is stored.
# @POST: Missing columns are added without data loss.
# @RELATION DEPENDS_ON -> [engine]
def _ensure_user_dashboard_preferences_columns(bind_engine):
with belief_scope("_ensure_user_dashboard_preferences_columns"):
table_name = "user_dashboard_preferences"
inspector = inspect(bind_engine)
if table_name not in inspector.get_table_names():
return
existing_columns = {
str(column.get("name") or "").strip()
for column in inspector.get_columns(table_name)
}
alter_statements = []
if "git_username" not in existing_columns:
alter_statements.append(
"ALTER TABLE user_dashboard_preferences ADD COLUMN git_username VARCHAR"
)
if "git_email" not in existing_columns:
alter_statements.append(
"ALTER TABLE user_dashboard_preferences ADD COLUMN git_email VARCHAR"
)
if "git_personal_access_token_encrypted" not in existing_columns:
alter_statements.append(
"ALTER TABLE user_dashboard_preferences "
"ADD COLUMN git_personal_access_token_encrypted VARCHAR"
)
if "start_page" not in existing_columns:
alter_statements.append(
"ALTER TABLE user_dashboard_preferences "
"ADD COLUMN start_page VARCHAR NOT NULL DEFAULT 'dashboards'"
)
if "auto_open_task_drawer" not in existing_columns:
alter_statements.append(
"ALTER TABLE user_dashboard_preferences "
"ADD COLUMN auto_open_task_drawer BOOLEAN NOT NULL DEFAULT TRUE"
)
if "dashboards_table_density" not in existing_columns:
alter_statements.append(
"ALTER TABLE user_dashboard_preferences "
"ADD COLUMN dashboards_table_density VARCHAR NOT NULL DEFAULT 'comfortable'"
)
if "show_only_slug_dashboards" not in existing_columns:
alter_statements.append(
"ALTER TABLE user_dashboard_preferences "
"ADD COLUMN show_only_slug_dashboards BOOLEAN NOT NULL DEFAULT TRUE"
)
if not alter_statements:
return
try:
with bind_engine.begin() as connection:
for statement in alter_statements:
connection.execute(text(statement))
except Exception as migration_error:
logger.explore(
"Profile preference additive migration failed",
extra={"src": "database", "error": str(migration_error)},
)
# #endregion _ensure_user_dashboard_preferences_columns
# #region _ensure_user_dashboard_preferences_health_columns [C:3] [TYPE Function]
# @BRIEF Applies additive schema upgrades for user_dashboard_preferences table (health fields).
# @RELATION DEPENDS_ON -> [engine]
def _ensure_user_dashboard_preferences_health_columns(bind_engine):
with belief_scope("_ensure_user_dashboard_preferences_health_columns"):
table_name = "user_dashboard_preferences"
inspector = inspect(bind_engine)
if table_name not in inspector.get_table_names():
return
existing_columns = {
str(column.get("name") or "").strip()
for column in inspector.get_columns(table_name)
}
alter_statements = []
if "telegram_id" not in existing_columns:
alter_statements.append(
"ALTER TABLE user_dashboard_preferences ADD COLUMN telegram_id VARCHAR"
)
if "email_address" not in existing_columns:
alter_statements.append(
"ALTER TABLE user_dashboard_preferences ADD COLUMN email_address VARCHAR"
)
if "notify_on_fail" not in existing_columns:
alter_statements.append(
"ALTER TABLE user_dashboard_preferences ADD COLUMN notify_on_fail BOOLEAN NOT NULL DEFAULT TRUE"
)
if not alter_statements:
return
try:
with bind_engine.begin() as connection:
for statement in alter_statements:
connection.execute(text(statement))
except Exception as migration_error:
logger.explore(
"Profile health preference additive migration failed",
extra={"src": "database", "error": str(migration_error)},
)
# #endregion _ensure_user_dashboard_preferences_health_columns
# #region _ensure_llm_validation_results_columns [C:3] [TYPE Function]
# @BRIEF Applies additive schema upgrades for llm_validation_results table.
# @RELATION DEPENDS_ON -> [engine]
def _ensure_llm_validation_results_columns(bind_engine):
with belief_scope("_ensure_llm_validation_results_columns"):
table_name = "llm_validation_results"
inspector = inspect(bind_engine)
if table_name not in inspector.get_table_names():
return
existing_columns = {
str(column.get("name") or "").strip()
for column in inspector.get_columns(table_name)
}
alter_statements = []
if "task_id" not in existing_columns:
alter_statements.append(
"ALTER TABLE llm_validation_results ADD COLUMN task_id VARCHAR"
)
if "environment_id" not in existing_columns:
alter_statements.append(
"ALTER TABLE llm_validation_results ADD COLUMN environment_id VARCHAR"
)
if not alter_statements:
return
try:
with bind_engine.begin() as connection:
for statement in alter_statements:
connection.execute(text(statement))
except Exception as migration_error:
logger.explore(
"ValidationRecord additive migration failed",
extra={"src": "database", "error": str(migration_error)},
)
# #endregion _ensure_llm_validation_results_columns
# #region _ensure_git_server_configs_columns [C:3] [TYPE Function]
# @BRIEF Applies additive schema upgrades for git_server_configs table.
# @PRE: bind_engine points to application database.
# @POST: Missing columns are added without data loss.
# @RELATION DEPENDS_ON -> [engine]
def _ensure_git_server_configs_columns(bind_engine):
with belief_scope("_ensure_git_server_configs_columns"):
table_name = "git_server_configs"
inspector = inspect(bind_engine)
if table_name not in inspector.get_table_names():
return
existing_columns = {
str(column.get("name") or "").strip()
for column in inspector.get_columns(table_name)
}
alter_statements = []
if "default_branch" not in existing_columns:
alter_statements.append(
"ALTER TABLE git_server_configs ADD COLUMN default_branch VARCHAR NOT NULL DEFAULT 'main'"
)
if not alter_statements:
return
try:
with bind_engine.begin() as connection:
for statement in alter_statements:
connection.execute(text(statement))
except Exception as migration_error:
logger.explore(
"GitServerConfig preference additive migration failed",
extra={"src": "database", "error": str(migration_error)},
)
# #endregion _ensure_git_server_configs_columns
# #region _ensure_auth_users_columns [C:3] [TYPE Function]
# @BRIEF Applies additive schema upgrades for auth users table.
# @PRE: bind_engine points to authentication database.
# @POST: Missing columns are added without data loss.
# @RELATION DEPENDS_ON -> [auth_engine]
def _ensure_auth_users_columns(bind_engine):
with belief_scope("_ensure_auth_users_columns"):
table_name = "users"
inspector = inspect(bind_engine)
if table_name not in inspector.get_table_names():
return
existing_columns = {
str(column.get("name") or "").strip()
for column in inspector.get_columns(table_name)
}
alter_statements = []
if "full_name" not in existing_columns:
alter_statements.append("ALTER TABLE users ADD COLUMN full_name VARCHAR")
if "is_ad_user" not in existing_columns:
alter_statements.append(
"ALTER TABLE users ADD COLUMN is_ad_user BOOLEAN NOT NULL DEFAULT FALSE"
)
if not alter_statements:
logger.reason(
"Auth users schema already up to date",
extra={"table": table_name, "columns": sorted(existing_columns)},
)
return
logger.reason(
"Applying additive auth users schema migration",
extra={"table": table_name, "statements": alter_statements},
)
try:
with bind_engine.begin() as connection:
for statement in alter_statements:
connection.execute(text(statement))
logger.reason(
"Auth users schema migration completed",
extra={
"table": table_name,
"added_columns": [
stmt.split(" ADD COLUMN ", 1)[1].split()[0]
for stmt in alter_statements
],
},
)
except Exception as migration_error:
logger.explore(
"Auth users additive migration failed",
extra={"src": "database", "error": str(migration_error)},
)
raise
# #endregion _ensure_auth_users_columns
# #region ensure_connection_configs_table [C:3] [TYPE Function]
# @BRIEF Ensures the external connection registry table exists in the main database.
# @PRE: bind_engine points to the application database.
# @POST: connection_configs table exists without dropping existing data.
# @RELATION DEPENDS_ON -> [ConnectionConfig]
def ensure_connection_configs_table(bind_engine):
with belief_scope("ensure_connection_configs_table"):
try:
ConnectionConfig.__table__.create(bind=bind_engine, checkfirst=True)
except Exception as migration_error:
logger.explore(
"ConnectionConfig table ensure failed",
extra={"src": "database", "error": str(migration_error)},
)
raise
# #endregion ensure_connection_configs_table
# #region _ensure_filter_source_enum_values [C:3] [TYPE Function]
# @BRIEF Adds missing FilterSource enum values to the PostgreSQL native filtersource type.
# @PRE: bind_engine points to application database with imported_filters table.
# @POST: New enum values are available without data loss.
# @RELATION DEPENDS_ON -> [engine]
def _ensure_filter_source_enum_values(bind_engine):
with belief_scope("_ensure_filter_source_enum_values"):
try:
with bind_engine.connect() as connection:
# Check if the native enum type exists
result = connection.execute(
text(
"SELECT t.typname FROM pg_type t "
"JOIN pg_namespace n ON t.typnamespace = n.oid "
"WHERE t.typname = 'filtersource' AND n.nspname = 'public'"
)
)
if result.fetchone() is None:
logger.reason(
"filtersource enum type does not exist yet; skipping migration"
)
return
# Get existing enum values
result = connection.execute(
text(
"SELECT e.enumlabel FROM pg_enum e "
"JOIN pg_type t ON e.enumtypid = t.oid "
"WHERE t.typname = 'filtersource' "
"ORDER BY e.enumsortorder"
)
)
existing_values = {row[0] for row in result.fetchall()}
required_values = ["SUPERSET_PERMALINK", "SUPERSET_NATIVE_FILTERS_KEY"]
missing_values = [
v for v in required_values if v not in existing_values
]
if not missing_values:
logger.reason(
"filtersource enum already up to date",
extra={"existing": sorted(existing_values)},
)
return
logger.reason(
"Adding missing values to filtersource enum",
extra={"missing": missing_values},
)
for value in missing_values:
connection.execute(
text(
f"ALTER TYPE filtersource ADD VALUE IF NOT EXISTS '{value}'"
)
)
connection.commit()
logger.reason(
"filtersource enum migration completed",
extra={"added": missing_values},
)
except Exception as migration_error:
logger.explore(
"FilterSource enum additive migration failed",
extra={"src": "database", "error": str(migration_error)},
)
# #endregion _ensure_filter_source_enum_values
# #region _ensure_dataset_review_session_columns [C:4] [TYPE Function]
# @BRIEF Apply additive schema upgrades for dataset review persistence required by optimistic-lock and recovery metadata semantics.
# @PRE: bind_engine points to the application database where dataset review tables are stored.
# @POST: Missing additive columns across legacy dataset review tables are created without removing existing data.
# @SIDE_EFFECT: Executes ALTER TABLE statements against dataset review tables in the application database.
# @RELATION DEPENDS_ON -> [DatasetReviewSession]
# @RELATION DEPENDS_ON -> [ImportedFilter]
def _ensure_translation_jobs_columns(bind_engine):
with belief_scope("_ensure_translation_jobs_columns"):
table_name = "translation_jobs"
inspector = inspect(bind_engine)
if table_name not in inspector.get_table_names():
return
existing_columns = {
str(column.get("name") or "").strip()
for column in inspector.get_columns(table_name)
}
if "environment_id" not in existing_columns:
try:
with bind_engine.begin() as connection:
connection.execute(
text(
"ALTER TABLE translation_jobs "
"ADD COLUMN environment_id VARCHAR"
)
)
logger.reflect(
"Added environment_id column to translation_jobs",
)
except Exception as migration_error:
logger.explore(
"Failed to add environment_id to translation_jobs",
extra={"error": str(migration_error)},
)
raise
if "target_database_id" not in existing_columns:
try:
with bind_engine.begin() as connection:
connection.execute(
text(
"ALTER TABLE translation_jobs "
"ADD COLUMN target_database_id VARCHAR"
)
)
logger.reflect(
"Added target_database_id column to translation_jobs",
)
except Exception as migration_error:
logger.explore(
"Failed to add target_database_id to translation_jobs",
extra={"error": str(migration_error)},
)
def _ensure_dataset_review_session_columns(bind_engine):
with belief_scope("_ensure_dataset_review_session_columns"):
inspector = inspect(bind_engine)
existing_tables = set(inspector.get_table_names())
migration_plan = {
"dataset_review_sessions": [
(
"version",
"ALTER TABLE dataset_review_sessions "
"ADD COLUMN version INTEGER NOT NULL DEFAULT 0",
)
],
"imported_filters": [
(
"raw_value_masked",
"ALTER TABLE imported_filters "
"ADD COLUMN raw_value_masked BOOLEAN NOT NULL DEFAULT FALSE",
)
],
}
for table_name, planned_columns in migration_plan.items():
if table_name not in existing_tables:
logger.reason(
"Dataset review table does not exist yet; skipping additive schema migration",
extra={"table": table_name},
)
continue
existing_columns = {
str(column.get("name") or "").strip()
for column in inspector.get_columns(table_name)
}
alter_statements = [
statement
for column_name, statement in planned_columns
if column_name not in existing_columns
]
if not alter_statements:
logger.reason(
"Dataset review table schema already up to date",
extra={"table": table_name, "columns": sorted(existing_columns)},
)
continue
logger.reason(
"Applying additive dataset review schema migration",
extra={"table": table_name, "statements": alter_statements},
)
try:
with bind_engine.begin() as connection:
for statement in alter_statements:
connection.execute(text(statement))
logger.reflect(
"Dataset review additive schema migration completed",
extra={
"table": table_name,
"added_columns": [
stmt.split(" ADD COLUMN ", 1)[1].split()[0]
for stmt in alter_statements
],
},
)
except Exception as migration_error:
logger.explore(
"Dataset review additive schema migration failed",
extra={"table": table_name, "error": str(migration_error)},
)
raise
# #endregion _ensure_dataset_review_session_columns
# #region init_db [C:3] [TYPE Function]
# @BRIEF Initializes the database by creating all tables.
# @PRE: engine, tasks_engine and auth_engine are initialized.
# @POST: Database tables created in all databases.
# @SIDE_EFFECT: Creates physical database files if they don't exist.
# @RELATION CALLS -> [ensure_connection_configs_table]
# @RELATION CALLS -> [_ensure_filter_source_enum_values]
# @RELATION CALLS -> [_ensure_dataset_review_session_columns]
def init_db():
with belief_scope("init_db"):
Base.metadata.create_all(bind=engine)
Base.metadata.create_all(bind=tasks_engine)
Base.metadata.create_all(bind=auth_engine)
_ensure_user_dashboard_preferences_columns(engine)
_ensure_llm_validation_results_columns(engine)
_ensure_user_dashboard_preferences_health_columns(engine)
_ensure_git_server_configs_columns(engine)
_ensure_auth_users_columns(auth_engine)
ensure_connection_configs_table(engine)
_ensure_filter_source_enum_values(engine)
_ensure_dataset_review_session_columns(engine)
_ensure_translation_jobs_columns(engine)
# #endregion init_db
# #region get_db [C:3] [TYPE Function]
# @BRIEF Dependency for getting a database session.
# @PRE: SessionLocal is initialized.
# @POST: Session is closed after use.
# @RELATION DEPENDS_ON -> [SessionLocal]
def get_db():
with belief_scope("get_db"):
db = SessionLocal()
try:
yield db
finally:
db.close()
# #endregion get_db
# #region get_tasks_db [C:3] [TYPE Function]
# @BRIEF Dependency for getting a tasks database session.
# @PRE: TasksSessionLocal is initialized.
# @POST: Session is closed after use.
# @RELATION DEPENDS_ON -> [TasksSessionLocal]
def get_tasks_db():
with belief_scope("get_tasks_db"):
db = TasksSessionLocal()
try:
yield db
finally:
db.close()
# #endregion get_tasks_db
# #region get_auth_db [C:3] [TYPE Function]
# @BRIEF Dependency for getting an authentication database session.
# @PRE: AuthSessionLocal is initialized.
# @POST: Session is closed after use.
# @DATA_CONTRACT: None -> Output[sqlalchemy.orm.Session]
# @RELATION DEPENDS_ON -> [AuthSessionLocal]
def get_auth_db():
with belief_scope("get_auth_db"):
db = AuthSessionLocal()
try:
yield db
finally:
db.close()
# #endregion get_auth_db
# #endregion DatabaseModule
Reference in New Issue View Git Blame Copy Permalink