From ec48cd6d0db895d822f79a1f890e0c16ed240752 Mon Sep 17 00:00:00 2001 From: Zac Pustejovsky Date: Fri, 8 May 2026 17:56:39 -0400 Subject: [PATCH 1/2] postgres connnection string support --- .../src/zeroshot_commons/config_utils.py | 18 ++++++++++ .../zeroshot_commons/postgres_connection.py | 15 +++++++- .../commons/tests/unit/test_config_utils.py | 36 +++++++++++++++++++ .../commons/tests/unit/test_postgres_utils.py | 31 ++++++++++++++++ uv.lock | 25 +++++++------ 5 files changed, 114 insertions(+), 11 deletions(-) diff --git a/packages/commons/src/zeroshot_commons/config_utils.py b/packages/commons/src/zeroshot_commons/config_utils.py index b2be687..e23b1a0 100644 --- a/packages/commons/src/zeroshot_commons/config_utils.py +++ b/packages/commons/src/zeroshot_commons/config_utils.py @@ -96,6 +96,24 @@ def load_config( config = _load_raw_config(config_path) merged_config = deep_merge(config, parse_env_variables()) + database_url = os.environ.get("DATABASE_URL") + if database_url: + from .postgres_connection import PostgresConnectionConfig + + pg = PostgresConnectionConfig.from_url(database_url) + merged_config = deep_merge( + merged_config, + { + "postgres": { + "host": pg.host, + "port": pg.port, + "username": pg.username, + "password": pg.password, + "database": pg.database, + } + }, + ) + if not config_key: return merged_config diff --git a/packages/commons/src/zeroshot_commons/postgres_connection.py b/packages/commons/src/zeroshot_commons/postgres_connection.py index 6d2efc9..2637624 100644 --- a/packages/commons/src/zeroshot_commons/postgres_connection.py +++ b/packages/commons/src/zeroshot_commons/postgres_connection.py @@ -3,7 +3,7 @@ from collections.abc import Mapping from dataclasses import dataclass from typing import Any -from urllib.parse import quote_plus +from urllib.parse import quote_plus, unquote, urlparse from .application_config import ApplicationConfig from .config_utils import load_config @@ -25,6 +25,19 @@ class PostgresConnectionConfig: POSTGRES_CONFIG_KEY = "postgres" + @classmethod + def from_url(cls, url: str) -> PostgresConnectionConfig: + parsed = urlparse(url) + if parsed.scheme not in ("postgresql", "postgres"): + raise ValueError(f"Unsupported scheme '{parsed.scheme}', expected 'postgresql' or 'postgres'") + return cls( + host=parsed.hostname or "localhost", + port=parsed.port or 5432, + username=unquote(parsed.username or ""), + password=unquote(parsed.password or ""), + database=(parsed.path or "/").lstrip("/"), + ) + @classmethod def from_mapping(cls, data: Mapping[str, Any]) -> PostgresConnectionConfig: return cls( diff --git a/packages/commons/tests/unit/test_config_utils.py b/packages/commons/tests/unit/test_config_utils.py index 0f7e917..1a2d30e 100644 --- a/packages/commons/tests/unit/test_config_utils.py +++ b/packages/commons/tests/unit/test_config_utils.py @@ -74,6 +74,42 @@ def test_load_config_reads_json_and_applies_env_overrides(tmp_path: Path) -> Non assert config["service"] == {"host": "base", "port": 8080} +def test_load_config_applies_database_url_override(tmp_path: Path) -> None: + package_root = tmp_path / "package" + main_dir = package_root / "src" + assets_dir = package_root / "assets" + main_dir.mkdir(parents=True) + assets_dir.mkdir() + (assets_dir / "config.json").write_text( + json.dumps( + { + "postgres": { + "host": "old-host", + "port": 5432, + "username": "old", + "password": "old", + "database": "old_db", + } + } + ), + encoding="utf-8", + ) + + config = run_with_env( + lambda: load_config( + str(main_dir), + config_file_path="assets/config.json", + ), + [("DATABASE_URL", "postgresql://newuser:newpass@newhost:6543/new_db")], + ) + + assert config["postgres"]["host"] == "newhost" + assert config["postgres"]["port"] == 6543 + assert config["postgres"]["username"] == "newuser" + assert config["postgres"]["password"] == "newpass" + assert config["postgres"]["database"] == "new_db" + + def test_load_config_reads_sub_config_and_application_config(tmp_path: Path) -> None: package_root = tmp_path / "package" main_dir = package_root / "src" diff --git a/packages/commons/tests/unit/test_postgres_utils.py b/packages/commons/tests/unit/test_postgres_utils.py index 51386a9..a296414 100644 --- a/packages/commons/tests/unit/test_postgres_utils.py +++ b/packages/commons/tests/unit/test_postgres_utils.py @@ -68,6 +68,37 @@ async def recovery() -> str: assert recovered is True +def test_from_url_parses_standard_postgres_url() -> None: + config = PostgresConnectionConfig.from_url("postgresql://user:pass@localhost:5432/mydb") + assert config.host == "localhost" + assert config.port == 5432 + assert config.username == "user" + assert config.password == "pass" + assert config.database == "mydb" + + +def test_from_url_handles_postgres_scheme() -> None: + config = PostgresConnectionConfig.from_url("postgres://u:p@host:1234/db") + assert config.host == "host" + assert config.port == 1234 + assert config.username == "u" + assert config.password == "p" + assert config.database == "db" + + +def test_from_url_decodes_url_encoded_credentials() -> None: + config = PostgresConnectionConfig.from_url( + "postgresql://user%40domain:p%40ss%3Aword@host:5432/db" + ) + assert config.username == "user@domain" + assert config.password == "p@ss:word" + + +def test_from_url_rejects_unsupported_scheme() -> None: + with pytest.raises(ValueError, match="Unsupported scheme"): + PostgresConnectionConfig.from_url("mysql://user:pass@host:3306/db") + + def test_postgres_connection_config_generates_expected_urls() -> None: config = PostgresConnectionConfig.from_mapping( { diff --git a/uv.lock b/uv.lock index decf986..d7bd476 100644 --- a/uv.lock +++ b/uv.lock @@ -101,7 +101,7 @@ wheels = [ [[package]] name = "buildkit-python-workspace" -version = "0.1.0" +version = "0.1.5" source = { virtual = "." } [package.dev-dependencies] @@ -1458,6 +1458,11 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/13/2d/26b8b30067d94339afee62c3edc9b803a6eb9332f521ba77d8aaab5de873/testcontainers-4.14.2-py3-none-any.whl", hash = "sha256:0d0522c3cd8f8d9627cda41f7a6b51b639fa57bdc492923c045117933c668d68", size = 125712, upload-time = "2026-03-18T05:19:15.29Z" }, ] +[package.optional-dependencies] +redis = [ + { name = "redis" }, +] + [[package]] name = "tqdm" version = "4.67.3" @@ -1636,7 +1641,7 @@ wheels = [ [[package]] name = "zeroshot-agent-experiments" -version = "0.1.0" +version = "0.1.5" source = { editable = "packages/agent-experiments" } dependencies = [ { name = "fpdf2" }, @@ -1659,7 +1664,7 @@ requires-dist = [ [[package]] name = "zeroshot-agentic-workflows" -version = "0.1.0" +version = "0.1.5" source = { editable = "packages/agentic-workflows" } dependencies = [ { name = "openai-agents" }, @@ -1676,7 +1681,7 @@ requires-dist = [ [[package]] name = "zeroshot-commons" -version = "0.1.0" +version = "0.1.5" source = { editable = "packages/commons" } dependencies = [ { name = "pyyaml" }, @@ -1691,7 +1696,7 @@ requires-dist = [ [[package]] name = "zeroshot-commons-injectors" -version = "0.1.0" +version = "0.1.5" source = { editable = "packages/commons-injectors" } dependencies = [ { name = "asyncpg" }, @@ -1712,22 +1717,22 @@ requires-dist = [ [[package]] name = "zeroshot-commons-testing" -version = "0.1.0" +version = "0.1.5" source = { editable = "packages/commons-testing" } dependencies = [ - { name = "testcontainers" }, + { name = "testcontainers", extra = ["redis"] }, { name = "zeroshot-commons" }, ] [package.metadata] requires-dist = [ - { name = "testcontainers", specifier = ">=4.9.0" }, + { name = "testcontainers", extras = ["postgres", "redis"], specifier = ">=4.9.0" }, { name = "zeroshot-commons", editable = "packages/commons" }, ] [[package]] name = "zeroshot-openai-utils" -version = "0.1.0" +version = "0.1.5" source = { editable = "packages/openai-utils" } dependencies = [ { name = "dependency-injector" }, @@ -1744,7 +1749,7 @@ requires-dist = [ [[package]] name = "zeroshot-sql-decorators" -version = "0.1.0" +version = "0.1.5" source = { editable = "packages/sql-decorators" } dependencies = [ { name = "asyncpg" }, From 9df547add2bb895003a34ac7d799e8e7522be643 Mon Sep 17 00:00:00 2001 From: Zac Pustejovsky Date: Fri, 8 May 2026 18:01:43 -0400 Subject: [PATCH 2/2] moar --- packages/commons/src/zeroshot_commons/postgres_connection.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/packages/commons/src/zeroshot_commons/postgres_connection.py b/packages/commons/src/zeroshot_commons/postgres_connection.py index 2637624..d2d0111 100644 --- a/packages/commons/src/zeroshot_commons/postgres_connection.py +++ b/packages/commons/src/zeroshot_commons/postgres_connection.py @@ -29,7 +29,9 @@ class PostgresConnectionConfig: def from_url(cls, url: str) -> PostgresConnectionConfig: parsed = urlparse(url) if parsed.scheme not in ("postgresql", "postgres"): - raise ValueError(f"Unsupported scheme '{parsed.scheme}', expected 'postgresql' or 'postgres'") + raise ValueError( + f"Unsupported scheme '{parsed.scheme}', expected 'postgresql' or 'postgres'" + ) return cls( host=parsed.hostname or "localhost", port=parsed.port or 5432,