"""Run parameterized SQL queries and return their rows as typed objects."""
import dataclasses
import datetime
import logging
import time
import uuid
from contextlib import contextmanager
from decimal import Decimal
from typing import Any, Callable, Dict, Generic, Iterator, List, Optional, Sequence, Tuple, Type, TypeVar
import pyodbc
from pydantic import BaseModel
from ace.config import settings
T = TypeVar("T")
SCALAR_TYPES = (int, float, str, bool, bytes, Decimal, datetime.date, datetime.time, uuid.UUID)
_READ_SESSION_USER = "SELECT CAST(SESSION_CONTEXT(N'UserId') AS INT)"
_SET_SESSION_USER = "EXEC sp_set_session_context 'UserId', ?"
[docs]
class SqlMappingError(ValueError):
"""A row of a query cannot be turned into the query's result type."""
# ---------------------------------------------------------------------------- result types
@dataclasses.dataclass(frozen=True)
class _Field:
key: str # column looked up, and the keyword the result type is built with
name: str # attribute name, also looked up when it differs from key (pydantic alias)
required: bool
def _annotated_fields(result_type: type) -> List[_Field]:
names: Dict[str, None] = {}
for klass in reversed(result_type.__mro__[:-1]):
for name, annotation in vars(klass).get("__annotations__", {}).items():
if not str(annotation).replace("typing.", "").startswith("ClassVar") and not name.startswith("_"):
names[name] = None
return [_Field(name, name, not hasattr(result_type, name)) for name in names]
def _describe(result_type: type) -> Tuple[List[_Field], Callable[[Dict[str, Any]], Any]]:
"""Return the fields of a result type and a function that builds it."""
if not isinstance(result_type, type):
raise TypeError(f"{result_type!r} is not a supported result type: pass a class, for example a dataclass")
if dataclasses.is_dataclass(result_type):
fields = [
_Field(f.name, f.name, f.default is dataclasses.MISSING and f.default_factory is dataclasses.MISSING)
for f in dataclasses.fields(result_type)
if f.init
]
return fields, lambda values: result_type(**values)
if issubclass(result_type, BaseModel):
fields = []
for name, info in result_type.model_fields.items():
key = info.validation_alias if isinstance(info.validation_alias, str) else (info.alias or name)
fields.append(_Field(key, name, info.is_required()))
return fields, result_type.model_validate
if issubclass(result_type, tuple) and hasattr(result_type, "_fields"):
defaults = getattr(result_type, "_field_defaults", {})
fields = [_Field(name, name, name not in defaults) for name in result_type._fields]
return fields, lambda values: result_type(**values)
fields = _annotated_fields(result_type)
if not fields:
raise TypeError(
f"{result_type.__name__} is not a supported result type: use a dataclass, a pydantic model, "
f"a NamedTuple, an annotated class, dict, tuple, a scalar type or None"
)
if result_type.__init__ is object.__init__:
def build(values: Dict[str, Any]) -> Any:
instance = result_type()
for name, value in values.items():
setattr(instance, name, value)
return instance
return fields, build
return fields, lambda values: result_type(**values)
def _positions(query_name: str, result_type: type, fields: List[_Field], columns: List[str], strict: bool) -> List[Tuple[str, int]]:
exact: Dict[str, List[int]] = {}
lowered: Dict[str, List[int]] = {}
for index, column in enumerate(columns):
exact.setdefault(column, []).append(index)
lowered.setdefault(column.lower(), []).append(index)
positions, used, missing = [], set(), []
for field in fields:
indexes = None
for candidate in dict.fromkeys((field.key, field.name)):
indexes = exact.get(candidate) or lowered.get(candidate.lower())
if indexes:
break
if not indexes:
if field.required:
missing.append(field.key)
continue
if len(indexes) > 1:
raise SqlMappingError(
f"Query {query_name}: column {columns[indexes[0]]} is returned {len(indexes)} times; "
f"give the columns distinct aliases"
)
positions.append((field.key, indexes[0]))
used.add(indexes[0])
if missing:
raise SqlMappingError(
f"Query {query_name}: no column for the required field(s) {', '.join(missing)} of "
f"{result_type.__name__} (columns: {', '.join(columns)})"
)
if strict:
extra = [column for index, column in enumerate(columns) if index not in used]
if extra:
raise SqlMappingError(
f"Query {query_name}: column(s) {', '.join(extra)} match no field of {result_type.__name__}"
)
return positions
def _to_scalar(query_name: str, result_type: type, value: Any) -> Any:
if value is None or isinstance(value, result_type):
return value
if result_type is int and isinstance(value, (Decimal, float)) and value != int(value):
raise SqlMappingError(f"Query {query_name}: a non-integral {type(value).__name__} cannot be returned as int")
try:
return result_type(value)
except (TypeError, ValueError, ArithmeticError) as e:
raise SqlMappingError(
f"Query {query_name}: a {type(value).__name__} value cannot be returned as {result_type.__name__}"
) from e
def _is_scalar_type(result_type: Any) -> bool:
return isinstance(result_type, type) and issubclass(result_type, SCALAR_TYPES)
# ---------------------------------------------------------------------------- cursors
@contextmanager
def _cursor(db_connection: Any) -> Iterator[Any]:
"""Open a cursor on the connection and close it on exit."""
if hasattr(db_connection, "conn"):
connection = db_connection.conn
if connection is None:
raise RuntimeError("The database connection is closed")
else:
connection = db_connection
cursor = connection.cursor()
try:
yield cursor
finally:
try:
cursor.close()
except pyodbc.ProgrammingError:
pass
def _name_the_query(error: BaseException, query_name: str) -> None:
"""Attach the query name to a database error without changing its type."""
note = f"SQL query: {query_name}"
if hasattr(error, "add_note"): # Python 3.11+
error.add_note(note)
else:
error.args = (*error.args, note)
def _columns(query_name: str, cursor: Any) -> List[str]:
if cursor.description is None:
raise SqlMappingError(f"Query {query_name} returned no rows to read; use execute() for statements")
return [column[0] for column in cursor.description]
# ---------------------------------------------------------------------------- queries
[docs]
class SqlQuery(Generic[T]):
"""
Named SQL query whose rows are returned as result_type.
Args:
name: Query name used in errors and logs
sql: SQL text with ? placeholders
result_type: Row type (dataclass, pydantic model, dict, tuple, scalar); None returns raw rows
strict: Reject columns that match no field
"""
[docs]
def __init__(self, name: str, sql: str, result_type: Optional[Type[T]] = None, *, strict: bool = False):
if not name or not isinstance(name, str):
raise ValueError("A query needs a name")
if not sql or not isinstance(sql, str):
raise ValueError(f"Query {name} has no SQL")
self.name = name
self.sql = sql
self.result_type = result_type
self.strict = strict
# checked here so a wrong type fails when the module is imported, not at the first call
self._described = None
if result_type not in (None, dict, tuple) and not _is_scalar_type(result_type):
self._described = _describe(result_type)
def __repr__(self) -> str:
return f"SqlQuery({self.name!r})"
[docs]
def fetch_all(self, db_connection: Any, *params: Any) -> List[T]:
"""Return all rows."""
return self._run(db_connection, params, self._read_all)
[docs]
def fetch_one(self, db_connection: Any, *params: Any) -> Optional[T]:
"""Return the first row, or None."""
return self._run(db_connection, params, self._read_one)
[docs]
def scalar(self, db_connection: Any, *params: Any) -> Any:
"""Return the first column of the first row, or None."""
return self._run(db_connection, params, self._read_scalar)
[docs]
def execute(self, db_connection: Any, *params: Any) -> int:
"""Run a statement without result rows and return the row count."""
return self._run(db_connection, params, lambda cursor: cursor.rowcount)
# -------------------------------------------------------------- internals
def _run(self, db_connection: Any, params: Sequence[Any], read: Callable[[Any], Any]) -> Any:
started = time.perf_counter()
with _cursor(db_connection) as cursor:
try:
cursor.execute(self.sql, *params)
result = read(cursor)
except pyodbc.Error as e:
_name_the_query(e, self.name)
raise
logging.debug("SQL query %s ran in %.1f ms", self.name, (time.perf_counter() - started) * 1000)
return result
def _mapper(self, cursor: Any) -> Callable[[Any], Any]:
result_type = self.result_type
if result_type is None:
return lambda row: row
columns = _columns(self.name, cursor)
if result_type is tuple:
return tuple
if result_type is dict:
duplicates = sorted({c for c in columns if columns.count(c) > 1})
if duplicates:
raise SqlMappingError(
f"Query {self.name}: column(s) {', '.join(duplicates)} returned more than once; "
f"give the columns distinct aliases"
)
return lambda row: dict(zip(columns, row))
if _is_scalar_type(result_type):
if self.strict and len(columns) != 1:
raise SqlMappingError(f"Query {self.name}: {len(columns)} columns returned for a {result_type.__name__}")
return lambda row: _to_scalar(self.name, result_type, row[0])
fields, build = self._described
positions = _positions(self.name, result_type, fields, columns, self.strict)
return lambda row: build({key: row[index] for key, index in positions})
def _read_all(self, cursor: Any) -> List[Any]:
rows = cursor.fetchall()
mapper = self._mapper(cursor)
return [mapper(row) for row in rows]
def _read_one(self, cursor: Any) -> Any:
row = cursor.fetchone()
if row is None:
return None
return self._mapper(cursor)(row)
def _read_scalar(self, cursor: Any) -> Any:
_columns(self.name, cursor)
row = cursor.fetchone()
if row is None:
return None
if _is_scalar_type(self.result_type):
return _to_scalar(self.name, self.result_type, row[0])
return row[0]
# ---------------------------------------------------------------------------- session context
[docs]
def current_session_user(db_connection: Any) -> Optional[int]:
"""Return the session UserId used by row-level security, or None."""
with _cursor(db_connection) as cursor:
cursor.execute(_READ_SESSION_USER)
row = cursor.fetchone()
return row[0] if row else None
def _set_session_user(db_connection: Any, user_id: Optional[int]) -> None:
with _cursor(db_connection) as cursor:
cursor.execute(_SET_SESSION_USER, user_id)
[docs]
@contextmanager
def session_user(db_connection: Any, user_id: Optional[int]) -> Iterator[None]:
"""Run the block as user_id for row-level security, then restore the previous user."""
previous = current_session_user(db_connection)
_set_session_user(db_connection, user_id)
try:
yield
finally:
_set_session_user(db_connection, previous)
[docs]
def session_service(db_connection: Any):
"""Run the block as the ACE service user, then restore the previous user."""
return session_user(db_connection, settings.ACE_SERVICE_USER_ID)
# ---------------------------------------------------------------------------- script copies
[docs]
class LegacySqlQuery:
"""
Compatibility class for scripts that defined their own SqlQuery helper.
New code should use SqlQuery and session_user.
"""
[docs]
def __init__(self, db_connection: Any, user_id: Optional[int] = None, as_admin: Optional[bool] = None):
self.db_connection = db_connection
self.user_id = user_id
self.as_admin = bool(as_admin)
self.use_current_context = not self.as_admin and user_id is None
[docs]
def set_context(self):
cursor = self.db_connection.cursor
if not self.use_current_context:
cursor.execute(_SET_SESSION_USER, settings.ACE_SERVICE_USER_ID if self.as_admin else self.user_id)
return cursor
[docs]
def query(self, sql_query: str, *params: Any):
cursor = self.set_context()
cursor.execute(sql_query, *params)
return cursor
[docs]
def query_and_fetch(self, sql_query: str, *params: Any):
return self.query(sql_query, *params).fetchall()
[docs]
def query_and_fetch_one(self, sql_query: str, *params: Any):
return self.query(sql_query, *params).fetchone()
[docs]
def query_and_fetch_list_of(self, sql_query: str, target_type: Type[T], *params: Any) -> List[T]:
"""Return all rows, built with the columns as keyword arguments."""
cursor = self.query(sql_query, *params)
if issubclass(target_type, (str, int, float, bool, complex, bytes)):
return [target_type(row[0]) for row in cursor.fetchall()]
return self.rows_to_type(cursor, target_type)
[docs]
def query_and_fetch_one_as(self, sql_query: str, target_type: Type[T], *params: Any) -> T:
"""Return the first row; raise IndexError when there is none."""
return self.query_and_fetch_list_of(sql_query, target_type, *params)[0]
[docs]
def query_and_fetch_as(self, sql_query: str, target_type: Type[T], *params: Any) -> List[T]:
"""Return all rows, setting the columns the type has as attributes."""
cursor = self.query(sql_query, *params)
results = cursor.fetchall()
if issubclass(target_type, (str, int, float, bool, complex, bytes)):
return [target_type(row[0]) for row in results]
return [self._row_to_type(target_type, cursor.description, row) for row in results]
[docs]
@staticmethod
def rows_to_type(cursor, target_type: Type[T]) -> List[T]:
columns = [c[0] for c in cursor.description]
return [target_type(**dict(zip(columns, row))) for row in cursor.fetchall()]
@staticmethod
def _row_to_type(target_type, cursor_description, row):
instance = target_type()
for index, column in enumerate(cursor_description):
if hasattr(instance, column[0]):
setattr(instance, column[0], row[index])
return instance
__all__ = [
"SqlQuery",
"SqlMappingError",
"session_user",
"session_service",
"current_session_user",
"LegacySqlQuery",
"SCALAR_TYPES",
]