#

# 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.

#

"""This module implements IO classes to read and write data on MongoDB.

Read from MongoDB

-----------------

:class:`ReadFromMongoDB` is a ``PTransform`` that reads from a configured

MongoDB source and returns a ``PCollection`` of dict representing MongoDB

documents.

To configure MongoDB source, the URI to connect to MongoDB server, database

name, collection name needs to be provided.

Example usage::

pipeline | ReadFromMongoDB(uri='mongodb://localhost:27017',

db='testdb',

coll='input')

To read from MongoDB Atlas, use ``bucket_auto`` option to enable

``@bucketAuto`` MongoDB aggregation instead of ``splitVector``

command which is a high-privilege function that cannot be assigned

to any user in Atlas.

Example usage::

pipeline | ReadFromMongoDB(uri='mongodb+srv://user:[email protected]',

db='testdb',

coll='input',

bucket_auto=True)

Write to MongoDB:

-----------------

:class:`WriteToMongoDB` is a ``PTransform`` that writes MongoDB documents to

configured sink, and the write is conducted through a mongodb bulk_write of

``ReplaceOne`` operations. If the document's _id field already existed in the

MongoDB collection, it results in an overwrite, otherwise, a new document

will be inserted.

Example usage::

pipeline | WriteToMongoDB(uri='mongodb://localhost:27017',

db='testdb',

coll='output',

batch_size=10)

No backward compatibility guarantees. Everything in this module is experimental.

"""

# pytype: skip-file

import itertools

import json

import logging

import math

import struct

from typing import Union

import apache_beam as beam

from apache_beam.io import iobase

from apache_beam.io.range_trackers import LexicographicKeyRangeTracker

from apache_beam.io.range_trackers import OffsetRangeTracker

from apache_beam.io.range_trackers import OrderedPositionRangeTracker

from apache_beam.transforms import DoFn

from apache_beam.transforms import PTransform

from apache_beam.transforms import Reshuffle

_LOGGER = logging.getLogger(__name__)

try:

# Mongodb has its own bundled bson, which is not compatible with bson package.

# (https://github.com/py-bson/bson/issues/82). Try to import objectid and if

# it fails because bson package is installed, MongoDB IO will not work but at

# least rest of the SDK will work.

from bson import json_util

from bson import objectid

from bson.objectid import ObjectId

# pymongo also internally depends on bson.

from pymongo import ASCENDING

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

ObjectId = None

ASCENDING = 1

DESCENDING = -1

MongoClient = None

ReplaceOne = None

DriverInfo = None

_LOGGER.warning("Could not find a compatible bson package.")

__all__ = ["ReadFromMongoDB", "WriteToMongoDB"]

class ReadFromMongoDB(PTransform):

"""A ``PTransform`` to read MongoDB documents into a ``PCollection``."""

def __init__(

self,

uri="mongodb://localhost:27017",

db=None,

coll=None,

filter=None,

projection=None,

extra_client_params=None,

bucket_auto=False,

):

"""Initialize a :class:`ReadFromMongoDB`

Args:

uri (str): The MongoDB connection string following the URI format.

db (str): The MongoDB database name.

coll (str): The MongoDB collection name.

filter: A `bson.SON

<https://api.mongodb.com/python/current/api/bson/son.html>`_ object

specifying elements which must be present for a document to be included

in the result set.

projection: A list of field names that should be returned in the result

set or a dict specifying the fields to include or exclude.

extra_client_params(dict): Optional `MongoClient

<https://api.mongodb.com/python/current/api/pymongo/mongo_client.html>`_

parameters.

bucket_auto (bool): If :data:`True`, use MongoDB `$bucketAuto` aggregation

to split collection into bundles instead of `splitVector` command,

which does not work with MongoDB Atlas.

If :data:`False` (the default), use `splitVector` command for bundling.

Returns:

:class:`~apache_beam.transforms.ptransform.PTransform`

"""

if extra_client_params is None:

extra_client_params = {}

if not isinstance(db, str):

raise ValueError("ReadFromMongDB db param must be specified as a string")

