From fece73853316eae70bef97d896bb14afa561ad2f Mon Sep 17 00:00:00 2001 From: Zeding Zhang Date: Tue, 3 Nov 2026 01:36:47 +0800 Subject: [PATCH] Python: Prefer injected MongoDB clients over connection settings --- python/packages/mongodb/README.md | 5 +- .../agent_framework_mongodb/_vector_store.py | 20 +++++--- .../mongodb/tests/mongodb/test_settings.py | 51 +++++++++++++++++-- 3 files changed, 65 insertions(+), 11 deletions(-) diff --git a/python/packages/mongodb/README.md b/python/packages/mongodb/README.md index e74b1bfa080..34df9992aff 100644 --- a/python/packages/mongodb/README.md +++ b/python/packages/mongodb/README.md @@ -22,8 +22,9 @@ Constructors resolve explicit arguments first, then an explicitly selected `.env` file, then process environment variables. URIs accept `str` or AF `SecretString` and are unwrapped only for PyMongo. -You can instead inject an `AsyncMongoClient` and pass `database_name`. An -injected client bypasses URI, app-name, and `.env` loading and always remains +You can instead inject an `AsyncMongoClient`. The database name is required and +resolves from an explicit `database_name`, then `MONGODB_DATABASE_NAME`. An injected +client bypasses URI, app-name, and `.env` loading and always remains caller-owned. Connector-created clients are closed by `close()` or async context exit. Collections created by a store borrow its already-resolved client. diff --git a/python/packages/mongodb/agent_framework_mongodb/_vector_store.py b/python/packages/mongodb/agent_framework_mongodb/_vector_store.py index edc739392f2..3ea074fe01e 100644 --- a/python/packages/mongodb/agent_framework_mongodb/_vector_store.py +++ b/python/packages/mongodb/agent_framework_mongodb/_vector_store.py @@ -60,6 +60,12 @@ class MongoDBSettings(TypedDict, total=False): app_name: str | None +class _MongoDBDatabaseSettings(TypedDict, total=False): + """Resolve the database selector independently of client connection settings.""" + + database_name: str | None + + _BSON_MAX_DOCUMENT_SIZE = 16 * 1024 * 1024 _BSON_INT64_MIN = -(2**63) _BSON_INT64_MAX = 2**63 - 1 @@ -104,8 +110,8 @@ def _create_client( env_file_encoding: str | None, ) -> tuple[MongoClient, MongoDatabase, str, bool]: if async_client is not None: - if any(value is not None for value in (uri, app_name, env_file_path, env_file_encoding)): - raise ValueError("async_client cannot be combined with uri, app_name, env_file_path, or env_file_encoding.") + if database_name is None: + database_name = load_settings(_MongoDBDatabaseSettings, env_prefix="MONGODB_").get("database_name") resolved_database_name = _validate_non_empty_string(database_name, "database_name") return ( async_client, @@ -606,9 +612,10 @@ def __init__( Args: record_type: Registered application model, or dict with an explicit definition. uri: MongoDB URI override; otherwise loaded from ``MONGODB_URI``. - database_name: Database name, or ``MONGODB_DATABASE_NAME`` for a connector-created client. + database_name: Database name, or ``MONGODB_DATABASE_NAME`` when omitted. app_name: Optional driver app name, or ``MONGODB_APP_NAME``. - async_client: Existing PyMongo async client. It bypasses settings and remains caller-owned. + async_client: Existing PyMongo async client. When supplied, URI, app-name, + and .env options are ignored. The client remains caller-owned. definition: Explicit dictionary schema. collection_name: Collection name, overriding the model definition. embedding_generator: Default local embedding client. @@ -1140,9 +1147,10 @@ def __init__( Args: uri: MongoDB URI override; otherwise loaded from ``MONGODB_URI``. - database_name: Database name, or ``MONGODB_DATABASE_NAME`` for a connector-created client. + database_name: Database name, or ``MONGODB_DATABASE_NAME`` when omitted. app_name: Optional driver app name, or ``MONGODB_APP_NAME``. - async_client: Existing PyMongo async client. It bypasses settings and remains caller-owned. + async_client: Existing PyMongo async client. When supplied, URI, app-name, + and .env options are ignored. The client remains caller-owned. embedding_generator: Default embedding client inherited by collections. env_file_path: Optional .env file selected explicitly. env_file_encoding: Encoding for the selected .env file. diff --git a/python/packages/mongodb/tests/mongodb/test_settings.py b/python/packages/mongodb/tests/mongodb/test_settings.py index 931b6dfec08..ca98f1c66c6 100644 --- a/python/packages/mongodb/tests/mongodb/test_settings.py +++ b/python/packages/mongodb/tests/mongodb/test_settings.py @@ -126,14 +126,59 @@ async def test_injected_client_bypasses_settings_and_remains_borrowed(connector_ {"env_file_encoding": "utf-16"}, ], ) -def test_injected_client_rejects_connection_settings(connector_factory, settings): +async def test_injected_client_ignores_connection_settings(connector_factory, settings): client = _client_mock() with ( patch("agent_framework_mongodb._vector_store.load_settings") as loader, - pytest.raises(ValueError, match="async_client cannot be combined"), + patch("agent_framework_mongodb._vector_store.AsyncMongoClient") as factory, ): - connector_factory(async_client=client, database_name="vectors", **settings) + async with connector_factory(async_client=client, database_name="vectors", **settings) as connector: + assert connector.async_client is client + assert connector.database_name == "vectors" loader.assert_not_called() + factory.assert_not_called() + client.get_database.assert_called_once_with("vectors") + client.close.assert_not_awaited() + + +@pytest.mark.parametrize( + ("database_name", "expected_database_name"), + [(None, "environment_db"), ("explicit_db", "explicit_db")], +) +async def test_injected_client_database_name_precedence( + connector_factory, monkeypatch, tmp_path, database_name, expected_database_name +): + monkeypatch.setenv("MONGODB_URI", "not-a-mongodb-uri") + monkeypatch.setenv("MONGODB_APP_NAME", "") + monkeypatch.setenv("MONGODB_DATABASE_NAME", "environment_db") + env_file = tmp_path / "mongodb.env" + env_file.write_text("MONGODB_DATABASE_NAME=file_db\n") + client = _client_mock() + with patch("agent_framework_mongodb._vector_store.AsyncMongoClient") as factory: + async with connector_factory( + async_client=client, + database_name=database_name, + uri="not-a-mongodb-uri", + app_name="", + env_file_path=str(env_file), + env_file_encoding="invalid-encoding", + ) as connector: + assert connector.async_client is client + assert connector.database_name == expected_database_name + factory.assert_not_called() + client.get_database.assert_called_once_with(expected_database_name) + client.close.assert_not_awaited() + + +def test_injected_client_requires_database_name(connector_factory): + client = _client_mock() + with ( + patch("agent_framework_mongodb._vector_store.AsyncMongoClient") as factory, + pytest.raises(ValueError, match="database_name"), + ): + connector_factory(async_client=client) + factory.assert_not_called() + client.get_database.assert_not_called() def test_store_collections_reuse_resolved_client_without_settings(definition):