svc-infra 0.1.589__py3-none-any.whl → 0.1.706__py3-none-any.whl

svc_infra/db/sql/utils.py

@@ -3,7 +3,7 @@ from __future__ import annotations
 import asyncio
 import os
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Optional, Sequence, Set, Union
+from typing import TYPE_CHECKING, Any, Optional, Sequence, Set, Union, cast
 
 from alembic.config import Config
 from dotenv import load_dotenv
@@ -17,19 +17,19 @@ if TYPE_CHECKING:
     from sqlalchemy.engine import Engine as SyncEngine
     from sqlalchemy.ext.asyncio import AsyncEngine as AsyncEngineType
 else:
-    SyncEngine = Any  # type: ignore
-    AsyncEngineType = Any  # type: ignore
+    SyncEngine = Any  # type: ignore[assignment]
+    AsyncEngineType = Any  # type: ignore[assignment]
 
 try:
     # Runtime import (may be missing if async extras aren’t installed)
-    from sqlalchemy.ext.asyncio import create_async_engine as _create_async_engine  # type: ignore
+    from sqlalchemy.ext.asyncio import create_async_engine as _create_async_engine
 except Exception:  # pragma: no cover - optional dep
-    _create_async_engine = None  # type: ignore
+    _create_async_engine = None  # type: ignore[assignment]
 
 try:
-    from sqlalchemy import create_engine as _create_engine  # type: ignore
+    from sqlalchemy import create_engine as _create_engine
 except Exception:  # pragma: no cover - optional env