if not isinstance(coll, str):

raise ValueError(

"ReadFromMongDB coll param must be specified as a string")

self._mongo_source = _BoundedMongoSource(

uri=uri,

db=db,

coll=coll,

filter=filter,

projection=projection,

extra_client_params=extra_client_params,

bucket_auto=bucket_auto,

)

def expand(self, pcoll):

return pcoll | iobase.Read(self._mongo_source)

class _ObjectIdRangeTracker(OrderedPositionRangeTracker):

"""RangeTracker for tracking mongodb _id of bson ObjectId type."""

def position_to_fraction(

self,

pos: ObjectId,

start: ObjectId,

end: ObjectId,

):

"""Returns the fraction of keys in the range [start, end) that

are less than the given key.

"""

pos_number = _ObjectIdHelper.id_to_int(pos)

start_number = _ObjectIdHelper.id_to_int(start)

end_number = _ObjectIdHelper.id_to_int(end)

return (pos_number - start_number) / (end_number - start_number)

def fraction_to_position(

self,

fraction: float,

start: ObjectId,

end: ObjectId,

):

"""Converts a fraction between 0 and 1

to a position between start and end.

"""

start_number = _ObjectIdHelper.id_to_int(start)

end_number = _ObjectIdHelper.id_to_int(end)

total = end_number - start_number

pos = int(total * fraction + start_number)

# make sure split position is larger than start position and smaller than

# end position.

if pos <= start_number:

return _ObjectIdHelper.increment_id(start, 1)

if pos >= end_number:

return _ObjectIdHelper.increment_id(end, -1)

return _ObjectIdHelper.int_to_id(pos)

class _BoundedMongoSource(iobase.BoundedSource):

"""A MongoDB source that reads a finite amount of input records.

This class defines following operations which can be used to read

MongoDB source efficiently.

* Size estimation - method ``estimate_size()`` may return an accurate

estimation in bytes for the size of the source.

* Splitting into bundles of a given size - method ``split()`` can be used to

split the source into a set of sub-sources (bundles) based on a desired

bundle size.

* Getting a RangeTracker - method ``get_range_tracker()`` should return a

``RangeTracker`` object for a given position range for the position type

of the records returned by the source.

* Reading the data - method ``read()`` can be used to read data from the

source while respecting the boundaries defined by a given

``RangeTracker``.

A runner will perform reading the source in two steps.

(1) Method ``get_range_tracker()`` will be invoked with start and end

positions to obtain a ``RangeTracker`` for the range of positions the

runner intends to read. Source must define a default initial start and end

position range. These positions must be used if the start and/or end

positions passed to the method ``get_range_tracker()`` are ``None``

(2) Method read() will be invoked with the ``RangeTracker`` obtained in the

previous step.

**Mutability**

A ``_BoundedMongoSource`` object should not be mutated while

its methods (for example, ``read()``) are being invoked by a runner. Runner

implementations may invoke methods of ``_BoundedMongoSource`` objects through

multi-threaded and/or reentrant execution modes.

"""

def __init__(

self,

uri=None,

db=None,

coll=None,

filter=None,

projection=None,

extra_client_params=None,

bucket_auto=False,

):

if extra_client_params is None:

extra_client_params = {}

if filter is None:

filter = {}

self.uri = uri

self.db = db

self.coll = coll

self.filter = filter

self.projection = projection

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")

def _estimate_average_document_size(self):

with MongoClient(self.uri, **self.spec) as client:

return client[self.db].command("collstats", self.coll).get("avgObjSize")

def split(

self,

desired_bundle_size: int,

start_position: Union[int, str, bytes, ObjectId] = None,

stop_position: Union[int, str, bytes, ObjectId] = None,

):

"""Splits the source into a set of bundles.

Bundles should be approximately of size ``desired_bundle_size`` bytes.

Args:

desired_bundle_size: the desired size (in bytes) of the bundles returned.

start_position: if specified the given position must be used as the

starting position of the first bundle.

stop_position: if specified the given position must be used as the ending

position of the last bundle.

Returns:

an iterator of objects of type 'SourceBundle' that gives information about

the generated bundles.

"""

