Python: Prefer injected DocumentDB clients over connection settings

python/packages/azure-documentdb/README.md

@@ -33,7 +33,13 @@ file, then process environment. Connection strings are held in Agent Framework
 Connector-created clients enforce TLS, disable retryable writes, and set an
 application name. They are closed by `aclose()` or an async context manager.
 An injected PyMongo `AsyncMongoClient`, `AsyncDatabase`, or `AsyncCollection`
-remains caller-owned and bypasses settings loading.
+remains caller-owned and takes precedence over connection strings, client
+options, and selected `.env` files. An account-level injected client resolves a
+missing `database_name` from `AZURE_DOCUMENTDB_DATABASE_NAME`. For a collection
+created from an injected client or database, `AZURE_DOCUMENTDB_COLLECTION_NAME`
+is used only when neither `collection_name` nor the model definition supplies a
+name. Connection settings and selected `.env` files are never read for these
+injected-object paths.
 
 The identity running `ensure_collection_exists()` needs permission to create
 collections and indexes. Each filtered data field must declare

python/packages/azure-documentdb/agent_framework_azure_documentdb/_vector_store.py

@@ -23,7 +23,12 @@
 )
 from agent_framework._telemetry import FeatureIndex, mark_feature_used
 from agent_framework._vector_filters import FilterExpression
-from agent_framework._vectors import EmbeddingClient, SearchType, Vector
+from agent_framework._vectors import (
+    _VECTOR_MODEL_REGISTRY,  # pyright: ignore[reportPrivateUsage]
+    EmbeddingClient,
+    SearchType,
+    Vector,
+)
 from agent_framework.exceptions import IntegrationException, IntegrationInvalidResponseException
 from bson import BSON
 from bson.errors import InvalidDocument
@@ -108,6 +113,35 @@ def _validate_collection_name(name: Any) -> str:
     return name
 
 
+class _AzureDocumentDBDatabaseSettings(TypedDict, total=False):
+    database_name: str | None
+
+
+class _AzureDocumentDBCollectionSettings(TypedDict, total=False):
+    collection_name: str | None
+
+
+def _resolve_resource_name(value: str | None, kind: str) -> str:
+    if value is None:
+        if kind == "database":
+            settings = load_settings(
+                _AzureDocumentDBDatabaseSettings,
+                env_prefix="AZURE_DOCUMENTDB_",
+                required_fields=["database_name"],
+            )
+            value = settings.get("database_name")
+        else:
+            settings = load_settings(
+                _AzureDocumentDBCollectionSettings,
+                env_prefix="AZURE_DOCUMENTDB_",
+                required_fields=["collection_name"],
+            )
+            value = settings.get("collection_name")
+    if kind == "database":
+        return _validate_database_name(value)
+    return _validate_collection_name(value)
+
+
 def _validate_storage_name(name: str, *, indexed: bool) -> str:
     encoded = name.encode("utf-8")
     if not name or "\0" in name or "." in name or name.startswith("$"):
@@ -163,29 +197,18 @@ def _create_client(
     env_file_encoding: str | None,
 ) -> _Client:
     if database is not None:
-        if any(
-            value is not None
-            for value in (
-                connection_string,
-                database_name,
-                client,
-                client_options,
-                env_file_path,
-                env_file_encoding,
-            )
-        ):
-            raise ValueError("database cannot be combined with a client or connection settings.")
+        if client is not None:
+            raise ValueError("Provide at most one of client or database.")
         if not isinstance(database, AsyncDatabase):
             raise TypeError("database must be a PyMongo AsyncDatabase.")
+        resolved_name = _validate_database_name(database.name)
+        if database_name is not None and database_name != resolved_name:
+            raise ValueError("database_name must match the injected database.")
         return _Client(database.client, database, owned=False)
     if client is not None:
-        if any(value is not None for value in (connection_string, client_options, env_file_path, env_file_encoding)):
-            raise ValueError("client cannot be combined with connection settings or client_options.")
         if not isinstance(client, AsyncMongoClient):
             raise TypeError("client must be a PyMongo AsyncMongoClient.")
