From ddaf121fcc949af0d2ccc9dd2c805f672e3828e1 Mon Sep 17 00:00:00 2001 From: Alex Bevilacqua Date: Wed, 24 Jun 2026 08:32:17 -0400 Subject: [PATCH] feat: add MongoDB driver handshake metadata MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Identify Burr in the MongoDB handshake so server-side telemetry can distinguish Burr traffic. A module-level constant _DRIVER_INFO = DriverInfo(name="Burr", version=_VERSION) is defined once in b_pymongo.py and used at every client construction site. b_pymongo.MongoDBBasePersister has three patterns: Pattern A – from_values() constructs MongoClient: mongo_client_kwargs.setdefault("driver", _DRIVER_INFO) guards against overriding a caller-supplied driver value. Pattern B – __init__() receives an existing client: if hasattr(client, "append_metadata"): client.append_metadata(_DRIVER_INFO) append_metadata is idempotent; the hasattr guard keeps this safe with older pymongo versions that predate the API. Pattern A – __setstate__() reconstructs MongoClient after unpickling: driver=_DRIVER_INFO passed directly. b_mongodb.py is a deprecated shim that constructs its own MongoClient instances in from_values() and __init__(). Both receive the same setdefault guard, importing _DRIVER_INFO from b_pymongo.py. Version is read via importlib.metadata("apache-burr") with a try/except fallback so it works whether or not the package is pip-installed. Adds tests/integrations/persisters/test_b_pymongo_driver_info.py (no live DB required) verifying all five behaviours: - _DRIVER_INFO has name="Burr" and matches the installed version - from_values() passes driver=_DRIVER_INFO to MongoClient - from_values() does not override a caller-supplied driver - __init__() calls append_metadata(_DRIVER_INFO) when the client supports it - __init__() does not raise when client lacks append_metadata Co-Authored-By: Claude Sonnet 4.6 (1M context) --- burr/integrations/persisters/b_mongodb.py | 5 ++ burr/integrations/persisters/b_pymongo.py | 16 +++- .../persisters/test_b_pymongo_driver_info.py | 83 +++++++++++++++++++ 3 files changed, 103 insertions(+), 1 deletion(-) create mode 100644 tests/integrations/persisters/test_b_pymongo_driver_info.py diff --git a/burr/integrations/persisters/b_mongodb.py b/burr/integrations/persisters/b_mongodb.py index f65599667..1e820877c 100644 --- a/burr/integrations/persisters/b_mongodb.py +++ b/burr/integrations/persisters/b_mongodb.py @@ -21,6 +21,9 @@ from pymongo import MongoClient +from burr.integrations.persisters.b_pymongo import ( + _DRIVER_INFO, +) from burr.integrations.persisters.b_pymongo import MongoDBBasePersister as PymongoPersister logger = logging.getLogger(__name__) @@ -50,6 +53,7 @@ def from_values( if mongo_client_kwargs is None: mongo_client_kwargs = {} + mongo_client_kwargs.setdefault("driver", _DRIVER_INFO) client = MongoClient(uri, **mongo_client_kwargs) return PymongoPersister( client=client, @@ -76,6 +80,7 @@ def __init__( """Initializes the MongoDBPersister class.""" if mongo_client_kwargs is None: mongo_client_kwargs = {} + mongo_client_kwargs.setdefault("driver", _DRIVER_INFO) client = MongoClient(uri, **mongo_client_kwargs) super(MongoDBPersister, self).__init__( client=client, diff --git a/burr/integrations/persisters/b_pymongo.py b/burr/integrations/persisters/b_pymongo.py index ba02438e9..85a80fba9 100644 --- a/burr/integrations/persisters/b_pymongo.py +++ b/burr/integrations/persisters/b_pymongo.py @@ -18,12 +18,21 @@ import json import logging from datetime import datetime, timezone +from importlib.metadata import version as get_version from typing import Literal, Optional from pymongo import MongoClient +from pymongo.driver_info import DriverInfo from burr.core import persistence, state +try: + _VERSION = get_version("apache-burr") +except Exception: + _VERSION = None + +_DRIVER_INFO = DriverInfo(name="Burr", version=_VERSION) + logger = logging.getLogger(__name__) @@ -69,6 +78,7 @@ def from_values( """Initializes the MongoDBBasePersister class.""" if mongo_client_kwargs is None: mongo_client_kwargs = {} + mongo_client_kwargs.setdefault("driver", _DRIVER_INFO) client = MongoClient(uri, **mongo_client_kwargs) return cls( client=client, @@ -92,6 +102,8 @@ def __init__( :param serde_kwargs: serializer/deserializer keyword arguments to pass to the state object """ self.client = client + if hasattr(client, "append_metadata"): + client.append_metadata(_DRIVER_INFO) self.db = self.client[db_name] self.collection = self.db[collection_name] self.serde_kwargs = serde_kwargs or {} @@ -215,7 +227,9 @@ def __getstate__(self) -> dict: def __setstate__(self, state: dict): connection_params = state.pop("connection_params") # we assume MongoClient. - self.client = MongoClient(connection_params["uri"], connection_params["port"]) + self.client = MongoClient( + connection_params["uri"], connection_params["port"], driver=_DRIVER_INFO + ) self.db = self.client[connection_params["db_name"]] self.collection = self.db[connection_params["collection_name"]] self.__dict__.update(state) diff --git a/tests/integrations/persisters/test_b_pymongo_driver_info.py b/tests/integrations/persisters/test_b_pymongo_driver_info.py new file mode 100644 index 000000000..3ba26c326 --- /dev/null +++ b/tests/integrations/persisters/test_b_pymongo_driver_info.py @@ -0,0 +1,83 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +"""Unit tests for MongoDB driver handshake metadata (no live DB required).""" + +from unittest.mock import MagicMock, patch + +from pymongo.driver_info import DriverInfo + +from burr.integrations.persisters.b_pymongo import ( + _DRIVER_INFO, + _VERSION, + MongoDBBasePersister, +) + + +def test_driver_info_name(): + assert isinstance(_DRIVER_INFO, DriverInfo) + assert _DRIVER_INFO.name == "Burr" + + +def test_driver_info_version_matches_package(): + assert _DRIVER_INFO.version == _VERSION + + +def test_from_values_passes_driver_info(): + """from_values() injects driver=_DRIVER_INFO into MongoClient.""" + with patch("burr.integrations.persisters.b_pymongo.MongoClient") as mock_ctor: + mock_client = MagicMock() + mock_client.__getitem__ = MagicMock(return_value=MagicMock()) + mock_ctor.return_value = mock_client + MongoDBBasePersister.from_values(uri="mongodb://localhost:27017") + _, kwargs = mock_ctor.call_args + assert "driver" in kwargs + assert isinstance(kwargs["driver"], DriverInfo) + assert kwargs["driver"].name == "Burr" + + +def test_from_values_does_not_override_caller_driver(): + """from_values() preserves a driver value supplied via mongo_client_kwargs.""" + custom = DriverInfo(name="MyApp", version="9.9") + with patch("burr.integrations.persisters.b_pymongo.MongoClient") as mock_ctor: + mock_client = MagicMock() + mock_client.__getitem__ = MagicMock(return_value=MagicMock()) + mock_ctor.return_value = mock_client + MongoDBBasePersister.from_values( + uri="mongodb://localhost:27017", + mongo_client_kwargs={"driver": custom}, + ) + _, kwargs = mock_ctor.call_args + assert kwargs["driver"] is custom + + +def test_init_calls_append_metadata_when_available(): + """__init__() calls append_metadata on a client that supports it.""" + mock_client = MagicMock() + mock_client.__getitem__ = MagicMock(return_value=MagicMock()) + MongoDBBasePersister(client=mock_client) + mock_client.append_metadata.assert_called_once_with(_DRIVER_INFO) + + +def test_init_skips_append_metadata_when_absent(): + """__init__() does not raise when client lacks append_metadata (older driver).""" + mock_client = MagicMock() + mock_client.__getitem__ = MagicMock(return_value=MagicMock()) + # Simulate a client without append_metadata by deleting the attribute + del mock_client.append_metadata + # Should not raise + MongoDBBasePersister(client=mock_client)