Skip to content

Commit 30fa494

Browse files
alexbeviclaude
andauthored
feat: add MongoDB driver handshake metadata (#821)
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) <noreply@anthropic.com>
1 parent 378ba7d commit 30fa494

3 files changed

Lines changed: 103 additions & 1 deletion

File tree

burr/integrations/persisters/b_mongodb.py

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,9 @@
2121

2222
from pymongo import MongoClient
2323

24+
from burr.integrations.persisters.b_pymongo import (
25+
_DRIVER_INFO,
26+
)
2427
from burr.integrations.persisters.b_pymongo import MongoDBBasePersister as PymongoPersister
2528

2629
logger = logging.getLogger(__name__)
@@ -50,6 +53,7 @@ def from_values(
5053

5154
if mongo_client_kwargs is None:
5255
mongo_client_kwargs = {}
56+
mongo_client_kwargs.setdefault("driver", _DRIVER_INFO)
5357
client = MongoClient(uri, **mongo_client_kwargs)
5458
return PymongoPersister(
5559
client=client,
@@ -76,6 +80,7 @@ def __init__(
7680
"""Initializes the MongoDBPersister class."""
7781
if mongo_client_kwargs is None:
7882
mongo_client_kwargs = {}
83+
mongo_client_kwargs.setdefault("driver", _DRIVER_INFO)
7984
client = MongoClient(uri, **mongo_client_kwargs)
8085
super(MongoDBPersister, self).__init__(
8186
client=client,

burr/integrations/persisters/b_pymongo.py

Lines changed: 15 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,12 +18,21 @@
1818
import json
1919
import logging
2020
from datetime import datetime, timezone
21+
from importlib.metadata import version as get_version
2122
from typing import Literal, Optional
2223

2324
from pymongo import MongoClient
25+
from pymongo.driver_info import DriverInfo
2426

2527
from burr.core import persistence, state
2628

29+
try:
30+
_VERSION = get_version("apache-burr")
31+
except Exception:
32+
_VERSION = None
33+
34+
_DRIVER_INFO = DriverInfo(name="Burr", version=_VERSION)
35+
2736
logger = logging.getLogger(__name__)
2837

2938

@@ -69,6 +78,7 @@ def from_values(
6978
"""Initializes the MongoDBBasePersister class."""
7079
if mongo_client_kwargs is None:
7180
mongo_client_kwargs = {}
81+
mongo_client_kwargs.setdefault("driver", _DRIVER_INFO)
7282
client = MongoClient(uri, **mongo_client_kwargs)
7383
return cls(
7484
client=client,
@@ -92,6 +102,8 @@ def __init__(
92102
:param serde_kwargs: serializer/deserializer keyword arguments to pass to the state object
93103
"""
94104
self.client = client
105+
if hasattr(client, "append_metadata"):
106+
client.append_metadata(_DRIVER_INFO)
95107
self.db = self.client[db_name]
96108
self.collection = self.db[collection_name]
97109
self.serde_kwargs = serde_kwargs or {}
@@ -215,7 +227,9 @@ def __getstate__(self) -> dict:
215227
def __setstate__(self, state: dict):
216228
connection_params = state.pop("connection_params")
217229
# we assume MongoClient.
218-
self.client = MongoClient(connection_params["uri"], connection_params["port"])
230+
self.client = MongoClient(
231+
connection_params["uri"], connection_params["port"], driver=_DRIVER_INFO
232+
)
219233
self.db = self.client[connection_params["db_name"]]
220234
self.collection = self.db[connection_params["collection_name"]]
221235
self.__dict__.update(state)
Lines changed: 83 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,83 @@
1+
# Licensed to the Apache Software Foundation (ASF) under one
2+
# or more contributor license agreements. See the NOTICE file
3+
# distributed with this work for additional information
4+
# regarding copyright ownership. The ASF licenses this file
5+
# to you under the Apache License, Version 2.0 (the
6+
# "License"); you may not use this file except in compliance
7+
# with the License. You may obtain a copy of the License at
8+
#
9+
# http://www.apache.org/licenses/LICENSE-2.0
10+
#
11+
# Unless required by applicable law or agreed to in writing,
12+
# software distributed under the License is distributed on an
13+
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
# KIND, either express or implied. See the License for the
15+
# specific language governing permissions and limitations
16+
# under the License.
17+
18+
"""Unit tests for MongoDB driver handshake metadata (no live DB required)."""
19+
20+
from unittest.mock import MagicMock, patch
21+
22+
from pymongo.driver_info import DriverInfo
23+
24+
from burr.integrations.persisters.b_pymongo import (
25+
_DRIVER_INFO,
26+
_VERSION,
27+
MongoDBBasePersister,
28+
)
29+
30+
31+
def test_driver_info_name():
32+
assert isinstance(_DRIVER_INFO, DriverInfo)
33+
assert _DRIVER_INFO.name == "Burr"
34+
35+
36+
def test_driver_info_version_matches_package():
37+
assert _DRIVER_INFO.version == _VERSION
38+
39+
40+
def test_from_values_passes_driver_info():
41+
"""from_values() injects driver=_DRIVER_INFO into MongoClient."""
42+
with patch("burr.integrations.persisters.b_pymongo.MongoClient") as mock_ctor:
43+
mock_client = MagicMock()
44+
mock_client.__getitem__ = MagicMock(return_value=MagicMock())
45+
mock_ctor.return_value = mock_client
46+
MongoDBBasePersister.from_values(uri="mongodb://localhost:27017")
47+
_, kwargs = mock_ctor.call_args
48+
assert "driver" in kwargs
49+
assert isinstance(kwargs["driver"], DriverInfo)
50+
assert kwargs["driver"].name == "Burr"
51+
52+
53+
def test_from_values_does_not_override_caller_driver():
54+
"""from_values() preserves a driver value supplied via mongo_client_kwargs."""
55+
custom = DriverInfo(name="MyApp", version="9.9")
56+
with patch("burr.integrations.persisters.b_pymongo.MongoClient") as mock_ctor:
57+
mock_client = MagicMock()
58+
mock_client.__getitem__ = MagicMock(return_value=MagicMock())
59+
mock_ctor.return_value = mock_client
60+
MongoDBBasePersister.from_values(
61+
uri="mongodb://localhost:27017",
62+
mongo_client_kwargs={"driver": custom},
63+
)
64+
_, kwargs = mock_ctor.call_args
65+
assert kwargs["driver"] is custom
66+
67+
68+
def test_init_calls_append_metadata_when_available():
69+
"""__init__() calls append_metadata on a client that supports it."""
70+
mock_client = MagicMock()
71+
mock_client.__getitem__ = MagicMock(return_value=MagicMock())
72+
MongoDBBasePersister(client=mock_client)
73+
mock_client.append_metadata.assert_called_once_with(_DRIVER_INFO)
74+
75+
76+
def test_init_skips_append_metadata_when_absent():
77+
"""__init__() does not raise when client lacks append_metadata (older driver)."""
78+
mock_client = MagicMock()
79+
mock_client.__getitem__ = MagicMock(return_value=MagicMock())
80+
# Simulate a client without append_metadata by deleting the attribute
81+
del mock_client.append_metadata
82+
# Should not raise
83+
MongoDBBasePersister(client=mock_client)

0 commit comments

Comments
 (0)