f77c0abebd
- Enhanced `.env.example` to include `SQL_APP_DATABASE` for managing additional advisory-related tables. - Refactored `config.yaml` to remove the `crops` block, shifting to a database-driven approach for job discovery. - Updated `README.md` to clarify the new database structure and job discovery process, emphasizing the role of the `AI_agrosupport_an_pathogen_models` and `AI_agrosupport_agro_models` tables. - Modified SQL queries across multiple files to accommodate the new database schema, ensuring correct table references and improving overall data handling. - Adjusted `pipeline` logic to utilize the new configuration, enhancing the efficiency of the advisory generation process.
141 lines
4.5 KiB
Python
141 lines
4.5 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
|
|
|
|
|
|
# Tables that live in SQL_APP_DATABASE. Every other name stays in SQL_DATABASE.
|
|
_APP_TABLES = ("advisories", "products", "product_uses", "label_chunks")
|
|
|
|
|
|
def _bracket(identifier: str) -> str:
|
|
"""Quote a SQL Server identifier. A closing bracket inside the name is doubled."""
|
|
return "[" + identifier.replace("]", "]]") + "]"
|
|
|
|
|
|
def _odbc_server(server: str) -> str:
|
|
"""
|
|
ODBC Driver expects host,port. Accept host:port and leave named instances alone.
|
|
|
|
`192.168.220.50:1433` becomes `192.168.220.50,1433`. `localhost\\SQLEXPRESS`
|
|
has no colon, so it is returned unchanged.
|
|
"""
|
|
host, colon, port = server.rpartition(":")
|
|
if colon and port.isdigit() and ":" not in host:
|
|
return f"{host},{port}"
|
|
return server
|
|
|
|
|
|
def app_table(settings: SqlSettings, table: str) -> str:
|
|
"""
|
|
Name of an application table on the open connection.
|
|
|
|
Bare when `app_database` is the connected database, so a single-database
|
|
setup keeps unqualified SQL. Otherwise a three-part name,
|
|
`[Agrosupport].dbo.[advisories]`, so one DBMeteo connection can read the
|
|
other database on the same instance.
|
|
"""
|
|
if table not in _APP_TABLES:
|
|
allowed = ", ".join(_APP_TABLES)
|
|
raise ValueError(f"Unknown application table {table!r}. Expected one of {allowed}.")
|
|
if (
|
|
not settings.app_database
|
|
or settings.app_database.casefold() == settings.database.casefold()
|
|
):
|
|
return table
|
|
return f"{_bracket(settings.app_database)}.dbo.{_bracket(table)}"
|
|
|
|
|
|
def qualify_app_sql(settings: SqlSettings, sql: str) -> str:
|
|
"""
|
|
Replace `{advisories}`, `{products}`, `{product_uses}`, and `{label_chunks}`.
|
|
|
|
Call this before `.format(placeholders=...)`. Only those four tokens are
|
|
substituted, so a `{placeholders}` brace list is left for the later format.
|
|
"""
|
|
for table in _APP_TABLES:
|
|
sql = sql.replace("{" + table + "}", app_table(settings, table))
|
|
return sql
|
|
|
|
|
|
def _connection_string(settings: SqlSettings) -> str:
|
|
return (
|
|
f"DRIVER={{{settings.driver}}};"
|
|
f"SERVER={_odbc_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
|