diff --git a/sdks/python/apache_beam/io/mongodbio.py b/sdks/python/apache_beam/io/mongodbio.py index e8602edd40cd..aede27674f8f 100644 --- a/sdks/python/apache_beam/io/mongodbio.py +++ b/sdks/python/apache_beam/io/mongodbio.py @@ -97,6 +97,7 @@ from pymongo import DESCENDING from pymongo import MongoClient from pymongo import ReplaceOne + from pymongo.driver_info import DriverInfo except ImportError: objectid = None json_util = None @@ -105,6 +106,7 @@ DESCENDING = -1 MongoClient = None ReplaceOne = None + DriverInfo = None _LOGGER.warning("Could not find a compatible bson package.") __all__ = ["ReadFromMongoDB", "WriteToMongoDB"] @@ -263,6 +265,12 @@ def __init__( self.spec = extra_client_params self.bucket_auto = bucket_auto + if "driver" not in self.spec: + self.spec["driver"] = DriverInfo( + name="Apache Beam", + version=beam.__version__, + ) + def estimate_size(self): with MongoClient(self.uri, **self.spec) as client: return client[self.db].command("collstats", self.coll).get("size") @@ -777,6 +785,12 @@ def __init__(self, uri=None, db=None, coll=None, extra_params=None): self.spec = extra_params self.client = None + if "driver" not in self.spec: + self.spec["driver"] = DriverInfo( + name="Apache Beam", + version=beam.__version__, + ) + def write(self, documents): if self.client is None: self.client = MongoClient(host=self.uri, **self.spec)