tegwick
dda93db0ac
Assistant: codex Assistant-Model: gpt-5.6-sol Assistant-Session: 01a053ff-1d6f-7fe2-ac1c-a6eb40a42a0c
52 lines
1.6 KiB
Python
52 lines
1.6 KiB
Python
from __future__ import annotations
|
|
|
|
import re
|
|
from typing import Protocol
|
|
|
|
|
|
class MigrationConnection(Protocol):
|
|
def exec_driver_sql(self, statement: str) -> object: ...
|
|
|
|
def commit(self) -> None: ...
|
|
|
|
|
|
def migration_role_statement(role: str | None) -> str | None:
|
|
"""Return a safely quoted SET ROLE statement for an admitted owner role."""
|
|
if not role:
|
|
return None
|
|
if not re.fullmatch(r"[a-z_][a-z0-9_]{0,62}", role):
|
|
raise ValueError("HUB_CORE_MIGRATION_ROLE is not a safe PostgreSQL role name")
|
|
return f'SET ROLE "{role}"'
|
|
|
|
|
|
def migration_schema_statement(schema: str | None) -> str | None:
|
|
"""Return an explicit, safely quoted migration search path."""
|
|
if not schema:
|
|
return None
|
|
if not re.fullmatch(r"[a-z_][a-z0-9_]{0,62}", schema):
|
|
raise ValueError("HUB_CORE_MIGRATION_SCHEMA is not a safe PostgreSQL schema name")
|
|
return f'SET search_path TO "{schema}", public'
|
|
|
|
|
|
def configure_migration_session(
|
|
connection: MigrationConnection,
|
|
*,
|
|
role: str | None,
|
|
schema: str | None,
|
|
) -> None:
|
|
"""Apply and commit session settings before Alembic opens its transaction."""
|
|
statements = tuple(
|
|
statement
|
|
for statement in (
|
|
migration_role_statement(role),
|
|
migration_schema_statement(schema),
|
|
)
|
|
if statement is not None
|
|
)
|
|
for statement in statements:
|
|
connection.exec_driver_sql(statement)
|
|
if statements:
|
|
# SQLAlchemy 2 autobegins on exec_driver_sql(). Without this commit,
|
|
# Alembic joins the setup transaction and its DDL is rolled back when
|
|
# the connection closes.
|
|
connection.commit()
|