Skip to content
Merged
Changes from 3 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
7 changes: 7 additions & 0 deletions sdks/python/apache_beam/io/mongodbio.py
Original file line number Diff line number Diff line change
Expand Up @@ -71,6 +71,7 @@
import logging
import math
import struct
from importlib import metadata
from typing import Union

import apache_beam as beam
Expand All @@ -97,6 +98,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 +107,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 +265,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", metadata.version("apache-beam")))

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