87 lines
2.6 KiB
Python
87 lines
2.6 KiB
Python
"""SQL Server connection helpers via pyodbc."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from contextlib import contextmanager
|
|
from typing import Any, Iterator
|
|
|
|
import pyodbc
|
|
|
|
from pipeline.config import SqlSettings
|
|
|
|
|
|
def _connection_string(settings: SqlSettings) -> str:
|
|
return (
|
|
f"DRIVER={{{settings.driver}}};"
|
|
f"SERVER={settings.server};"
|
|
f"DATABASE={settings.database};"
|
|
f"UID={settings.username};"
|
|
f"PWD={settings.password};"
|
|
"TrustServerCertificate=yes;"
|
|
)
|
|
|
|
|
|
def connect_raw(settings: SqlSettings) -> pyodbc.Connection:
|
|
"""
|
|
Open a new pyodbc connection without any context-manager lifecycle.
|
|
|
|
Used by the batch orchestrator to hand each worker thread its own
|
|
long-lived connection (pyodbc connections must not be shared across
|
|
threads); the caller is responsible for closing it.
|
|
"""
|
|
return pyodbc.connect(_connection_string(settings), autocommit=True)
|
|
|
|
|
|
@contextmanager
|
|
def connect(settings: SqlSettings) -> Iterator[pyodbc.Connection]:
|
|
"""Yield an open pyodbc connection that is closed on exit."""
|
|
conn = connect_raw(settings)
|
|
try:
|
|
yield conn
|
|
finally:
|
|
conn.close()
|
|
|
|
|
|
def fetch_one(conn: pyodbc.Connection, sql: str, params: tuple[Any, ...] = ()) -> pyodbc.Row | None:
|
|
cursor = conn.cursor()
|
|
cursor.execute(sql, params)
|
|
return cursor.fetchone()
|
|
|
|
|
|
def fetch_all(conn: pyodbc.Connection, sql: str, params: tuple[Any, ...] = ()) -> list[pyodbc.Row]:
|
|
cursor = conn.cursor()
|
|
cursor.execute(sql, params)
|
|
return list(cursor.fetchall())
|
|
|
|
|
|
def execute(conn: pyodbc.Connection, sql: str, params: tuple[Any, ...] = ()) -> int:
|
|
"""Run a write statement and return the affected row count (autocommit is on)."""
|
|
cursor = conn.cursor()
|
|
cursor.execute(sql, params)
|
|
return cursor.rowcount
|
|
|
|
|
|
@contextmanager
|
|
def transaction(conn: pyodbc.Connection) -> Iterator[None]:
|
|
"""
|
|
Run a block of statements atomically on a connection that normally runs
|
|
with autocommit on.
|
|
|
|
Used by `insert_advice`'s DELETE-then-INSERT pair: under concurrent
|
|
workers, autocommit would let a crash between the two statements leave a
|
|
day's advisory missing. This temporarily disables autocommit, commits on
|
|
success, rolls back on any exception, and always restores the previous
|
|
autocommit mode afterwards so the connection is safe to reuse for the
|
|
next job.
|
|
"""
|
|
previous_autocommit = conn.autocommit
|
|
conn.autocommit = False
|
|
try:
|
|
yield
|
|
conn.commit()
|
|
except Exception:
|
|
conn.rollback()
|
|
raise
|
|
finally:
|
|
conn.autocommit = previous_autocommit
|