Skip to content
Merged
Changes from 4 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 6 additions & 0 deletions sdks/python/apache_beam/io/mongodbio.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -105,6 +106,7 @@
DESCENDING = -1
MongoClient = None
ReplaceOne = None
DriverInfo = None
_LOGGER.warning("Could not find a compatible bson package.")

__all__ = ["ReadFromMongoDB", "WriteToMongoDB"]
Expand Down Expand Up @@ -262,6 +264,10 @@ def __init__(
self.projection = projection
self.spec = extra_client_params
self.bucket_auto = bucket_auto
self.client = MongoClient(
self.uri,
**self.spec,
driver=DriverInfo("Apache Beam", beam.__version__))

def estimate_size(self):
with MongoClient(self.uri, **self.spec) as client:
Expand Down
Loading