desired_bundle_size_in_mb = desired_bundle_size // 1024 // 1024

# for desired bundle size, if desired chunk size smaller than 1mb, use

# MongoDB default split size of 1mb.

desired_bundle_size_in_mb = max(desired_bundle_size_in_mb, 1)

is_initial_split = start_position is None and stop_position is None

start_position, stop_position = self._replace_none_positions(

start_position, stop_position

)

if self.bucket_auto:

# Use $bucketAuto for bundling

split_keys = []

weights = []

for bucket in self._get_auto_buckets(

desired_bundle_size_in_mb,

start_position,

stop_position,

is_initial_split,

):

split_keys.append({"_id": bucket["_id"]["max"]})

weights.append(bucket["count"])

else:

# Use splitVector for bundling

split_keys = self._get_split_keys(

desired_bundle_size_in_mb, start_position, stop_position)

weights = itertools.cycle((desired_bundle_size_in_mb, ))

bundle_start = start_position

for split_key_id, weight in zip(split_keys, weights):

if bundle_start >= stop_position:

break

bundle_end = min(stop_position, split_key_id["_id"])

yield iobase.SourceBundle(

weight=weight,

source=self,

start_position=bundle_start,

stop_position=bundle_end,

)

bundle_start = bundle_end

# add range of last split_key to stop_position

if bundle_start < stop_position:

# bucket_auto mode can come here if not split due to single document

weight = 1 if self.bucket_auto else desired_bundle_size_in_mb

yield iobase.SourceBundle(

weight=weight,

source=self,

start_position=bundle_start,

stop_position=stop_position,

)

def get_range_tracker(

self,

start_position: Union[int, str, ObjectId] = None,

stop_position: Union[int, str, ObjectId] = None,

) -> Union[

_ObjectIdRangeTracker, OffsetRangeTracker, LexicographicKeyRangeTracker]:

"""Returns a RangeTracker for a given position range depending on type.

Args:

start_position: starting position of the range. If 'None' default start

position of the source must be used.

stop_position: ending position of the range. If 'None' default stop

position of the source must be used.

Returns:

a ``_ObjectIdRangeTracker``, ``OffsetRangeTracker``

or ``LexicographicKeyRangeTracker`` depending on the given position range.

"""

start_position, stop_position = self._replace_none_positions(

start_position, stop_position

)

if isinstance(start_position, ObjectId):

return _ObjectIdRangeTracker(start_position, stop_position)

if isinstance(start_position, int):

return OffsetRangeTracker(start_position, stop_position)

if isinstance(start_position, str):

return LexicographicKeyRangeTracker(start_position, stop_position)

raise NotImplementedError(

f"RangeTracker for {type(start_position)} not implemented!")

def read(self, range_tracker):

"""Returns an iterator that reads data from the source.

The returned set of data must respect the boundaries defined by the given

``RangeTracker`` object. For example:

* Returned set of data must be for the range

``[range_tracker.start_position, range_tracker.stop_position)``. Note

that a source may decide to return records that start after

``range_tracker.stop_position``. See documentation in class

``RangeTracker`` for more details. Also, note that framework might

invoke ``range_tracker.try_split()`` to perform dynamic split

operations. range_tracker.stop_position may be updated

dynamically due to successful dynamic split operations.

* Method ``range_tracker.try_split()`` must be invoked for every record

that starts at a split point.

* Method ``range_tracker.record_current_position()`` may be invoked for

records that do not start at split points.

Args:

range_tracker: a ``RangeTracker`` whose boundaries must be respected

when reading data from the source. A runner that reads this

source muss pass a ``RangeTracker`` object that is not

``None``.

Returns:

an iterator of data read by the source.

"""

with MongoClient(self.uri, **self.spec) as client:

all_filters = self._merge_id_filter(

range_tracker.start_position(), range_tracker.stop_position())

