diff --git a/.env.example b/.env.example index 38d4acc6..e110d9b9 100644 --- a/.env.example +++ b/.env.example @@ -20,7 +20,10 @@ RABBITMQ_HOST=localhost NEO4J_HOST=localhost NEO4J_USERNAME=neo4j NEO4J_PASSWORD=discogsography +# POSTGRES_HOST may include a port (e.g. pgbouncer:6432) to target a pooler; +# otherwise POSTGRES_PORT (default 5432) is used. POSTGRES_HOST=localhost +# POSTGRES_PORT=5432 POSTGRES_USERNAME=discogsography POSTGRES_PASSWORD=discogsography POSTGRES_DATABASE=discogsography diff --git a/api/admin_setup.py b/api/admin_setup.py index eef15bd9..38b8ecb8 100644 --- a/api/admin_setup.py +++ b/api/admin_setup.py @@ -14,15 +14,14 @@ import psycopg from api.auth import _hash_password -from common.config import get_secret +from common.config import get_secret, parse_postgres_host_port def _build_conninfo() -> str: """Build PostgreSQL connection string from environment variables.""" - host = getenv("POSTGRES_HOST", "localhost") - port = getenv("POSTGRES_PORT", "5432") - if ":" in host: - host, port = host.rsplit(":", 1) + # POSTGRES_HOST may include an optional :port suffix (e.g. a pooler) + default_port = int(getenv("POSTGRES_PORT", "5432") or "5432") + host, port = parse_postgres_host_port(getenv("POSTGRES_HOST", "localhost"), default_port) user = get_secret("POSTGRES_USERNAME") or "postgres" password = get_secret("POSTGRES_PASSWORD") or "postgres" database = getenv("POSTGRES_DATABASE", "discogsography") diff --git a/api/api.py b/api/api.py index 40e04e00..8e215cb5 100644 --- a/api/api.py +++ b/api/api.py @@ -66,7 +66,7 @@ fetch_discogs_identity, request_oauth_token, ) -from common import AsyncPostgreSQLPool, AsyncResilientNeo4jDriver, HealthServer, neo4j_security_kwargs, setup_logging +from common import AsyncPostgreSQLPool, AsyncResilientNeo4jDriver, HealthServer, neo4j_security_kwargs, parse_postgres_host_port, setup_logging from common.config import ApiConfig from common.query_debug import execute_sql @@ -213,16 +213,12 @@ async def lifespan(_app: FastAPI) -> AsyncGenerator[None]: # pragma: no cover health_srv.start_background() logger.info("🏥 Health server started", port=API_HEALTH_PORT) - # Parse postgres address (format: host:port) - if ":" in _config.postgres_host: - host, port_str = _config.postgres_host.rsplit(":", 1) - else: - host = _config.postgres_host - port_str = "5432" + # Parse postgres address (POSTGRES_HOST may embed a port, e.g. a pooler) + host, port = parse_postgres_host_port(_config.postgres_host) _pool = AsyncPostgreSQLPool( connection_params={ "host": host, - "port": int(port_str), + "port": port, "dbname": _config.postgres_database, "user": _config.postgres_username, "password": _config.postgres_password, diff --git a/api/setup.py b/api/setup.py index c7fd985b..67cd7c3d 100644 --- a/api/setup.py +++ b/api/setup.py @@ -13,7 +13,7 @@ from psycopg.rows import dict_row from api.auth import decrypt_oauth_token, encrypt_oauth_token, get_oauth_encryption_key -from common.config import get_secret +from common.config import get_secret, parse_postgres_host_port def _build_conninfo() -> str: @@ -37,12 +37,9 @@ def _build_conninfo() -> str: print(f"❌ Missing required environment variables: {', '.join(missing)}", file=sys.stderr) sys.exit(1) - # POSTGRES_HOST may include an optional :port suffix - if ":" in address: - host, port = address.rsplit(":", 1) - else: - host = address - port = os.getenv("POSTGRES_PORT", "5432") + # POSTGRES_HOST may include an optional :port suffix (e.g. a pooler) + default_port = int(os.getenv("POSTGRES_PORT", "5432") or "5432") + host, port = parse_postgres_host_port(address, default_port) return f"host={host} port={port} user={username} password={password} dbname={database}" diff --git a/brainztableinator/brainztableinator.py b/brainztableinator/brainztableinator.py index c6954a37..f873fd29 100644 --- a/brainztableinator/brainztableinator.py +++ b/brainztableinator/brainztableinator.py @@ -19,6 +19,7 @@ AsyncResilientRabbitMQ, BrainztableinatorConfig, HealthServer, + parse_postgres_host_port, setup_logging, ) from orjson import loads @@ -946,13 +947,8 @@ async def main() -> None: logger.error("❌ Configuration error", error=str(e)) return - # Parse host and port from address - if ":" in config.postgres_host: - host, port_str = config.postgres_host.split(":", 1) - port = int(port_str) - else: - host = config.postgres_host - port = 5432 + # Parse host and port from address (POSTGRES_HOST may embed a port, e.g. a pooler) + host, port = parse_postgres_host_port(config.postgres_host) # Set connection parameters connection_params = { diff --git a/common/__init__.py b/common/__init__.py index 0b881ca1..e6bca8de 100644 --- a/common/__init__.py +++ b/common/__init__.py @@ -21,6 +21,7 @@ TableinatorConfig, get_config, neo4j_security_kwargs, + parse_postgres_host_port, setup_logging, ) from common.data_normalizer import ( @@ -128,6 +129,7 @@ "log_sql_query", "neo4j_security_kwargs", "normalize_record", + "parse_postgres_host_port", "process_message_with_retry", "resilient_connection", "setup_logging", diff --git a/common/config.py b/common/config.py index 1078c06a..a0e78c39 100644 --- a/common/config.py +++ b/common/config.py @@ -66,10 +66,70 @@ def _build_neo4j_uri() -> str: return f"bolt://{host}:7687" +def _coerce_port(value: str | None, default_port: int) -> int: + """Parse a port string to int, falling back to default_port on anything invalid.""" + try: + return int(str(value).strip()) + except (TypeError, ValueError): + return default_port + + +def parse_postgres_host_port(value: str | None, default_port: int = 5432) -> tuple[str, int]: + """Split a POSTGRES_HOST value into a (host, port) pair. + + POSTGRES_HOST may carry an embedded port (e.g. a PgBouncer pooler configured + as ``"pgbouncer:6432"``). When a port is embedded it always wins; otherwise + ``default_port`` (normally POSTGRES_PORT, falling back to 5432) is used. The + two ports are never concatenated. + + Accepted forms: + + - ``"host"`` -> ``(host, default_port)`` + - ``"host:6432"`` -> ``(host, 6432)`` (embedded port wins) + - ``"[::1]"`` -> ``("::1", default_port)`` + - ``"[::1]:6432"`` -> ``("::1", 6432)`` (IPv6 in brackets) + - ``"::1"`` -> ``("::1", default_port)`` (bare IPv6 literal, no port) + - ``""`` / ``None`` -> ``("localhost", default_port)`` + """ + raw = (value or "").strip() + if not raw: + return "localhost", default_port + + # IPv6 in brackets: "[host]" or "[host]:port" + if raw.startswith("["): + end = raw.find("]") + if end != -1: + host = raw[1:end] + rest = raw[end + 1 :] + if rest.startswith(":") and rest[1:]: + return host, _coerce_port(rest[1:], default_port) + return host, default_port + # Malformed bracket — fall through and treat the whole value as a host. + + # Bare IPv6 literal (more than one colon, no brackets) — no port to extract. + if raw.count(":") > 1: + return raw, default_port + + if ":" in raw: + host, _, port_str = raw.partition(":") + return (host or "localhost"), _coerce_port(port_str, default_port) + + return raw, default_port + + def _build_postgres_connstr() -> str: - """Build PostgreSQL host:port connection string from plain hostname environment variable.""" - host = getenv("POSTGRES_HOST", "localhost") - return f"{host}:5432" + """Build a canonical ``host:port`` connection string for PostgreSQL. + + Reads POSTGRES_HOST (which may embed a port, e.g. ``"pgbouncer:6432"``) and + POSTGRES_PORT (default 5432). An embedded port in POSTGRES_HOST takes + precedence over POSTGRES_PORT; the two are never concatenated. IPv6 hosts are + bracketed so the result round-trips through ``parse_postgres_host_port``. + """ + default_port = _coerce_port(getenv("POSTGRES_PORT", "5432"), 5432) + host, port = parse_postgres_host_port(getenv("POSTGRES_HOST", "localhost"), default_port) + if ":" in host: # IPv6 literal — bracket it for safe round-tripping. + return f"[{host}]:{port}" + return f"{host}:{port}" def _build_redis_url() -> str: diff --git a/dashboard/dashboard.py b/dashboard/dashboard.py index 33b0e91e..d9bab80a 100755 --- a/dashboard/dashboard.py +++ b/dashboard/dashboard.py @@ -27,6 +27,7 @@ AsyncResilientRabbitMQ, get_config, neo4j_security_kwargs, + parse_postgres_host_port, setup_logging, ) from dashboard.admin_proxy import configure as configure_admin_proxy, router as admin_router @@ -166,13 +167,8 @@ async def startup(self) -> None: logger.info("🔗 Connected to Neo4j with resilient driver") # Initialize resilient PostgreSQL connection - # Parse host and port from address - if ":" in self.config.postgres_host: - host, port_str = self.config.postgres_host.split(":", 1) - port = int(port_str) - else: - host = self.config.postgres_host - port = 5432 + # Parse host and port from address (POSTGRES_HOST may embed a port, e.g. a pooler) + host, port = parse_postgres_host_port(self.config.postgres_host) self.postgres_conn = AsyncResilientPostgreSQL( connection_params={ diff --git a/docs/configuration.md b/docs/configuration.md index 829063a8..5511f3a5 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -225,15 +225,21 @@ Bolt traffic crosses an untrusted network (separate host/VM, overlay network, cl ### PostgreSQL Configuration -| Variable | Description | Default | Required | -| ------------------- | ------------------- | ---------------- | -------- | -| `POSTGRES_HOST` | PostgreSQL hostname | `localhost` | Yes | -| `POSTGRES_USERNAME` | PostgreSQL username | (none) | Yes | -| `POSTGRES_PASSWORD` | PostgreSQL password | (none) | Yes | -| `POSTGRES_DATABASE` | Database name | `discogsography` | Yes | +| Variable | Description | Default | Required | +| ------------------- | -------------------------------------------- | ---------------- | -------- | +| `POSTGRES_HOST` | PostgreSQL hostname, optionally `host:port` | `localhost` | Yes | +| `POSTGRES_PORT` | PostgreSQL port (used when not in host) | `5432` | No | +| `POSTGRES_USERNAME` | PostgreSQL username | (none) | Yes | +| `POSTGRES_PASSWORD` | PostgreSQL password | (none) | Yes | +| `POSTGRES_DATABASE` | Database name | `discogsography` | Yes | **Used By**: Tableinator, Dashboard, API, Insights, Brainztableinator, Schema-Init +> **Note**: `POSTGRES_HOST` may include a port, e.g. `pgbouncer:6432` (useful when +> connecting through a connection pooler). An embedded port always takes precedence +> over `POSTGRES_PORT`; if no port is present, `POSTGRES_PORT` (default `5432`) is used. +> IPv6 hosts may be bracketed, e.g. `[::1]:6432`. + **Connection Details**: - Protocol: PostgreSQL wire protocol @@ -262,6 +268,12 @@ POSTGRES_HOST="db.example.com" POSTGRES_USERNAME="app_user" POSTGRES_PASSWORD="secure-password" POSTGRES_DATABASE="discogsography_prod" + +# Through a connection pooler (port embedded in host) +POSTGRES_HOST="pgbouncer:6432" +POSTGRES_USERNAME="app_user" +POSTGRES_PASSWORD="secure-password" +POSTGRES_DATABASE="discogsography_prod" ``` **Performance Tuning**: diff --git a/insights/insights.py b/insights/insights.py index 15424e87..b4be4cf4 100644 --- a/insights/insights.py +++ b/insights/insights.py @@ -22,7 +22,7 @@ import structlog import uvicorn -from common import AsyncPostgreSQLPool, HealthServer, setup_logging +from common import AsyncPostgreSQLPool, HealthServer, parse_postgres_host_port, setup_logging from common.config import InsightsConfig from insights.cache import InsightsCache from insights.computations import run_all_computations @@ -111,16 +111,12 @@ async def lifespan(_app: FastAPI) -> AsyncGenerator[None]: health_srv.start_background() logger.info("🏥 Health server started", port=INSIGHTS_HEALTH_PORT) - # Initialize PostgreSQL - if ":" in _config.postgres_host: - host, port_str = _config.postgres_host.rsplit(":", 1) - else: - host = _config.postgres_host - port_str = "5432" + # Initialize PostgreSQL (POSTGRES_HOST may embed a port, e.g. a pooler) + host, port = parse_postgres_host_port(_config.postgres_host) _pool = AsyncPostgreSQLPool( connection_params={ "host": host, - "port": int(port_str), + "port": port, "dbname": _config.postgres_database, "user": _config.postgres_username, "password": _config.postgres_password, diff --git a/schema-init/schema_init.py b/schema-init/schema_init.py index e6b45f82..b0a2b2f0 100755 --- a/schema-init/schema_init.py +++ b/schema-init/schema_init.py @@ -24,6 +24,7 @@ AsyncPostgreSQLPool, AsyncResilientNeo4jDriver, neo4j_security_kwargs, + parse_postgres_host_port, setup_logging, ) from common.config import get_secret @@ -47,13 +48,9 @@ def _postgres_connection_params() -> dict[str, Any]: - """Parse POSTGRES_HOST into psycopg connection params.""" - if ":" in POSTGRES_HOST: - host, port_str = POSTGRES_HOST.split(":", 1) - port = int(port_str) - else: - host = POSTGRES_HOST - port = 5432 + """Parse POSTGRES_HOST (which may embed a port, e.g. a pooler) into psycopg connection params.""" + default_port = int(os.environ.get("POSTGRES_PORT", "5432") or "5432") + host, port = parse_postgres_host_port(POSTGRES_HOST, default_port) return { "host": host, "port": port, diff --git a/tableinator/tableinator.py b/tableinator/tableinator.py index c9123080..1f4420ea 100644 --- a/tableinator/tableinator.py +++ b/tableinator/tableinator.py @@ -20,6 +20,7 @@ HealthServer, TableinatorConfig, normalize_record, + parse_postgres_host_port, setup_logging, ) from orjson import loads @@ -874,13 +875,8 @@ async def main() -> None: logger.error("❌ Configuration error", error=str(e)) return - # Parse host and port from address - if ":" in config.postgres_host: - host, port_str = config.postgres_host.split(":", 1) - port = int(port_str) - else: - host = config.postgres_host - port = 5432 + # Parse host and port from address (POSTGRES_HOST may embed a port, e.g. a pooler) + host, port = parse_postgres_host_port(config.postgres_host) # Set connection parameters connection_params = { diff --git a/tests/common/test_config.py b/tests/common/test_config.py index 86cde187..41f00797 100644 --- a/tests/common/test_config.py +++ b/tests/common/test_config.py @@ -14,7 +14,7 @@ neo4j_security_kwargs, setup_logging, ) -from common.config import get_secret +from common.config import _build_postgres_connstr, get_secret, parse_postgres_host_port class TestExtractorConfig: @@ -1225,3 +1225,107 @@ def test_non_true_values_stay_disabled(self, monkeypatch: pytest.MonkeyPatch) -> for value in ["1", "yes", "on", ""]: monkeypatch.setenv("NEO4J_TLS_ENABLED", value) assert neo4j_security_kwargs() == {} + + +class TestParsePostgresHostPort: + """Test parse_postgres_host_port — robust POSTGRES_HOST host/port splitting.""" + + def test_host_with_embedded_port(self) -> None: + """A pooler address like 'pgbouncer:6432' yields the embedded port.""" + assert parse_postgres_host_port("pgbouncer:6432") == ("pgbouncer", 6432) + + def test_embedded_port_wins_over_default(self) -> None: + """An embedded port takes precedence over the supplied default port.""" + assert parse_postgres_host_port("pgbouncer:6432", 5432) == ("pgbouncer", 6432) + + def test_host_without_port_uses_default(self) -> None: + """A bare host falls back to the default port (POSTGRES_PORT).""" + assert parse_postgres_host_port("postgres", 5432) == ("postgres", 5432) + + def test_host_without_port_uses_builtin_default(self) -> None: + """With no default supplied, a bare host falls back to 5432.""" + assert parse_postgres_host_port("postgres") == ("postgres", 5432) + + def test_host_with_explicit_5432(self) -> None: + """'postgres:5432' with no separate port still yields 5432.""" + assert parse_postgres_host_port("postgres:5432") == ("postgres", 5432) + + def test_empty_value_falls_back_to_localhost(self) -> None: + """An empty value yields localhost on the default port.""" + assert parse_postgres_host_port("", 5432) == ("localhost", 5432) + + def test_none_value_falls_back_to_localhost(self) -> None: + """A None value yields localhost on the default port.""" + assert parse_postgres_host_port(None) == ("localhost", 5432) + + def test_whitespace_is_stripped(self) -> None: + """Surrounding whitespace is stripped before parsing.""" + assert parse_postgres_host_port(" pgbouncer:6432 ") == ("pgbouncer", 6432) + + def test_malformed_port_falls_back_to_default(self) -> None: + """A non-numeric port falls back to the default rather than raising.""" + assert parse_postgres_host_port("pgbouncer:notaport", 5432) == ("pgbouncer", 5432) + + def test_trailing_colon_uses_default_port(self) -> None: + """'host:' (empty port) falls back to the default port.""" + assert parse_postgres_host_port("postgres:", 5432) == ("postgres", 5432) + + def test_ipv6_bracketed_with_port(self) -> None: + """A bracketed IPv6 host with a port is parsed correctly.""" + assert parse_postgres_host_port("[::1]:6432") == ("::1", 6432) + + def test_ipv6_bracketed_without_port(self) -> None: + """A bracketed IPv6 host without a port uses the default port.""" + assert parse_postgres_host_port("[2001:db8::1]", 5432) == ("2001:db8::1", 5432) + + def test_ipv6_bare_literal_no_port(self) -> None: + """A bare (unbracketed) IPv6 literal is treated as host-only.""" + assert parse_postgres_host_port("::1", 5432) == ("::1", 5432) + + def test_never_concatenates_ports(self) -> None: + """Regression: the embedded port is never concatenated with the default.""" + _host, port = parse_postgres_host_port("pgbouncer:6432", 5432) + assert port == 6432 + assert ":" not in str(port) + + +class TestBuildPostgresConnstr: + """Test _build_postgres_connstr — canonical host:port assembly from env.""" + + def test_embedded_port_in_host(self, monkeypatch: pytest.MonkeyPatch) -> None: + """POSTGRES_HOST carrying a port produces a clean host:port with no double port.""" + monkeypatch.setenv("POSTGRES_HOST", "pgbouncer:6432") + monkeypatch.delenv("POSTGRES_PORT", raising=False) + assert _build_postgres_connstr() == "pgbouncer:6432" + + def test_host_plus_postgres_port(self, monkeypatch: pytest.MonkeyPatch) -> None: + """A bare host with POSTGRES_PORT set uses that port.""" + monkeypatch.setenv("POSTGRES_HOST", "postgres") + monkeypatch.setenv("POSTGRES_PORT", "5432") + assert _build_postgres_connstr() == "postgres:5432" + + def test_host_without_port_defaults_5432(self, monkeypatch: pytest.MonkeyPatch) -> None: + """A bare host with no POSTGRES_PORT defaults to 5432.""" + monkeypatch.setenv("POSTGRES_HOST", "postgres") + monkeypatch.delenv("POSTGRES_PORT", raising=False) + assert _build_postgres_connstr() == "postgres:5432" + + def test_embedded_port_wins_over_postgres_port(self, monkeypatch: pytest.MonkeyPatch) -> None: + """An embedded port in POSTGRES_HOST beats POSTGRES_PORT — never concatenated.""" + monkeypatch.setenv("POSTGRES_HOST", "pgbouncer:6432") + monkeypatch.setenv("POSTGRES_PORT", "5432") + assert _build_postgres_connstr() == "pgbouncer:6432" + + def test_ipv6_host_is_bracketed(self, monkeypatch: pytest.MonkeyPatch) -> None: + """An IPv6 host round-trips as a bracketed host:port string.""" + monkeypatch.setenv("POSTGRES_HOST", "[::1]:6432") + monkeypatch.delenv("POSTGRES_PORT", raising=False) + result = _build_postgres_connstr() + assert result == "[::1]:6432" + assert parse_postgres_host_port(result) == ("::1", 6432) + + def test_default_host(self, monkeypatch: pytest.MonkeyPatch) -> None: + """No POSTGRES_HOST defaults to localhost:5432.""" + monkeypatch.delenv("POSTGRES_HOST", raising=False) + monkeypatch.delenv("POSTGRES_PORT", raising=False) + assert _build_postgres_connstr() == "localhost:5432"