diff --git a/python/packages/postgres/README.md b/python/packages/postgres/README.md index dd3efa6c13..5a3cef265b 100644 --- a/python/packages/postgres/README.md +++ b/python/packages/postgres/README.md @@ -39,7 +39,7 @@ The `schema` argument defaults to `public`. A connector-created pool is closed by `close()` or an async context manager. Alternatively, inject an open Psycopg `AsyncConnection` or `AsyncConnectionPool` using `client`; it remains caller-owned and bypasses settings loading. -Injected clients cannot be combined with connection-string or `.env` options. +Connection-string and `.env` options are ignored when a client is injected. Collections created by a store borrow its pool, so keep the store open while using them. diff --git a/python/packages/postgres/agent_framework_postgres/_vector_store.py b/python/packages/postgres/agent_framework_postgres/_vector_store.py index 7d8f797f1b..e4857cb5db 100644 --- a/python/packages/postgres/agent_framework_postgres/_vector_store.py +++ b/python/packages/postgres/agent_framework_postgres/_vector_store.py @@ -122,8 +122,6 @@ def _create_client( env_file_encoding: str | None, ) -> _Client: if client is not None: - if connection_string is not None or env_file_path is not None or env_file_encoding is not None: - raise ValueError("client cannot be combined with connection_string, env_file_path, or env_file_encoding.") return _Client(None, client) settings = load_settings( PostgresSettings, @@ -430,13 +428,14 @@ def __init__( record_type: Decorated/registered model type, or ``dict``. connection_string: Psycopg conninfo or URI for an owned lazy pool. Falls back to ``POSTGRES_CONNECTION_STRING`` in the selected .env file or process environment. - client: An open caller-owned async connection or pool, instead of conninfo. + client: An open caller-owned async connection or pool. When provided, + connection_string and .env options are ignored. schema: Existing PostgreSQL schema containing the table. definition: Explicit field definition for dictionary records. collection_name: Table name, overriding the model's collection name. embedding_generator: Default local embedding client. - env_file_path: Optional .env file to read; cannot be combined with an injected client. - env_file_encoding: Encoding of the .env file; cannot be combined with an injected client. + env_file_path: Optional .env file to read when creating a pool. + env_file_encoding: Encoding of the .env file used when creating a pool. """ super().__init__( record_type, @@ -877,11 +876,12 @@ def __init__( Args: connection_string: Psycopg conninfo or URI for an owned lazy pool. Falls back to ``POSTGRES_CONNECTION_STRING`` in the selected .env file or process environment. - client: An open caller-owned async connection or pool. + client: An open caller-owned async connection or pool. When provided, + connection_string and .env options are ignored. schema: Existing schema used by all collection clients. embedding_generator: Default embedding generator for collections. - env_file_path: Optional .env file to read; cannot be combined with an injected client. - env_file_encoding: Encoding of the .env file; cannot be combined with an injected client. + env_file_path: Optional .env file to read when creating a pool. + env_file_encoding: Encoding of the .env file used when creating a pool. """ super().__init__(embedding_generator=embedding_generator, managed_client=client is None) _prepare_identifier(schema) diff --git a/python/packages/postgres/tests/test_postgres.py b/python/packages/postgres/tests/test_postgres.py index e876cde013..0aa9671155 100644 --- a/python/packages/postgres/tests/test_postgres.py +++ b/python/packages/postgres/tests/test_postgres.py @@ -205,7 +205,7 @@ async def test_keys_are_strict(definition_factory, record_factory, key_type, val [ {}, {"connection_string": ""}, - {"connection_string": "x", "client": MagicMock(spec=AsyncConnection)}, + {"connection_string": "x", "client": object()}, {"client": object()}, ], ) diff --git a/python/packages/postgres/tests/test_postgres_settings.py b/python/packages/postgres/tests/test_postgres_settings.py index d8626f63a6..ecf6ad94cc 100644 --- a/python/packages/postgres/tests/test_postgres_settings.py +++ b/python/packages/postgres/tests/test_postgres_settings.py @@ -144,15 +144,36 @@ async def test_injected_client_bypasses_loader_and_remains_borrowed(constructor, {"env_file_path": ""}, {"env_file_encoding": "utf-16"}, {"env_file_encoding": ""}, + { + "connection_string": "invalid conninfo", + "env_file_path": "missing.env", + "env_file_encoding": "invalid-encoding", + }, ], ) -def test_injected_client_rejects_conflicting_settings_without_reading(constructor, kwargs): - client = MagicMock(spec=AsyncConnection) +@pytest.mark.parametrize("client_type", [AsyncConnection, AsyncConnectionPool]) +async def test_injected_client_ignores_conflicting_settings_without_reading( + constructor, monkeypatch, kwargs, client_type +): + monkeypatch.setenv("POSTGRES_CONNECTION_STRING", "not valid conninfo; must not be read") + client = MagicMock(spec=client_type) + client.close = AsyncMock() with ( patch.object(module, "load_settings", side_effect=AssertionError("Must not read settings")), - pytest.raises(ValueError, match="cannot be combined"), + patch.object( + AsyncConnectionPool, "__init__", side_effect=AssertionError("Must not create a pool") + ) as create_pool, ): - constructor(client=client, **kwargs) + instance = constructor(client=client, schema="caller_schema", **kwargs) + assert instance._client.client is client + assert not instance.managed_client + assert instance.schema == "caller_schema" + await instance.close() + create_pool.assert_not_called() + client.close.assert_not_awaited() + if client_type is AsyncConnection: + client.commit.assert_not_called() + client.rollback.assert_not_called() async def test_store_children_reuse_resolved_client_without_loading_again(monkeypatch, definition_factory):