docs_cursor = (

client[self.db][self.coll].find(

filter=all_filters,

projection=self.projection).sort([("_id", ASCENDING)]))

for doc in docs_cursor:

if not range_tracker.try_claim(doc["_id"]):

return

yield doc

def display_data(self):

"""Returns the display data associated to a pipeline component."""

res = super().display_data()

res["database"] = self.db

res["collection"] = self.coll

res["filter"] = json.dumps(self.filter, default=json_util.default)

res["projection"] = str(self.projection)

res["bucket_auto"] = self.bucket_auto

return res

@staticmethod

def _range_is_not_splittable(

start_pos: Union[int, str, ObjectId],

end_pos: Union[int, str, ObjectId],

):

"""Return `True` if splitting range doesn't make sense

(single document is not splittable),

Return `False` otherwise.

"""

return ((

isinstance(start_pos, ObjectId) and

start_pos >= _ObjectIdHelper.increment_id(end_pos, -1)) or

(isinstance(start_pos, int) and start_pos >= end_pos - 1) or

(isinstance(start_pos, str) and start_pos >= end_pos))

def _get_split_keys(

self,

desired_chunk_size_in_mb: int,

start_pos: Union[int, str, ObjectId],

end_pos: Union[int, str, ObjectId],

):

"""Calls MongoDB `splitVector` command

to get document ids at split position.

"""

# single document not splittable

if self._range_is_not_splittable(start_pos, end_pos):

return []

with MongoClient(self.uri, **self.spec) as client:

name_space = "%s.%s" % (self.db, self.coll)

return client[self.db].command(

"splitVector",

name_space,

keyPattern={"_id": 1}, # Ascending index

min={"_id": start_pos},

max={"_id": end_pos},

maxChunkSize=desired_chunk_size_in_mb,

)["splitKeys"]

def _get_auto_buckets(

self,

desired_chunk_size_in_mb: int,

start_pos: Union[int, str, ObjectId],

end_pos: Union[int, str, ObjectId],

is_initial_split: bool,

) -> list:

"""Use MongoDB `$bucketAuto` aggregation to split collection into bundles

instead of `splitVector` command, which does not work with MongoDB Atlas.

"""

# single document not splittable

if self._range_is_not_splittable(start_pos, end_pos):

return []

if is_initial_split and not self.filter:

# total collection size in MB

size_in_mb = self.estimate_size() / float(1 << 20)

else:

# size of documents within start/end id range and possibly filtered

documents_count = self._count_id_range(start_pos, end_pos)

avg_document_size = self._estimate_average_document_size()

size_in_mb = documents_count * avg_document_size / float(1 << 20)

if size_in_mb == 0:

# no documents not splittable (maybe a result of filtering)

return []

bucket_count = math.ceil(size_in_mb / desired_chunk_size_in_mb)

with beam.io.mongodbio.MongoClient(self.uri, **self.spec) as client:

pipeline = [

{

# filter by positions and by the custom filter if any

"$match": self._merge_id_filter(start_pos, end_pos)

},

{

"$bucketAuto": {

"groupBy": "$_id", "buckets": bucket_count

}

},

]

buckets = list(

# Use `allowDiskUse` option to avoid aggregation limit of 100 Mb RAM

client[self.db][self.coll].aggregate(pipeline, allowDiskUse=True))

if buckets:

buckets[-1]["_id"]["max"] = end_pos

return buckets

def _merge_id_filter(

self,

start_position: Union[int, str, bytes, ObjectId],

stop_position: Union[int, str, bytes, ObjectId] = None,

) -> dict:

"""Merge the default filter (if any) with refined _id field range

of range_tracker.

$gte specifies start position (inclusive)

and $lt specifies the end position (exclusive),

see more at

https://docs.mongodb.com/manual/reference/operator/query/gte/ and

https://docs.mongodb.com/manual/reference/operator/query/lt/

"""

if stop_position is None:

id_filter = {"_id": {"$gte": start_position}}

else:

id_filter = {"_id": {"$gte": start_position, "$lt": stop_position}}

if self.filter:

