Compare commits

...
Author SHA1 Message Date
Nicolò Boschi 65d0f0f6e1 test(db): pin current production URL shapes as regression guard 2026-04-23 14:20:26 +02:00
Nicolò Boschi 3469b8fac7 fix(db): accept asyncpg-style URLs for external PostgreSQL
Fixes #1216. External PostgreSQL deployments (Cloud SQL, RDS, etc.)
configured with a SQLAlchemy-style URL like
`postgresql+asyncpg://user:pass@host/db?ssl=require` failed in two
places:

1. Five sync `create_engine(database_url)` call sites in migrations.py
   — psycopg2 doesn't understand the asyncpg dialect, and it expects
   `sslmode=require` rather than `ssl=require`.
2. `asyncpg.create_pool(self.db_url)` in memory_engine.py — asyncpg
   doesn't parse the `postgresql+asyncpg://` scheme directly.

Adds a single `to_libpq_url()` helper (urllib.parse-based, idempotent,
safe on passwords containing `+`) and applies it at:

- All five `create_engine()` sites in migrations.py (including the
  run_migrations advisory-lock connection)
- `asyncpg.create_pool()` in memory_engine.py
- The ad-hoc scheme rewrite in alembic/env.py (replaced by the helper)

Existing configs (`pg0`, plain `postgresql://`, `sslmode=require`,
`postgresql+psycopg2://`) are returned byte-identical — no behaviour
change for current users.
2026-04-23 13:53:46 +02:00
5 changed files with 213 additions and 11 deletions
@@ -12,6 +12,7 @@ from dotenv import load_dotenv
from sqlalchemy import engine_from_config, pool
# Import your models here
from hindsight_api.db_url import to_libpq_url
from hindsight_api.models import Base
@@ -65,11 +66,11 @@ def get_database_url() -> str:
"Set HINDSIGHT_API_DATABASE_URL environment variable or pass database_url to run_migrations()."
)
# For migrations, use psycopg2 (sync driver) to avoid pgbouncer prepared statement issues
if database_url.startswith("postgresql+asyncpg://"):
database_url = database_url.replace("postgresql+asyncpg://", "postgresql://", 1)
elif database_url.startswith("postgres+asyncpg://"):
database_url = database_url.replace("postgres+asyncpg://", "postgresql://", 1)
# For migrations, use the sync psycopg2 driver (avoids pgbouncer prepared
# statement issues and is required since create_engine is the sync API).
# Also translates ?ssl=require (SQLAlchemy asyncpg style) to ?sslmode=require
# (libpq style) for external-PostgreSQL deployments.
database_url = to_libpq_url(database_url)
# Update config with processed URL for engine_from_config to use
config.set_main_option("sqlalchemy.url", database_url)
@@ -0,0 +1,58 @@
"""Database URL normalization.
Hindsight accepts SQLAlchemy-style URLs like ``postgresql+asyncpg://...?ssl=require``
for its async engine, but the same string cannot be handed directly to synchronous
SQLAlchemy (psycopg2) or to :func:`asyncpg.create_pool`, which both expect a
libpq-compatible URL (``postgresql://...?sslmode=require``).
:func:`to_libpq_url` performs that translation. It is idempotent and safe to
apply to URLs that are already libpq-compatible, to the ``pg0`` embedded-PG
marker, or to any non-PostgreSQL string (returned unchanged).
"""
from __future__ import annotations
from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit
_ASYNCPG_SCHEMES = ("postgresql+asyncpg", "postgres+asyncpg")
_POSTGRES_SCHEMES = ("postgresql", "postgres") + _ASYNCPG_SCHEMES
def to_libpq_url(url: str) -> str:
"""Normalize a PostgreSQL URL for libpq-style consumers.
Accepts a SQLAlchemy URL (``postgresql+asyncpg://...``) or a plain libpq
URL and returns a form suitable for:
- :func:`sqlalchemy.create_engine` (sync / psycopg2)
- :func:`asyncpg.create_pool`
Transformations:
- ``postgresql+asyncpg`` / ``postgres+asyncpg`` / ``postgres`` → ``postgresql``
- Query param ``ssl=<mode>`` → ``sslmode=<mode>`` (SQLAlchemy's asyncpg
dialect uses ``ssl=``; libpq uses ``sslmode=``)
Any non-PostgreSQL input (e.g. the ``pg0`` embedded-PG marker, a sqlite
URL, an empty string) is returned unchanged. Already-normalized URLs are
returned unchanged.
"""
if not url or "://" not in url:
return url
parts = urlsplit(url)
if parts.scheme not in _POSTGRES_SCHEMES:
return url
new_scheme = "postgresql"
new_query_pairs = [
("sslmode", v) if k == "ssl" else (k, v)
for k, v in parse_qsl(parts.query, keep_blank_values=True)
]
new_query = urlencode(new_query_pairs)
if new_scheme == parts.scheme and new_query == parts.query:
return url
return urlunsplit((new_scheme, parts.netloc, parts.path, new_query, parts.fragment))
@@ -31,6 +31,7 @@ from ..config import (
DEFAULT_REFLECT_SOURCE_FACTS_MAX_TOKENS,
get_config,
)
from ..db_url import to_libpq_url
from ..metrics import get_metrics_collector
from ..tracing import create_operation_span
from ..utils import mask_network_location
@@ -1871,7 +1872,7 @@ class MemoryEngine(MemoryEngineInterface):
await conn.execute(f"SET statement_timeout = '{stmt_timeout_s}s'")
self._pool = await asyncpg.create_pool(
self.db_url,
to_libpq_url(self.db_url),
min_size=self._pool_min_size,
max_size=self._pool_max_size,
command_timeout=self._db_command_timeout,
@@ -27,6 +27,7 @@ from alembic.config import Config
from alembic.script.revision import ResolutionError
from sqlalchemy import Connection, create_engine, text
from .db_url import to_libpq_url
from .utils import mask_network_location
logger = logging.getLogger(__name__)
@@ -220,7 +221,7 @@ def run_migrations(
# ineffective when the app URL goes through a pooler. Configure
# HINDSIGHT_API_MIGRATION_DATABASE_URL to the direct PostgreSQL endpoint
# (e.g. hindsight-pg-rw) to restore correct locking behaviour.
migration_url = migration_database_url or database_url
migration_url = to_libpq_url(migration_database_url or database_url)
try:
# Determine script location
@@ -450,7 +451,7 @@ def check_migration_status(
return None, None
# Get current revision from database
engine = create_engine(database_url)
engine = create_engine(to_libpq_url(database_url))
with engine.connect() as connection:
context = MigrationContext.configure(connection)
current_rev = context.get_current_revision()
@@ -624,7 +625,7 @@ def ensure_embedding_dimension(
"""
schema_name = schema or "public"
engine = create_engine(database_url)
engine = create_engine(to_libpq_url(database_url))
with engine.connect() as conn:
# Check if memory_units table exists (proxy for schema being initialized)
table_exists = conn.execute(
@@ -673,7 +674,7 @@ def ensure_vector_extension(
"""
schema_name = schema or "public"
engine = create_engine(database_url)
engine = create_engine(to_libpq_url(database_url))
with engine.connect() as conn:
# Detect which vector extension should be used
target_ext = _detect_vector_extension(conn, vector_extension)
@@ -894,7 +895,7 @@ def ensure_text_search_extension(
"""
schema_name = schema or "public"
engine = create_engine(database_url)
engine = create_engine(to_libpq_url(database_url))
with engine.connect() as conn:
# Tables with search_vector columns to check
tables_to_check = [
+141
View File
@@ -0,0 +1,141 @@
"""Tests for ``hindsight_api.db_url.to_libpq_url``.
Covers backward compatibility (existing configs must pass through unchanged)
and the two transformations needed to support external PostgreSQL deployments
that use SQLAlchemy-style ``postgresql+asyncpg://...?ssl=require`` URLs:
1. strip the ``+asyncpg`` dialect suffix,
2. rename the ``ssl=`` query parameter to ``sslmode=``.
"""
from __future__ import annotations
import pytest
from hindsight_api.db_url import to_libpq_url
class TestPassthrough:
"""Inputs that must be returned unchanged — protects existing configs."""
@pytest.mark.parametrize(
"url",
[
"pg0",
"",
"postgresql://user:pass@host:5432/db",
"postgresql://user:pass@host:5432/db?sslmode=require",
"postgresql://user:pass@host/db?sslmode=verify-full&connect_timeout=10",
"sqlite:///./test.db",
"postgresql+psycopg2://user:pass@host/db",
],
)
def test_unchanged(self, url: str) -> None:
assert to_libpq_url(url) == url
class TestSchemeNormalization:
def test_asyncpg_scheme_stripped(self) -> None:
assert (
to_libpq_url("postgresql+asyncpg://user:pass@host:5432/db")
== "postgresql://user:pass@host:5432/db"
)
def test_postgres_asyncpg_scheme_normalized(self) -> None:
assert (
to_libpq_url("postgres+asyncpg://user:pass@host/db")
== "postgresql://user:pass@host/db"
)
def test_bare_postgres_scheme_normalized_to_postgresql(self) -> None:
assert to_libpq_url("postgres://user:pass@host/db") == "postgresql://user:pass@host/db"
class TestSslParamRename:
def test_ssl_require_to_sslmode_require(self) -> None:
assert (
to_libpq_url("postgresql+asyncpg://user:pass@host:5432/db?ssl=require")
== "postgresql://user:pass@host:5432/db?sslmode=require"
)
@pytest.mark.parametrize("mode", ["disable", "allow", "prefer", "require", "verify-ca", "verify-full"])
def test_all_ssl_modes_translated(self, mode: str) -> None:
result = to_libpq_url(f"postgresql+asyncpg://h/d?ssl={mode}")
assert result == f"postgresql://h/d?sslmode={mode}"
def test_ssl_rename_on_libpq_url(self) -> None:
"""Someone accidentally using SQLAlchemy-style ssl= on a libpq URL is also fixed."""
assert to_libpq_url("postgresql://h/d?ssl=require") == "postgresql://h/d?sslmode=require"
def test_ssl_param_preserved_among_other_params(self) -> None:
result = to_libpq_url(
"postgresql+asyncpg://h/d?ssl=require&application_name=hindsight&connect_timeout=10"
)
assert result.startswith("postgresql://h/d?")
# Query order should be preserved; ssl renamed, others untouched.
assert "sslmode=require" in result
assert "application_name=hindsight" in result
assert "connect_timeout=10" in result
assert "ssl=" not in result.split("?", 1)[1].replace("sslmode=", "")
def test_sslmode_not_double_renamed(self) -> None:
"""An already-correct sslmode= param must not be altered."""
assert (
to_libpq_url("postgresql+asyncpg://h/d?sslmode=require")
== "postgresql://h/d?sslmode=require"
)
class TestProductionConfigs:
"""Regression guard: current production URL shapes must pass through unchanged.
These are the exact shapes currently set for HINDSIGHT_API_DATABASE_URL,
HINDSIGHT_API_CONTROL_DATABASE_URL and HINDSIGHT_API_MIGRATION_DATABASE_URL
in production. The helper must be a pure no-op for them so this change is
truly backward-compatible.
"""
@pytest.mark.parametrize(
"url",
[
"postgresql://app:[email protected]:5432/appdb?sslmode=disable",
"postgresql://app:[email protected]:5432/appdb_control?sslmode=disable",
"postgresql://app:[email protected]:5432/appdb?sslmode=disable",
],
)
def test_prod_urls_object_identical(self, url: str) -> None:
# Not just equal — must be the exact same object (early-out path),
# guaranteeing no parse/reassembly and no subtle mutation.
assert to_libpq_url(url) is url
class TestEdgeCases:
def test_idempotent(self) -> None:
original = "postgresql+asyncpg://user:pass@host:5432/db?ssl=require"
once = to_libpq_url(original)
twice = to_libpq_url(once)
assert once == twice
def test_password_with_plus_is_preserved(self) -> None:
"""A naive str.replace('+asyncpg', ...) would corrupt passwords containing '+'.
urllib.parse operates on the parsed scheme only, so this stays safe.
"""
url = "postgresql+asyncpg://user:pa%2Bsswd@host/db?ssl=require"
result = to_libpq_url(url)
assert result == "postgresql://user:pa%2Bsswd@host/db?sslmode=require"
def test_password_literal_asyncpg_in_password(self) -> None:
"""Even a password that literally contains '+asyncpg' must survive."""
url = "postgresql+asyncpg://user:my%2Basyncpgpass@host/db"
result = to_libpq_url(url)
assert result == "postgresql://user:my%2Basyncpgpass@host/db"
def test_url_without_query_string(self) -> None:
assert (
to_libpq_url("postgresql+asyncpg://user:pass@host/db")
== "postgresql://user:pass@host/db"
)
def test_url_with_port_and_path_only(self) -> None:
assert to_libpq_url("postgresql+asyncpg://host:5432/db") == "postgresql://host:5432/db"