intermittent commit
This commit is contained in:
@@ -3,6 +3,7 @@ from __future__ import annotations
|
||||
from collections.abc import Mapping
|
||||
from dataclasses import dataclass, replace
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
from pathlib import Path
|
||||
import re
|
||||
@@ -25,6 +26,8 @@ from govoplan_core.server.config import ManifestFactory
|
||||
from govoplan_core.server.registry import available_module_manifests, build_platform_registry
|
||||
from govoplan_core.settings import settings
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Historic development databases could be created partly through Alembic and
|
||||
# partly through Base.metadata.create_all(). In that state Alembic still says
|
||||
# "2c..." while the 3d/4e file-storage tables already exist, so a normal
|
||||
@@ -166,6 +169,15 @@ _CREATE_ALL_THROUGH_HIERARCHICAL_INDEXES = {
|
||||
}
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _LegacyCreateAllSchemaState:
|
||||
current: str | None
|
||||
has_no_revision: bool
|
||||
has_file_storage: bool
|
||||
has_file_folders: bool
|
||||
has_create_all_hierarchical_schema: bool
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class MigrationResult:
|
||||
previous_revision: str | None
|
||||
@@ -233,10 +245,7 @@ def run_registered_module_migration_tasks(
|
||||
dry_run: bool = False,
|
||||
manifest_factories: tuple[ManifestFactory, ...] = (),
|
||||
) -> tuple[dict[str, object], ...]:
|
||||
active_phases = tuple(dict.fromkeys(str(item).strip() for item in phases if str(item).strip()))
|
||||
invalid_phases = tuple(phase for phase in active_phases if phase not in MIGRATION_TASK_PHASES)
|
||||
if invalid_phases:
|
||||
raise ValueError("Unsupported module migration task phase(s): " + ", ".join(invalid_phases))
|
||||
active_phases = _normalized_migration_task_phases(phases)
|
||||
url = database_url or settings.database_url
|
||||
registry = _registered_module_registry(
|
||||
database_url=url,
|
||||
@@ -244,91 +253,192 @@ def run_registered_module_migration_tasks(
|
||||
manifest_factories=manifest_factories,
|
||||
)
|
||||
manifests = {manifest.id: manifest for manifest in registry.manifests()}
|
||||
ordered_ids = tuple(dict.fromkeys([
|
||||
*(str(item).strip() for item in (migration_module_order or ()) if str(item).strip()),
|
||||
*manifests.keys(),
|
||||
]))
|
||||
ordered_ids = _ordered_migration_task_module_ids(migration_module_order, manifests)
|
||||
records: list[dict[str, object]] = []
|
||||
database = get_database()
|
||||
with database.SessionLocal() as session:
|
||||
for phase in active_phases:
|
||||
for module_id in ordered_ids:
|
||||
manifest = manifests.get(module_id)
|
||||
if manifest is None or manifest.migration_spec is None:
|
||||
continue
|
||||
for task in manifest.migration_spec.migration_tasks:
|
||||
if task.phase != phase:
|
||||
continue
|
||||
record: dict[str, object] = {
|
||||
"module_id": manifest.id,
|
||||
"task_id": task.task_id,
|
||||
"phase": task.phase,
|
||||
"summary": task.summary,
|
||||
"task_version": task.task_version,
|
||||
"safety": task.safety,
|
||||
"idempotent": task.idempotent,
|
||||
"dry_run": dry_run,
|
||||
}
|
||||
if task.timeout_seconds is not None:
|
||||
record["timeout_seconds"] = task.timeout_seconds
|
||||
if not task.idempotent:
|
||||
record.update({"status": "blocked", "message": "Task is not idempotent."})
|
||||
records.append(record)
|
||||
raise ModuleMigrationTaskExecutionError(
|
||||
f"Module migration task {manifest.id}/{task.task_id} is not idempotent.",
|
||||
records=tuple(records),
|
||||
)
|
||||
if dry_run:
|
||||
record.update({"status": "skipped", "message": "Dry run; executor was not called."})
|
||||
records.append(record)
|
||||
continue
|
||||
if task.executor is None:
|
||||
record.update({"status": "blocked", "message": "Task has no executor."})
|
||||
records.append(record)
|
||||
raise ModuleMigrationTaskExecutionError(
|
||||
f"Module migration task {manifest.id}/{task.task_id} has no executor.",
|
||||
records=tuple(records),
|
||||
)
|
||||
context = ModuleMigrationTaskContext(
|
||||
module_id=manifest.id,
|
||||
task_id=task.task_id,
|
||||
phase=task.phase,
|
||||
database_url=url,
|
||||
target_version=manifest.version,
|
||||
session=session,
|
||||
dry_run=dry_run,
|
||||
metadata=task.metadata,
|
||||
)
|
||||
try:
|
||||
result = task.executor(context)
|
||||
normalized = _normalize_migration_task_result(result)
|
||||
except Exception as exc:
|
||||
session.rollback()
|
||||
record.update({
|
||||
"status": "blocked",
|
||||
"message": f"{type(exc).__name__}: {exc}",
|
||||
})
|
||||
records.append(record)
|
||||
raise ModuleMigrationTaskExecutionError(
|
||||
f"Module migration task {manifest.id}/{task.task_id} failed.",
|
||||
records=tuple(records),
|
||||
) from exc
|
||||
record.update({
|
||||
"status": normalized.status,
|
||||
"message": normalized.message,
|
||||
"details": _jsonable_migration_task_details(normalized.details),
|
||||
})
|
||||
records.append(record)
|
||||
if normalized.status == "blocked":
|
||||
session.rollback()
|
||||
raise ModuleMigrationTaskExecutionError(
|
||||
normalized.message or f"Module migration task {manifest.id}/{task.task_id} blocked migration.",
|
||||
records=tuple(records),
|
||||
)
|
||||
session.commit()
|
||||
for manifest, task in _iter_phase_migration_tasks(phase, ordered_ids=ordered_ids, manifests=manifests):
|
||||
_run_registered_module_migration_task(
|
||||
session,
|
||||
manifest=manifest,
|
||||
task=task,
|
||||
database_url=url,
|
||||
dry_run=dry_run,
|
||||
records=records,
|
||||
)
|
||||
return tuple(records)
|
||||
|
||||
|
||||
def _normalized_migration_task_phases(phases: tuple[str, ...] | list[str]) -> tuple[str, ...]:
|
||||
active_phases = tuple(dict.fromkeys(str(item).strip() for item in phases if str(item).strip()))
|
||||
invalid_phases = tuple(phase for phase in active_phases if phase not in MIGRATION_TASK_PHASES)
|
||||
if invalid_phases:
|
||||
raise ValueError("Unsupported module migration task phase(s): " + ", ".join(invalid_phases))
|
||||
return active_phases
|
||||
|
||||
|
||||
def _ordered_migration_task_module_ids(
|
||||
migration_module_order: tuple[str, ...] | list[str] | None,
|
||||
manifests: Mapping[str, object],
|
||||
) -> tuple[str, ...]:
|
||||
return tuple(dict.fromkeys([
|
||||
*(str(item).strip() for item in (migration_module_order or ()) if str(item).strip()),
|
||||
*manifests.keys(),
|
||||
]))
|
||||
|
||||
|
||||
def _iter_phase_migration_tasks(
|
||||
phase: str,
|
||||
*,
|
||||
ordered_ids: tuple[str, ...],
|
||||
manifests: Mapping[str, object],
|
||||
) -> Iterable[tuple[object, object]]:
|
||||
for module_id in ordered_ids:
|
||||
manifest = manifests.get(module_id)
|
||||
migration_spec = getattr(manifest, "migration_spec", None)
|
||||
if manifest is None or migration_spec is None:
|
||||
continue
|
||||
for task in migration_spec.migration_tasks:
|
||||
if task.phase == phase:
|
||||
yield manifest, task
|
||||
|
||||
|
||||
def _run_registered_module_migration_task(
|
||||
session: object,
|
||||
*,
|
||||
manifest: object,
|
||||
task: object,
|
||||
database_url: str,
|
||||
dry_run: bool,
|
||||
records: list[dict[str, object]],
|
||||
) -> None:
|
||||
record = _migration_task_record(manifest, task, dry_run=dry_run)
|
||||
_validate_migration_task_can_run(manifest, task, record, records, dry_run=dry_run)
|
||||
if dry_run:
|
||||
record.update({"status": "skipped", "message": "Dry run; executor was not called."})
|
||||
records.append(record)
|
||||
return
|
||||
normalized = _execute_module_migration_task(
|
||||
session,
|
||||
manifest=manifest,
|
||||
task=task,
|
||||
record=record,
|
||||
records=records,
|
||||
database_url=database_url,
|
||||
dry_run=dry_run,
|
||||
)
|
||||
record.update({
|
||||
"status": normalized.status,
|
||||
"message": normalized.message,
|
||||
"details": _jsonable_migration_task_details(normalized.details),
|
||||
})
|
||||
records.append(record)
|
||||
if normalized.status == "blocked":
|
||||
session.rollback()
|
||||
raise ModuleMigrationTaskExecutionError(
|
||||
normalized.message or f"Module migration task {manifest.id}/{task.task_id} blocked migration.",
|
||||
records=tuple(records),
|
||||
)
|
||||
session.commit()
|
||||
|
||||
|
||||
def _migration_task_record(manifest: object, task: object, *, dry_run: bool) -> dict[str, object]:
|
||||
record: dict[str, object] = {
|
||||
"module_id": manifest.id,
|
||||
"task_id": task.task_id,
|
||||
"phase": task.phase,
|
||||
"summary": task.summary,
|
||||
"task_version": task.task_version,
|
||||
"safety": task.safety,
|
||||
"idempotent": task.idempotent,
|
||||
"dry_run": dry_run,
|
||||
}
|
||||
if task.timeout_seconds is not None:
|
||||
record["timeout_seconds"] = task.timeout_seconds
|
||||
return record
|
||||
|
||||
|
||||
def _validate_migration_task_can_run(
|
||||
manifest: object,
|
||||
task: object,
|
||||
record: dict[str, object],
|
||||
records: list[dict[str, object]],
|
||||
*,
|
||||
dry_run: bool,
|
||||
) -> None:
|
||||
if not task.idempotent:
|
||||
_block_migration_task(
|
||||
manifest,
|
||||
task,
|
||||
record,
|
||||
records,
|
||||
message="Task is not idempotent.",
|
||||
error=f"Module migration task {manifest.id}/{task.task_id} is not idempotent.",
|
||||
)
|
||||
if not dry_run and task.executor is None:
|
||||
_block_migration_task(
|
||||
manifest,
|
||||
task,
|
||||
record,
|
||||
records,
|
||||
message="Task has no executor.",
|
||||
error=f"Module migration task {manifest.id}/{task.task_id} has no executor.",
|
||||
)
|
||||
|
||||
|
||||
def _block_migration_task(
|
||||
manifest: object,
|
||||
task: object,
|
||||
record: dict[str, object],
|
||||
records: list[dict[str, object]],
|
||||
*,
|
||||
message: str,
|
||||
error: str,
|
||||
) -> None:
|
||||
record.update({"status": "blocked", "message": message})
|
||||
records.append(record)
|
||||
raise ModuleMigrationTaskExecutionError(
|
||||
error,
|
||||
records=tuple(records),
|
||||
)
|
||||
|
||||
|
||||
def _execute_module_migration_task(
|
||||
session: object,
|
||||
*,
|
||||
manifest: object,
|
||||
task: object,
|
||||
record: dict[str, object],
|
||||
records: list[dict[str, object]],
|
||||
database_url: str,
|
||||
dry_run: bool,
|
||||
) -> ModuleMigrationTaskResult:
|
||||
context = ModuleMigrationTaskContext(
|
||||
module_id=manifest.id,
|
||||
task_id=task.task_id,
|
||||
phase=task.phase,
|
||||
database_url=database_url,
|
||||
target_version=manifest.version,
|
||||
session=session,
|
||||
dry_run=dry_run,
|
||||
metadata=task.metadata,
|
||||
)
|
||||
try:
|
||||
return _normalize_migration_task_result(task.executor(context))
|
||||
except Exception as exc:
|
||||
session.rollback()
|
||||
record.update({
|
||||
"status": "blocked",
|
||||
"message": f"{type(exc).__name__}: {exc}",
|
||||
})
|
||||
records.append(record)
|
||||
raise ModuleMigrationTaskExecutionError(
|
||||
f"Module migration task {manifest.id}/{task.task_id} failed.",
|
||||
records=tuple(records),
|
||||
) from exc
|
||||
|
||||
|
||||
def _normalize_migration_task_result(result: ModuleMigrationTaskResult | None) -> ModuleMigrationTaskResult:
|
||||
if result is None:
|
||||
return ModuleMigrationTaskResult()
|
||||
@@ -524,18 +634,21 @@ def _backfill_user_lock_state_for_create_all_schema(database_url: str) -> None:
|
||||
|
||||
def _row_count(connection, table_name: str) -> int:
|
||||
quoted = _quoted_table_name(connection, table_name)
|
||||
return int(connection.execute(text(f"SELECT COUNT(*) FROM {quoted}")).scalar_one()) # nosemgrep: python.sqlalchemy.security.audit.avoid-sqlalchemy-text.avoid-sqlalchemy-text
|
||||
statement = text(f"SELECT COUNT(*) FROM {quoted}") # nosec B608 # nosemgrep: python.sqlalchemy.security.audit.avoid-sqlalchemy-text.avoid-sqlalchemy-text
|
||||
return int(connection.execute(statement).scalar_one())
|
||||
|
||||
|
||||
def _drop_table(connection, table_name: str) -> None:
|
||||
quoted = _quoted_table_name(connection, table_name)
|
||||
connection.execute(text(f"DROP TABLE {quoted}")) # nosemgrep: python.sqlalchemy.security.audit.avoid-sqlalchemy-text.avoid-sqlalchemy-text
|
||||
statement = text(f"DROP TABLE {quoted}") # nosec B608 # nosemgrep: python.sqlalchemy.security.audit.avoid-sqlalchemy-text.avoid-sqlalchemy-text
|
||||
connection.execute(statement)
|
||||
|
||||
|
||||
def _rename_table(connection, old_name: str, new_name: str) -> None:
|
||||
quoted_old = _quoted_table_name(connection, old_name)
|
||||
quoted_new = _quoted_table_name(connection, new_name)
|
||||
connection.execute(text(f"ALTER TABLE {quoted_old} RENAME TO {quoted_new}")) # nosemgrep: python.sqlalchemy.security.audit.avoid-sqlalchemy-text.avoid-sqlalchemy-text
|
||||
statement = text(f"ALTER TABLE {quoted_old} RENAME TO {quoted_new}") # nosec B608 # nosemgrep: python.sqlalchemy.security.audit.avoid-sqlalchemy-text.avoid-sqlalchemy-text
|
||||
connection.execute(statement)
|
||||
|
||||
|
||||
def _quoted_table_name(connection, table_name: str) -> str:
|
||||
@@ -715,7 +828,8 @@ def reconcile_covered_alembic_dependency_heads(
|
||||
for revision_id in current:
|
||||
try:
|
||||
revision = script.get_revision(revision_id)
|
||||
except Exception:
|
||||
except Exception as exc:
|
||||
logger.debug("Skipping Alembic revision %s while pruning dependency heads: %s", revision_id, exc, exc_info=True)
|
||||
continue
|
||||
ancestors = {
|
||||
item.revision
|
||||
@@ -749,58 +863,68 @@ def reconcile_legacy_create_all_schema(
|
||||
"""
|
||||
|
||||
url = database_url or settings.database_url
|
||||
engine = create_engine(url)
|
||||
try:
|
||||
with engine.connect() as connection:
|
||||
heads = MigrationContext.configure(connection).get_current_heads()
|
||||
current = heads[0] if len(heads) == 1 else None
|
||||
has_no_revision = len(heads) == 0
|
||||
schema = inspect(connection)
|
||||
tables = set(schema.get_table_names())
|
||||
|
||||
has_file_storage = _FILE_STORAGE_TABLES.issubset(tables) and all(
|
||||
_has_columns(schema, table, _FILE_STORAGE_COLUMNS[table])
|
||||
for table in _FILE_STORAGE_TABLES
|
||||
)
|
||||
has_file_folders = _FILE_FOLDER_TABLES.issubset(tables) and _has_columns(
|
||||
schema,
|
||||
"file_folders",
|
||||
_FILE_STORAGE_COLUMNS["file_folders"],
|
||||
)
|
||||
has_create_all_hierarchical_schema = _has_create_all_schema_through_hierarchical_settings(schema, tables)
|
||||
finally:
|
||||
engine.dispose()
|
||||
|
||||
target: str | None = None
|
||||
if current == REVISION_AUTH_RBAC and has_file_storage and has_file_folders:
|
||||
target = REVISION_FILE_FOLDERS
|
||||
elif current == REVISION_AUTH_RBAC and has_file_storage:
|
||||
target = REVISION_FILE_STORAGE
|
||||
elif current == REVISION_FILE_STORAGE and has_file_folders:
|
||||
target = REVISION_FILE_FOLDERS
|
||||
elif current == REVISION_FILE_FOLDERS and has_create_all_hierarchical_schema:
|
||||
# Development DBs may be stamped at 4e after earlier create_all
|
||||
# reconciliation even though later tables/columns were already created
|
||||
# by a newer model set. Skip only when that complete known schema is
|
||||
# present, then let newer migrations such as mail credential usernames
|
||||
# run normally.
|
||||
_backfill_user_lock_state_for_create_all_schema(url)
|
||||
target = REVISION_HIERARCHICAL_SETTINGS
|
||||
elif has_no_revision and has_create_all_hierarchical_schema:
|
||||
_backfill_user_lock_state_for_create_all_schema(url)
|
||||
target = REVISION_HIERARCHICAL_SETTINGS
|
||||
elif has_no_revision and has_file_storage and has_file_folders:
|
||||
# This is the other create_all-only development shape. The strict
|
||||
# column checks above ensure that we only stamp a complete known schema.
|
||||
target = REVISION_FILE_FOLDERS
|
||||
|
||||
state = _legacy_create_all_schema_state(url)
|
||||
target = _legacy_create_all_reconciliation_target(state)
|
||||
if target is None:
|
||||
return None
|
||||
if target == REVISION_HIERARCHICAL_SETTINGS:
|
||||
_backfill_user_lock_state_for_create_all_schema(url)
|
||||
|
||||
command.stamp(alembic_config(database_url=url, migration_track=migration_track), target)
|
||||
return target
|
||||
|
||||
|
||||
def _legacy_create_all_schema_state(database_url: str) -> _LegacyCreateAllSchemaState:
|
||||
engine = create_engine(database_url)
|
||||
try:
|
||||
with engine.connect() as connection:
|
||||
heads = MigrationContext.configure(connection).get_current_heads()
|
||||
schema = inspect(connection)
|
||||
tables = set(schema.get_table_names())
|
||||
return _legacy_create_all_schema_state_from_inspection(heads, schema, tables)
|
||||
finally:
|
||||
engine.dispose()
|
||||
|
||||
|
||||
def _legacy_create_all_schema_state_from_inspection(
|
||||
heads: tuple[str, ...],
|
||||
schema: object,
|
||||
tables: set[str],
|
||||
) -> _LegacyCreateAllSchemaState:
|
||||
has_file_storage = _FILE_STORAGE_TABLES.issubset(tables) and all(
|
||||
_has_columns(schema, table, _FILE_STORAGE_COLUMNS[table])
|
||||
for table in _FILE_STORAGE_TABLES
|
||||
)
|
||||
has_file_folders = _FILE_FOLDER_TABLES.issubset(tables) and _has_columns(
|
||||
schema,
|
||||
"file_folders",
|
||||
_FILE_STORAGE_COLUMNS["file_folders"],
|
||||
)
|
||||
return _LegacyCreateAllSchemaState(
|
||||
current=heads[0] if len(heads) == 1 else None,
|
||||
has_no_revision=len(heads) == 0,
|
||||
has_file_storage=has_file_storage,
|
||||
has_file_folders=has_file_folders,
|
||||
has_create_all_hierarchical_schema=_has_create_all_schema_through_hierarchical_settings(schema, tables),
|
||||
)
|
||||
|
||||
|
||||
def _legacy_create_all_reconciliation_target(state: _LegacyCreateAllSchemaState) -> str | None:
|
||||
if state.current == REVISION_AUTH_RBAC and state.has_file_storage and state.has_file_folders:
|
||||
return REVISION_FILE_FOLDERS
|
||||
if state.current == REVISION_AUTH_RBAC and state.has_file_storage:
|
||||
return REVISION_FILE_STORAGE
|
||||
if state.current == REVISION_FILE_STORAGE and state.has_file_folders:
|
||||
return REVISION_FILE_FOLDERS
|
||||
if state.current == REVISION_FILE_FOLDERS and state.has_create_all_hierarchical_schema:
|
||||
return REVISION_HIERARCHICAL_SETTINGS
|
||||
if state.has_no_revision and state.has_create_all_hierarchical_schema:
|
||||
return REVISION_HIERARCHICAL_SETTINGS
|
||||
if state.has_no_revision and state.has_file_storage and state.has_file_folders:
|
||||
return REVISION_FILE_FOLDERS
|
||||
return None
|
||||
|
||||
|
||||
def migrate_database(
|
||||
*,
|
||||
database_url: str | None = None,
|
||||
|
||||
100
src/govoplan_core/db/query_metrics.py
Normal file
100
src/govoplan_core/db/query_metrics.py
Normal file
@@ -0,0 +1,100 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from contextlib import contextmanager
|
||||
from contextvars import ContextVar
|
||||
from dataclasses import dataclass
|
||||
import time
|
||||
import weakref
|
||||
from collections.abc import Iterator
|
||||
|
||||
from sqlalchemy import event
|
||||
from sqlalchemy.engine import Connection, Engine
|
||||
from sqlalchemy.engine.interfaces import ExecutionContext
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class QueryMetrics:
|
||||
query_count: int = 0
|
||||
total_ms: float = 0.0
|
||||
slowest_ms: float = 0.0
|
||||
error_count: int = 0
|
||||
|
||||
def record(self, elapsed_ms: float, *, failed: bool = False) -> None:
|
||||
self.query_count += 1
|
||||
self.total_ms += elapsed_ms
|
||||
self.slowest_ms = max(self.slowest_ms, elapsed_ms)
|
||||
if failed:
|
||||
self.error_count += 1
|
||||
|
||||
|
||||
_current_metrics: ContextVar[QueryMetrics | None] = ContextVar("govoplan_query_metrics", default=None)
|
||||
_instrumented_engines: weakref.WeakSet[Engine] = weakref.WeakSet()
|
||||
_START_STACK_KEY = "govoplan_query_metric_starts"
|
||||
|
||||
|
||||
@contextmanager
|
||||
def collect_query_metrics() -> Iterator[QueryMetrics]:
|
||||
metrics = QueryMetrics()
|
||||
token = _current_metrics.set(metrics)
|
||||
try:
|
||||
yield metrics
|
||||
finally:
|
||||
_current_metrics.reset(token)
|
||||
|
||||
|
||||
def current_query_metrics() -> QueryMetrics | None:
|
||||
return _current_metrics.get()
|
||||
|
||||
|
||||
def instrument_engine(engine: Engine) -> Engine:
|
||||
if engine in _instrumented_engines:
|
||||
return engine
|
||||
|
||||
event.listen(engine, "before_cursor_execute", _before_cursor_execute)
|
||||
event.listen(engine, "after_cursor_execute", _after_cursor_execute)
|
||||
event.listen(engine, "handle_error", _handle_error)
|
||||
_instrumented_engines.add(engine)
|
||||
return engine
|
||||
|
||||
|
||||
def _before_cursor_execute(
|
||||
conn: Connection,
|
||||
_cursor: object,
|
||||
_statement: str,
|
||||
_parameters: object,
|
||||
_context: ExecutionContext,
|
||||
_executemany: bool,
|
||||
) -> None:
|
||||
if current_query_metrics() is None:
|
||||
return
|
||||
starts = conn.info.setdefault(_START_STACK_KEY, [])
|
||||
starts.append(time.perf_counter())
|
||||
|
||||
|
||||
def _after_cursor_execute(
|
||||
conn: Connection,
|
||||
_cursor: object,
|
||||
_statement: str,
|
||||
_parameters: object,
|
||||
_context: ExecutionContext,
|
||||
_executemany: bool,
|
||||
) -> None:
|
||||
_record_elapsed(conn, failed=False)
|
||||
|
||||
|
||||
def _handle_error(exception_context: object) -> None:
|
||||
conn = getattr(exception_context, "connection", None)
|
||||
if isinstance(conn, Connection):
|
||||
_record_elapsed(conn, failed=True)
|
||||
|
||||
|
||||
def _record_elapsed(conn: Connection, *, failed: bool) -> None:
|
||||
metrics = current_query_metrics()
|
||||
if metrics is None:
|
||||
return
|
||||
starts = conn.info.get(_START_STACK_KEY)
|
||||
if not starts:
|
||||
return
|
||||
started_at = starts.pop()
|
||||
elapsed_ms = (time.perf_counter() - started_at) * 1000
|
||||
metrics.record(elapsed_ms, failed=failed)
|
||||
@@ -7,6 +7,8 @@ from sqlalchemy import create_engine
|
||||
from sqlalchemy.engine import Engine
|
||||
from sqlalchemy.orm import Session, sessionmaker
|
||||
|
||||
from govoplan_core.db.query_metrics import instrument_engine
|
||||
|
||||
|
||||
def default_connect_args(database_url: str) -> dict[str, Any]:
|
||||
return {"check_same_thread": False} if database_url.startswith("sqlite") else {}
|
||||
@@ -22,13 +24,13 @@ def create_database_engine(
|
||||
merged_connect_args = dict(default_connect_args(database_url))
|
||||
if connect_args:
|
||||
merged_connect_args.update(connect_args)
|
||||
return create_engine(database_url, pool_pre_ping=pool_pre_ping, connect_args=merged_connect_args, **kwargs)
|
||||
return instrument_engine(create_engine(database_url, pool_pre_ping=pool_pre_ping, connect_args=merged_connect_args, **kwargs))
|
||||
|
||||
|
||||
class DatabaseHandle:
|
||||
def __init__(self, database_url: str, *, engine: Engine | None = None) -> None:
|
||||
self.database_url = database_url
|
||||
self.engine = engine or create_database_engine(database_url)
|
||||
self.engine = instrument_engine(engine) if engine is not None else create_database_engine(database_url)
|
||||
self.SessionLocal = sessionmaker(bind=self.engine, autoflush=False, autocommit=False, expire_on_commit=False)
|
||||
|
||||
def session(self) -> Session:
|
||||
|
||||
Reference in New Issue
Block a user