all_filters = {

# see more at

# https://docs.mongodb.com/manual/reference/operator/query/and/

"$and": [self.filter.copy(), id_filter]

}

else:

all_filters = id_filter

return all_filters

def _get_head_document_id(self, sort_order):

with MongoClient(self.uri, **self.spec) as client:

cursor = (

client[self.db][self.coll].find(filter={}, projection=[]).sort([

("_id", sort_order)

]).limit(1))

try:

return cursor[0]["_id"]

except IndexError:

raise ValueError("Empty Mongodb collection")

def _replace_none_positions(self, start_position, stop_position):

if start_position is None:

start_position = self._get_head_document_id(ASCENDING)

if stop_position is None:

last_doc_id = self._get_head_document_id(DESCENDING)

# increment last doc id binary value by 1 to make sure the last document

# is not excluded

if isinstance(last_doc_id, ObjectId):

stop_position = _ObjectIdHelper.increment_id(last_doc_id, 1)

elif isinstance(last_doc_id, int):

stop_position = last_doc_id + 1

elif isinstance(last_doc_id, str):

stop_position = last_doc_id + '\x00'

return start_position, stop_position

def _count_id_range(self, start_position, stop_position):

"""Number of documents between start_position (inclusive)

and stop_position (exclusive), respecting the custom filter if any.

"""

with MongoClient(self.uri, **self.spec) as client:

return client[self.db][self.coll].count_documents(

filter=self._merge_id_filter(start_position, stop_position))

class _ObjectIdHelper:

"""A Utility class to manipulate bson object ids."""

@classmethod

def id_to_int(cls, _id: Union[int, ObjectId]) -> int:

"""

Args:

_id: ObjectId required for each MongoDB document _id field.

Returns: Converted integer value of ObjectId's 12 bytes binary value.

"""

if isinstance(_id, int):

return _id

# converts object id binary to integer

# id object is bytes type with size of 12

ints = struct.unpack(">III", _id.binary)

return (ints[0] << 64) + (ints[1] << 32) + ints[2]

@classmethod

def int_to_id(cls, number):

"""

Args:

number(int): The integer value to be used to convert to ObjectId.

Returns: The ObjectId that has the 12 bytes binary converted from the

integer value.

"""

# converts integer value to object id. Int value should be less than

# (2 ^ 96) so it can be convert to 12 bytes required by object id.

if number < 0 or number >= (1 << 96):

raise ValueError("number value must be within [0, %s)" % (1 << 96))

ints = [

(number & 0xFFFFFFFF0000000000000000) >> 64,

(number & 0x00000000FFFFFFFF00000000) >> 32,

number & 0x0000000000000000FFFFFFFF,

]

number_bytes = struct.pack(">III", *ints)

return ObjectId(number_bytes)

@classmethod

def increment_id(

cls,

_id: ObjectId,

inc: int,

) -> ObjectId:

"""

Increment object_id binary value by inc value and return new object id.

Args:

_id: The `_id` to change.

inc(int): The incremental int value to be added to `_id`.

Returns:

`_id` incremented by `inc` value

"""

id_number = _ObjectIdHelper.id_to_int(_id)

new_number = id_number + inc

if new_number < 0 or new_number >= (1 << 96):

raise ValueError(

"invalid incremental, inc value must be within ["

"%s, %s)" % (0 - id_number, 1 << 96 - id_number))

return _ObjectIdHelper.int_to_id(new_number)

class WriteToMongoDB(PTransform):

"""WriteToMongoDB is a ``PTransform`` that writes a ``PCollection`` of

mongodb document to the configured MongoDB server.

In order to make the document writes idempotent so that the bundles are

retry-able without creating duplicates, the PTransform added 2 transformations

before final write stage:

a ``GenerateId`` transform and a ``Reshuffle`` transform.::

-----------------------------------------------

Pipeline --> |GenerateId --> Reshuffle --> WriteToMongoSink|

-----------------------------------------------

(WriteToMongoDB)

The ``GenerateId`` transform adds a random and unique*_id* field to the

documents if they don't already have one, it uses the same format as MongoDB

default. The ``Reshuffle`` transform makes sure that no fusion happens between

``GenerateId`` and the final write stage transform,so that the set of

documents and their unique IDs are not regenerated if final write step is

retried due to a failure. This prevents duplicate writes of the same document

with different unique IDs.

"""

