[GitHub] [beam] Abacn commented on a diff in pull request #22400: [MongoDBio] custom writeFn for mongodb io python sdk

[GitBox <gi...@apache.org>]

Abacn commented on code in PR #22400:
URL: https://github.com/apache/beam/pull/22400#discussion_r983982931


##########
sdks/python/apache_beam/io/mongodbio.py:
##########
@@ -772,36 +796,46 @@ def display_data(self):
 
 
 class _MongoSink:
-  def __init__(self, uri=None, db=None, coll=None, extra_params=None):
+  def __init__(
+      self, uri=None, db=None, coll=None, extra_params=None, writeFn=None):
     if extra_params is None:
       extra_params = {}
     self.uri = uri
     self.db = db
     self.coll = coll
     self.spec = extra_params
     self.client = None
+    self.writeFn = writeFn
+    if writeFn is None:
+      self.writeFn = self._defaultWriteFn
+
+  @staticmethod
+  def _defaultWriteFn(client, db, coll, documents, logger):

Review Comment:
   We can name this as _ReplaceOneWriteFunc, and also provide a _UpdateOneWriteFunc



##########
sdks/python/apache_beam/io/mongodbio.py:
##########
@@ -680,6 +680,7 @@ def __init__(
       coll=None,
       batch_size=100,
       extra_client_params=None,
+      writeFn=None,

Review Comment:
   Generally in Bean xxxFn implies a DoFn. This is just a callable used in write. Consider renaming it to 'write_func' or 'write_callback' to avoid ambiguity and in_line_with_python_nomenclature



##########
sdks/python/apache_beam/io/mongodbio.py:
##########
@@ -732,13 +750,18 @@ def process(self, element, *args, **kwargs):
       # 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):
+      self,
+      uri=None,
+      db=None,
+      coll=None,
+      batch_size=100,
+      extra_params=None,
+      writeFn=None):

Review Comment:
   We can define the typehint of write_func argument as Optional[str, typing.Callable]. If it is None or "ReplaceOne" then set write_func= _ReplaceOneWriteFunc; if it is "UpdateOne" then set write_func= _UpdateOneWriteFunc, and document this behavior.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org