-    _create_engine = None  # type: ignore
+    _create_engine = None  # type: ignore[assignment]
 
 
 def prepare_process_env(
@@ -83,7 +83,9 @@ def _compose_url_from_parts() -> Optional[str]:
       DB_PARAMS (raw query string like 'sslmode=require&connect_timeout=5')
     """
     dialect = os.getenv("DB_DIALECT", "").strip() or "postgresql"
-    driver = os.getenv("DB_DRIVER", "").strip()  # e.g. asyncpg, psycopg, pymysql, aiosqlite
+    driver = os.getenv(
+        "DB_DRIVER", ""
+    ).strip()  # e.g. asyncpg, psycopg, pymysql, aiosqlite
     host = os.getenv("DB_HOST", "").strip() or None
     port = os.getenv("DB_PORT", "").strip() or None
     db = os.getenv("DB_NAME", "").strip() or None
@@ -124,18 +126,62 @@ def _compose_url_from_parts() -> Optional[str]:
 # ---------- Environment helpers ----------
 
 
+def _normalize_database_url(url: str) -> str:
+    """
+    Normalize database URL for SQLAlchemy compatibility.
+
+    Handles:
+      - postgres:// → postgresql:// (Heroku/Railway legacy format)
+      - Strips whitespace
+
+    Args:
+        url: Raw database URL string
+
+    Returns:
+        Normalized URL suitable for SQLAlchemy
+    """
+    url = url.strip()
+    # Heroku and Railway historically use 'postgres://' which SQLAlchemy doesn't accept
+    if url.startswith("postgres://"):
+        url = "postgresql://" + url[len("postgres://") :]
+    return url
+
+
 def get_database_url_from_env(
     required: bool = True,
     env_vars: Sequence[str] = DEFAULT_DB_ENV_VARS,
+    normalize: bool = True,
 ) -> Optional[str]:
     """
     Resolve the database connection string, with support for:
-      - Primary env vars (in order): DEFAULT_DB_ENV_VARS (e.g. SQL_URL, PRIVATE_SQL_URL, DB_URL).
+      - Primary env vars (in order): DEFAULT_DB_ENV_VARS
+        (SQL_URL, DB_URL, DATABASE_URL, DATABASE_URL_PRIVATE, PRIVATE_SQL_URL)
       - Secret file pointers: <NAME>_FILE (reads file contents).
       - Well-known locations: SQL_URL_FILE, /run/secrets/database_url.
       - Composed from parts: DB_* (host, port, name, user, password, params).
-    When a value is found, it is also written back into os.environ["SQL_URL"] for downstream code.
+
+    When a value is found, it is also written back into os.environ["SQL_URL"]
+    for downstream code.
+
+    Args:
+        required: If True, raise RuntimeError when no URL is found
+        env_vars: Sequence of environment variable names to check
+        normalize: If True, convert postgres:// to postgresql:// (default: True)
+
+    Returns:
+        Database URL string, or None if not found and not required
+
+    Raises:
+        RuntimeError: If required=True and no URL is found
     """
+
+    def _finalize(url: str) -> str:
+        """Normalize and cache the URL."""
+        if normalize:
+            url = _normalize_database_url(url)
+        os.environ["SQL_URL"] = url
+        return url
+
     # Load .env without clobbering existing process env
     load_dotenv(override=False)
 
@@ -150,10 +196,8 @@ def get_database_url_from_env(
             if os.path.isabs(s) and Path(s).exists():
                 file_val = _read_secret_from_file(s)
                 if file_val:
-                    os.environ["SQL_URL"] = file_val
-                    return file_val
-            os.environ["SQL_URL"] = s
-            return s
+                    return _finalize(file_val)
+            return _finalize(s)
 
         # Companion NAME_FILE secret path (e.g., SQL_URL_FILE)
         file_key = f"{key}_FILE"
@@ -161,32 +205,28 @@ def get_database_url_from_env(
         if file_path:
             file_val = _read_secret_from_file(file_path)
             if file_val:
-                os.environ["SQL_URL"] = file_val
-                return file_val
+                return _finalize(file_val)
 
     # 2) Conventional secret envs
     file_path = os.getenv("SQL_URL_FILE")
     if file_path:
         file_val = _read_secret_from_file(file_path)
         if file_val:
-            os.environ["SQL_URL"] = file_val
-            return file_val
+            return _finalize(file_val)
 
     # 3) Docker/K8s default secret mount
     file_val = _read_secret_from_file("/run/secrets/database_url")
     if file_val:
-        os.environ["SQL_URL"] = file_val
-        return file_val
+        return _finalize(file_val)
 
     # 4) Compose from parts (DB_DIALECT/DB_DRIVER/DB_HOST/.../DB_PARAMS)
     composed = _compose_url_from_parts()
     if composed:
-        os.environ["SQL_URL"] = composed
-        return composed
+        return _finalize(composed)
 
     if required:
         raise RuntimeError(
-            "Database URL not set. Set SQL_URL (or PRIVATE_SQL_URL / DB_URL), "
+            "Database URL not set. Set SQL_URL, DATABASE_URL, or DATABASE_URL_PRIVATE, "
             "or provide DB_* parts (DB_HOST, DB_NAME, etc.), or a *_FILE secret."
         )
     return None
@@ -196,10 +236,17 @@ def _ensure_timeout_default(u: URL) -> URL:
     """
     Ensure a conservative connection timeout is present for libpq-based drivers.
     For psycopg/psycopg2, 'connect_timeout' is honored via the query string.
+    For asyncpg, timeout is set via connect_args (not query string).
     """
     backend = (u.get_backend_name() or "").lower()
     if backend not in ("postgresql", "postgres"):
         return u
+
+    # asyncpg doesn't support connect_timeout in query string - use connect_args instead
+    dn = (u.drivername or "").lower()
+    if "+asyncpg" in dn:
+        return u
+
     if "connect_timeout" in u.query:
         return u
     # Default 10s unless overridden
@@ -329,17 +376,22 @@ def _ensure_ssl_default(u: URL) -> URL:
     driver = (u.drivername or "").lower()
 
     # If any SSL hint already present, do nothing
-    if any(k in u.query for k in ("sslmode", "ssl", "sslrootcert", "sslcert", "sslkey")):
+    if any(
+        k in u.query for k in ("sslmode", "ssl", "sslrootcert", "sslcert", "sslkey")
+    ):
         return u
 
     # Allow env override; support both common spellings
-    mode_env = os.getenv("DB_SSLMODE_DEFAULT") or os.getenv("PGSSLMODE") or os.getenv("PGSSL_MODE")
+    mode_env = (
+        os.getenv("DB_SSLMODE_DEFAULT")
+        or os.getenv("PGSSLMODE")
+        or os.getenv("PGSSL_MODE")
+    )
     mode = (mode_env or "").strip()
 
     if "+asyncpg" in driver:
-        # asyncpg: use 'ssl=true' (SQLAlchemy forwards to asyncpg)
-        if "ssl" not in u.query:
-            return u.set(query={**u.query, "ssl": "true"})
+        # asyncpg: SSL is handled in connect_args in build_engine(), not in URL query
+        # Do not add ssl parameter to URL query for asyncpg
         return u
     else:
         # libpq-based drivers: use sslmode (default 'require' for hosted PG)
@@ -351,7 +403,9 @@ def _ensure_ssl_default_async(u: URL) -> URL:
     backend = (u.get_backend_name() or "").lower()
     if backend in ("postgresql", "postgres"):
         # asyncpg prefers 'ssl=true' via SQLAlchemy param; if already present, keep it
-        if any(k in u.query for k in ("ssl", "sslmode", "sslrootcert", "sslcert", "sslkey")):
+        if any(
+            k in u.query for k in ("ssl", "sslmode", "sslrootcert", "sslcert", "sslkey")
+        ):
             return u
         return u.set(query={**u.query, "ssl": "true"})
     return u
@@ -366,7 +420,9 @@ def _certifi_ca() -> str | None:
         return None
 
 
-def build_engine(url: URL | str, echo: bool = False) -> Union[SyncEngine, AsyncEngineType]:
+def build_engine(
+    url: URL | str, echo: bool = False
+) -> Union[SyncEngine, AsyncEngineType]:
     u = make_url(url) if isinstance(url, str) else url
 
     # Keep your existing PG helpers
@@ -382,10 +438,20 @@ def build_engine(url: URL | str, echo: bool = False) -> Union[SyncEngine, AsyncE
                 "Async driver URL provided but SQLAlchemy async extras are not available."
             )
 
-        # asyncpg: honor connection timeout
+        # asyncpg: honor connection timeout only (NOT connect_timeout)
         if "+asyncpg" in (u.drivername or ""):
             connect_args["timeout"] = int(os.getenv("DB_CONNECT_TIMEOUT", "10"))
 
+            # asyncpg doesn't accept sslmode or ssl=true in query params
+            # Remove these and set ssl='require' in connect_args
+            if "ssl" in u.query or "sslmode" in u.query:
+                new_query = {
+                    k: v for k, v in u.query.items() if k not in ("ssl", "sslmode")
+                }
+                u = u.set(query=new_query)
+            # Set ssl in connect_args - 'require' is safest for hosted databases
+            connect_args["ssl"] = "require"
+
         # NEW: aiomysql SSL default
         if "+aiomysql" in (u.drivername or "") and not any(
             k in u.query for k in ("ssl", "ssl_ca", "sslmode")
@@ -395,7 +461,11 @@ def build_engine(url: URL | str, echo: bool = False) -> Union[SyncEngine, AsyncE
                 import ssl
 
                 ca = _certifi_ca()
-                ctx = ssl.create_default_context(cafile=ca) if ca else ssl.create_default_context()
+                ctx = (
+                    ssl.create_default_context(cafile=ca)
+                    if ca
+                    else ssl.create_default_context()
+                )
                 # if your host uses a public CA, verification works;
                 # if not, you can relax verification (not recommended):
                 #   ctx.check_hostname = False
@@ -404,15 +474,17 @@ def build_engine(url: URL | str, echo: bool = False) -> Union[SyncEngine, AsyncE
             except Exception:
                 connect_args["ssl"] = True  # minimal hint to enable TLS
 
-        kwargs: dict[str, Any] = {"echo": echo, "pool_pre_ping": True}
+        async_engine_kwargs: dict[str, Any] = {"echo": echo, "pool_pre_ping": True}
         if connect_args:
-            kwargs["connect_args"] = connect_args
-        return _create_async_engine(u, **kwargs)
+            async_engine_kwargs["connect_args"] = connect_args
+        return _create_async_engine(u, **async_engine_kwargs)
 
     # ----------------- SYNC -----------------
     u = _coerce_sync_driver(u)
     if _create_engine is None:
-        raise RuntimeError("SQLAlchemy create_engine is not available in this environment.")
+        raise RuntimeError(
+            "SQLAlchemy create_engine is not available in this environment."
+        )
 
     dn = (u.drivername or "").lower()
 
@@ -435,10 +507,10 @@ def build_engine(url: URL | str, echo: bool = False) -> Union[SyncEngine, AsyncE
         # Optional: if your provider requires it, you can also add:
         # connect_args.setdefault("client_flag", 0)
 
-    kwargs: dict[str, Any] = {"echo": echo, "pool_pre_ping": True}
+    sync_engine_kwargs: dict[str, Any] = {"echo": echo, "pool_pre_ping": True}
     if connect_args:
-        kwargs["connect_args"] = connect_args
-    return _create_engine(u, **kwargs)
+        sync_engine_kwargs["connect_args"] = connect_args
+    return _create_engine(u, **sync_engine_kwargs)
 
 
 # ---------- Identifier quoting helpers ----------
@@ -485,7 +557,7 @@ async def _pg_create_database_async(url: URL) -> None:
             )
             if not exists:
                 quoted = _pg_quote_ident(target_db)
-                await conn.execution_options(isolation_level="AUTOCOMMIT").execute(
+                await conn.execution_options(isolation_level="AUTOCOMMIT").execute(  # type: ignore[attr-defined]
                     text(f'CREATE DATABASE "{quoted}"')
                 )
     except DBAPIError as e:
@@ -545,7 +617,9 @@ async def _mysql_create_database_async(url: URL) -> None:
     engine: AsyncEngineType = build_engine(base_url)  # type: ignore[assignment]
     async with engine.begin() as conn:
         exists = await conn.scalar(
-            text("SELECT SCHEMA_NAME FROM INFORMATION_SCHEMA.SCHEMATA WHERE SCHEMA_NAME = :name"),
+            text(
+                "SELECT SCHEMA_NAME FROM INFORMATION_SCHEMA.SCHEMATA WHERE SCHEMA_NAME = :name"
+            ),
             {"name": target_db},
         )
         if not exists:
@@ -562,7 +636,9 @@ def _mysql_create_database_sync(url: URL) -> None:
     engine: SyncEngine = build_engine(base_url)  # type: ignore[assignment]
     with engine.begin() as conn:
         exists = conn.scalar(
-            text("SELECT SCHEMA_NAME FROM INFORMATION_SCHEMA.SCHEMATA WHERE SCHEMA_NAME = :name"),
+            text(
+                "SELECT SCHEMA_NAME FROM INFORMATION_SCHEMA.SCHEMATA WHERE SCHEMA_NAME = :name"
+            ),
             {"name": target_db},
         )
         if not exists:
@@ -789,17 +865,19 @@ def ensure_database_exists(url: URL | str) -> None:
     try:
         eng = build_engine(u)
         if is_async_url(u):
+            async_eng = cast(AsyncEngineType, eng)
 
             async def _ping_and_dispose():
-                async with eng.begin() as conn:  # type: ignore[call-arg]
+                async with async_eng.begin() as conn:
                     await conn.execute(text("SELECT 1"))
-                await eng.dispose()  # type: ignore[attr-defined]
+                await async_eng.dispose()
 
             asyncio.run(_ping_and_dispose())
         else:
-            with eng.begin() as conn:  # type: ignore[call-arg]
+            sync_eng = cast(SyncEngine, eng)
+            with sync_eng.begin() as conn:
                 conn.execute(text("SELECT 1"))
-            eng.dispose()  # type: ignore[attr-defined]
+            sync_eng.dispose()
     except OperationalError as exc:  # pragma: no cover (depends on env)
         raise RuntimeError(f"Failed to connect to database: {exc}") from exc
 
@@ -811,7 +889,10 @@ def repair_alembic_state_if_needed(cfg: Config) -> None:
         return
 
     # Gather local revision ids from versions/
-    script_location = Path(cfg.get_main_option("script_location"))
+    script_location_str = cfg.get_main_option("script_location")
+    if not script_location_str:
+        return
+    script_location = Path(script_location_str)
     versions_dir = script_location / "versions"
     local_ids: Set[str] = set()
     if versions_dir.exists():
@@ -831,7 +912,7 @@ def repair_alembic_state_if_needed(cfg: Config) -> None:
     if is_async_url(url_obj):
 
         async def _run() -> None:
-            eng = build_engine(url_obj)  # AsyncEngine
+            eng = cast(AsyncEngineType, build_engine(url_obj))
             try:
                 async with eng.begin() as conn:
                     # Do sync-y inspector / SQL via run_sync
@@ -846,7 +927,9 @@ def repair_alembic_state_if_needed(cfg: Config) -> None:
                         )
                         missing = any((ver not in local_ids) for (ver,) in rows)
                         if missing:
-                            sync_conn.execute(text("DROP TABLE IF EXISTS alembic_version"))
+                            sync_conn.execute(
+                                text("DROP TABLE IF EXISTS alembic_version")
+                            )
 
                     await conn.run_sync(_check_and_maybe_drop)
             finally:
@@ -854,13 +937,17 @@ def repair_alembic_state_if_needed(cfg: Config) -> None:
 
         asyncio.run(_run())
     else:
-        eng = build_engine(url_obj)  # sync Engine
+        eng = cast(SyncEngine, build_engine(url_obj))
         try:
             with eng.begin() as c:
                 insp = inspect(c)
                 if not insp.has_table("alembic_version"):
                     return
-                rows = list(c.execute(text("SELECT version_num FROM alembic_version")).fetchall())
+                rows = list(
+                    c.execute(
+                        text("SELECT version_num FROM alembic_version")
+                    ).fetchall()
+                )
                 missing = any((ver not in local_ids) for (ver,) in rows)
                 if missing:
                     c.execute(text("DROP TABLE IF EXISTS alembic_version"))

svc_infra/db/sql/versioning.py

@@ -0,0 +1,14 @@
+from __future__ import annotations
+
+from sqlalchemy import Integer
+from sqlalchemy.orm import Mapped, mapped_column
+
+
+class Versioned:
+    """Mixin for optimistic locking with integer version.
+
+    - Initialize version=1 on insert (via default=1)
+    - Bump version in app code before commit to detect mismatches.
+    """
+
+    version: Mapped[int] = mapped_column(Integer, nullable=False, default=1)