def __init__(

self,

uri="mongodb://localhost:27017",

db=None,

coll=None,

batch_size=100,

extra_client_params=None,

):

"""

Args:

uri (str): The MongoDB connection string following the URI format

db (str): The MongoDB database name

coll (str): The MongoDB collection name

batch_size(int): Number of documents per bulk_write to MongoDB,

default to 100

extra_client_params(dict): Optional `MongoClient

<https://api.mongodb.com/python/current/api/pymongo/mongo_client.html>`_

parameters as keyword arguments

Returns:

:class:`~apache_beam.transforms.ptransform.PTransform`

"""

if extra_client_params is None:

extra_client_params = {}

if not isinstance(db, str):

raise ValueError("WriteToMongoDB db param must be specified as a string")

if not isinstance(coll, str):

raise ValueError(

"WriteToMongoDB coll param must be specified as a string")

self._uri = uri

self._db = db

self._coll = coll

self._batch_size = batch_size

self._spec = extra_client_params

def expand(self, pcoll):

return (

pcoll

| beam.ParDo(_GenerateObjectIdFn())

| Reshuffle()

| beam.ParDo(

_WriteMongoFn(

self._uri, self._db, self._coll, self._batch_size, self._spec)))

class _GenerateObjectIdFn(DoFn):

def process(self, element, *args, **kwargs):

# if _id field already exist we keep it as it is, otherwise the ptransform

# generates a new _id field to achieve idempotent write to mongodb.

if "_id" not in element:

# object.ObjectId() generates a unique identifier that follows mongodb

# default format, if _id is not present in document, mongodb server

# generates it with this same function upon write. However the

# uniqueness of generated id may not be guaranteed if the work load are

# distributed across too many processes. See more on the ObjectId format

# https://docs.mongodb.com/manual/reference/bson-types/#objectid.

element["_id"] = objectid.ObjectId()

yield element

class _WriteMongoFn(DoFn):

def __init__(

self, uri=None, db=None, coll=None, batch_size=100, extra_params=None):

if extra_params is None:

extra_params = {}

self.uri = uri

self.db = db

self.coll = coll

self.spec = extra_params

self.batch_size = batch_size

self.batch = []

def finish_bundle(self):

self._flush()

def process(self, element, *args, **kwargs):

self.batch.append(element)

if len(self.batch) >= self.batch_size:

self._flush()

def _flush(self):

if len(self.batch) == 0:

return

with _MongoSink(self.uri, self.db, self.coll, self.spec) as sink:

sink.write(self.batch)

self.batch = []

def display_data(self):

res = super().display_data()

res["database"] = self.db

res["collection"] = self.coll

res["batch_size"] = self.batch_size

return res

class _MongoSink:

def __init__(self, uri=None, db=None, coll=None, extra_params=None):

if extra_params is None:

extra_params = {}

self.uri = uri

self.db = db

self.coll = coll

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)

requests = []

for doc in documents:

# match document based on _id field, if not found in current collection,

# insert new one, otherwise overwrite it.

requests.append(

ReplaceOne(

filter={"_id": doc.get("_id", None)},

replacement=doc,

upsert=True))

resp = self.client[self.db][self.coll].bulk_write(requests)

_LOGGER.debug(

"BulkWrite to MongoDB result in nModified:%d, nUpserted:%d, "

"nMatched:%d, Errors:%s" % (

resp.modified_count,

resp.upserted_count,

resp.matched_count,

resp.bulk_api_result.get("writeErrors"),

))

def __enter__(self):

if self.client is None:

self.client = MongoClient(host=self.uri, **self.spec)

return self

def __exit__(self, exc_type, exc_val, exc_tb):

if self.client is not None:

self.client.close()