-        if database_name is None:
-            raise ValueError("database_name is required with an injected client.")
-        resolved_database_name = _validate_database_name(database_name)
+        resolved_database_name = _resolve_resource_name(database_name, "database")
         return _Client(client, client[resolved_database_name], owned=False)
 
     settings = load_settings(
@@ -610,38 +633,41 @@ def __init__(
         Args:
             record_type: Registered model type, or ``dict`` with an explicit definition.
             connection_string: MongoDB connection string for a connector-owned client.
-            database_name: Database selected for a created or injected client.
-            client: Caller-owned PyMongo async client. Requires ``database_name``.
-            database: Caller-owned resolved PyMongo async database.
-            collection: Caller-owned resolved PyMongo async collection.
+                Ignored when a client, database, or collection is injected.
+            database_name: Database selected for a created or injected client. An injected account
+                client falls back to ``AZURE_DOCUMENTDB_DATABASE_NAME`` when this is omitted.
+            client: Caller-owned PyMongo async client taking precedence over connection settings.
+            database: Caller-owned resolved PyMongo async database taking precedence over connection settings.
+            collection: Caller-owned resolved PyMongo async collection taking precedence over connection settings.
             client_options: Created-client options limited to ``tls``, ``retryWrites``, and ``appname``.
+                Ignored when a client, database, or collection is injected.
             definition: Explicit dictionary schema.
-            collection_name: Collection name overriding the model definition.
+            collection_name: Collection name overriding the model definition. With an injected client
+                or database, ``AZURE_DOCUMENTDB_COLLECTION_NAME`` is the fallback when neither an
+                explicit value nor the model definition supplies a name.
             embedding_generator: Default local embedding generator.
-            env_file_path: Optional selected settings file.
-            env_file_encoding: Encoding of the selected settings file.
+            env_file_path: Optional selected settings file, ignored for injected objects.
+            env_file_encoding: Encoding of the selected settings file, ignored for injected objects.
         """
         if collection is not None:
             if not isinstance(collection, AsyncCollection):
                 raise TypeError("collection must be a PyMongo AsyncCollection.")
+            if client is not None or database is not None:
+                raise ValueError("Provide at most one of client, database, or collection.")
             if collection_name is not None and collection_name != collection.name:
                 raise ValueError("collection_name must match the injected collection.")
-            if any(
-                value is not None
-                for value in (
-                    connection_string,
-                    database_name,
-                    client,
-                    database,
-                    client_options,
-                    env_file_path,
-                    env_file_encoding,
-                )
-            ):
-                raise ValueError("collection cannot be combined with a client, database, or connection settings.")
             resolved_collection_name = collection.name
         else:
             resolved_collection_name = collection_name
+        if collection is not None or client is not None or database is not None:
+            registration = _VECTOR_MODEL_REGISTRY.get(record_type)
+            registered_definition = definition or (registration.definition if registration is not None else None)
+            configured_name = (
+                resolved_collection_name
+                if resolved_collection_name is not None
+                else getattr(registered_definition, "collection_name", None)
+            )
+            resolved_collection_name = _resolve_resource_name(configured_name, "collection")
         super().__init__(
             record_type,
             definition=definition,
@@ -1116,13 +1142,16 @@ def __init__(
 
         Args:
             connection_string: MongoDB connection string for a connector-owned client.
-            database_name: Database selected for a created or injected client.
-            client: Caller-owned PyMongo async client. Requires ``database_name``.
-            database: Caller-owned resolved PyMongo async database.
+                Ignored when a client or database is injected.
+            database_name: Database selected for a created or injected client. An injected account
+                client falls back to ``AZURE_DOCUMENTDB_DATABASE_NAME`` when this is omitted.
+            client: Caller-owned PyMongo async client taking precedence over connection settings.
+            database: Caller-owned resolved PyMongo async database taking precedence over connection settings.
             client_options: Created-client options limited to ``tls``, ``retryWrites``, and ``appname``.
+                Ignored when a client or database is injected.
             embedding_generator: Default embedding generator for child collections.
-            env_file_path: Optional selected settings file.
-            env_file_encoding: Encoding of the selected settings file.
+            env_file_path: Optional selected settings file, ignored for injected objects.
+            env_file_encoding: Encoding of the selected settings file, ignored for injected objects.
         """
         super().__init__(
             embedding_generator=embedding_generator,

python/packages/azure-documentdb/tests/azure_documentdb/test_settings.py

@@ -2,11 +2,12 @@
 
 from __future__ import annotations
 
-from typing import get_type_hints
+from typing import Any, get_type_hints
 from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
-from agent_framework import SecretString, load_settings
+from agent_framework import SecretString, VectorStoreCollectionDefinition, load_settings, register_vectorstoremodel
+from agent_framework.exceptions import SettingNotFoundError
 from pymongo import AsyncMongoClient
 from pymongo.asynchronous.database import AsyncDatabase
 
@@ -129,20 +130,157 @@ async def test_injected_database_and_collection_bypass_settings(definition_facto
 
 
 @pytest.mark.parametrize(
-    "kwargs",
-    [
-        {"connection_string": "mongodb://explicit"},
-        {"client_options": {"tls": True}},
-        {"env_file_path": "selected.env"},
-        {"env_file_encoding": "utf-16"},
-    ],
+    "target",
+    ["store-account", "store-database", "collection-account", "collection-database", "collection-collection"],
 )
-def test_injected_client_rejects_connection_settings(kwargs, pymongo_objects):
-    with (
-        patch.object(module, "load_settings", side_effect=AssertionError("settings must be bypassed")),
-        pytest.raises(ValueError, match="cannot be combined"),
-    ):
-        AzureDocumentDBStore(client=pymongo_objects[0], database_name="db", **kwargs)
+async def test_injected_objects_ignore_connection_settings(
+    target, definition_factory, pymongo_objects, monkeypatch, tmp_path
+):
+    client, database, collection = pymongo_objects
+    monkeypatch.setenv("AZURE_DOCUMENTDB_CONNECTION_STRING", "")
+    monkeypatch.setenv("AZURE_DOCUMENTDB_DATABASE_NAME", "x" * 64)
+    monkeypatch.setenv("AZURE_DOCUMENTDB_COLLECTION_NAME", "x" * 256)
+    kwargs: dict[str, Any] = {
+        "connection_string": "mongodb://ignored",
+        "client_options": {"tls": False},
+        "env_file_path": str(tmp_path / "missing.env"),
+        "env_file_encoding": "invalid-encoding",
+    }
+    kind = target.split("-")[1]
+    if kind == "account":
+        kwargs.update(client=client, database_name="explicit_db")
+    elif kind == "database":
+        kwargs.update(database=database, database_name="test_database")
+    else:
+        kwargs.update(collection=collection, database_name="ignored", collection_name="documents")
+    if target.startswith("collection") and kind != "collection":
+        kwargs["collection_name"] = "explicit_collection"
+
+    with patch.object(module, "load_settings", side_effect=AssertionError("settings must be bypassed")):
+        resource: Any
+        if target.startswith("store"):
+            resource = AzureDocumentDBStore(**kwargs)
+        else:
+            resource = AzureDocumentDBCollection(dict, definition=definition_factory(), **kwargs)
+        await resource.aclose()
+
+    assert not resource.managed_client
+    client.close.assert_not_awaited()
+
+
+@pytest.mark.parametrize("target", ["store", "collection"])
+@pytest.mark.parametrize("explicit", [False, True])
+def test_injected_account_resolves_database_independently(
+    target, explicit, definition_factory, pymongo_objects, monkeypatch
+):
+    client, database, _ = pymongo_objects
+    monkeypatch.setenv("AZURE_DOCUMENTDB_DATABASE_NAME", "environment_db")
+    kwargs: dict[str, Any] = {"client": client}
+    if explicit:
+        kwargs["database_name"] = "explicit_db"
+    resource: Any
+    if target == "store":
+        resource = AzureDocumentDBStore(**kwargs)
+    else:
+        resource = AzureDocumentDBCollection(dict, definition=definition_factory(), **kwargs)
+
+    expected = "explicit_db" if explicit else "environment_db"
+    client.__getitem__.assert_called_once_with(expected)
+    assert resource._client.database is database
+
+
+@pytest.mark.parametrize("kind", ["account", "database"])
+@pytest.mark.parametrize("explicit", [False, True])
+def test_injected_collection_resolves_collection_independently(
+    kind, explicit, definition_factory, pymongo_objects, monkeypatch
+):
+    client, database, _ = pymongo_objects
+    monkeypatch.setenv("AZURE_DOCUMENTDB_COLLECTION_NAME", "environment_collection")
+    definition = VectorStoreCollectionDefinition(definition_factory().fields)
+    kwargs = {"client": client, "database_name": "db"} if kind == "account" else {"database": database}
+    resource = AzureDocumentDBCollection(
+        dict,
+        definition=definition,
+        collection_name="explicit_collection" if explicit else None,
+        **kwargs,
+    )
+
+    assert resource.collection_name == ("explicit_collection" if explicit else "environment_collection")
+
+
+@pytest.mark.parametrize("env_name", [None, "environment_collection"])
+def test_injected_collection_uses_registered_definition_before_environment(
+    env_name, definition_factory, pymongo_objects, monkeypatch
+):
+    if env_name is None:
+        monkeypatch.delenv("AZURE_DOCUMENTDB_COLLECTION_NAME", raising=False)
+    else:
+        monkeypatch.setenv("AZURE_DOCUMENTDB_COLLECTION_NAME", env_name)
+    definition = VectorStoreCollectionDefinition(definition_factory().fields, collection_name="registered_collection")
+    resource = AzureDocumentDBCollection(
+        dict,
+        definition=definition,
+        client=pymongo_objects[0],
+        database_name="db",
+    )
+
+    assert resource.collection_name == "registered_collection"
+    assert resource.definition is definition
+
+
+@pytest.mark.parametrize("kind", ["account", "database"])
+def test_registered_model_definition_precedes_environment_without_explicit_definition(
+    kind, definition_factory, pymongo_objects, monkeypatch
+):
+    class RegisteredRecord:
+        pass
+
+    definition = VectorStoreCollectionDefinition(definition_factory().fields, collection_name="registered_collection")
+    register_vectorstoremodel(RegisteredRecord, definition=definition, decoder=lambda data: RegisteredRecord())
+    monkeypatch.setenv("AZURE_DOCUMENTDB_COLLECTION_NAME", "environment_collection")
+    client, database, _ = pymongo_objects
+    kwargs: dict[str, Any] = {"client": client, "database_name": "db"} if kind == "account" else {"database": database}
+
+    resource = AzureDocumentDBCollection(RegisteredRecord, **kwargs)
+
+    assert resource.collection_name == "registered_collection"
+    assert resource.definition is definition
+
+
+@pytest.mark.parametrize("target", ["store-account", "collection-account", "collection-database"])
+def test_injected_objects_report_missing_selector(target, definition_factory, pymongo_objects, monkeypatch):
+    monkeypatch.delenv("AZURE_DOCUMENTDB_DATABASE_NAME", raising=False)
+    monkeypatch.delenv("AZURE_DOCUMENTDB_COLLECTION_NAME", raising=False)
+    client, database, _ = pymongo_objects
+    kwargs = {"database": database} if target.endswith("database") else {"client": client}
+    selector = "AZURE_DOCUMENTDB_DATABASE_NAME" if target == "store-account" else "AZURE_DOCUMENTDB_COLLECTION_NAME"
+    with pytest.raises(SettingNotFoundError, match=selector) as exc:
+        if target.startswith("store"):
+            AzureDocumentDBStore(**kwargs)
+        else:
+            AzureDocumentDBCollection(
+                dict,
+                definition=VectorStoreCollectionDefinition(definition_factory().fields),
+                **kwargs,
+            )
+    parameter = "database_name" if target == "store-account" else "collection_name"
+    assert parameter in str(exc.value)
+
+
+@pytest.mark.parametrize("target", ["store-database", "collection-database", "collection-collection"])
+def test_injected_proxy_rejects_mismatched_name(target, definition_factory, pymongo_objects):
+    _, database, collection = pymongo_objects
+    kwargs: dict[str, Any] = {"collection": collection, "collection_name": "other"}
+    message = "collection_name must match the injected collection."
+    if target.endswith("database"):
+        kwargs = {"database": database, "database_name": "other"}
+        message = "database_name must match the injected database."
+    with pytest.raises(ValueError) as exc:
+        if target.startswith("store"):
+            AzureDocumentDBStore(**kwargs)
+        else:
+            AzureDocumentDBCollection(dict, definition=definition_factory(), **kwargs)
+    assert str(exc.value) == message
 
 
 async def test_store_children_reuse_resolved_client(definition_factory, pymongo_objects):