Fokko commented on a change in pull request #3890: [AIRFLOW-3049] Add extra
operations for Mongo hook
URL: https://github.com/apache/incubator-airflow/pull/3890#discussion_r217894023
##########
File path: airflow/contrib/hooks/mongo_hook.py
##########
@@ -130,3 +130,147 @@ def insert_many(self, mongo_collection, docs,
mongo_db=None, **kwargs):
collection = self.get_collection(mongo_collection, mongo_db=mongo_db)
return collection.insert_many(docs, **kwargs)
+
+ def update_one(self, mongo_collection, filter_doc, update_doc,
+ mongo_db=None, **kwargs):
+ """
+ Updates a single document in a mongo collection.
+
https://api.mongodb.com/python/current/api/pymongo/collection.html#pymongo.collection.Collection.update_one
+
+ :param mongo_collection: The name of the collection to update.
+ :type mongo_collection: str
+ :param filter_doc: A query that matches the documents to update.
+ :type filter_doc: dict
+ :param update_doc: The modifications to apply.
+ :type update_doc: dict
+ :param mongo_db: The name of the database to use.
+ Can be omitted; then the database from the connection string is
used.
+ :type mongo_db: str
+
+ """
+ collection = self.get_collection(mongo_collection, mongo_db=mongo_db)
+
+ return collection.update_one(filter_doc, update_doc, **kwargs)
+
+ def update_many(self, mongo_collection, filter_doc, update_doc,
+ mongo_db=None, **kwargs):
+ """
+ Updates one or more documents in a mongo collection.
+
https://api.mongodb.com/python/current/api/pymongo/collection.html#pymongo.collection.Collection.update_many
+
+ :param mongo_collection: The name of the collection to update.
+ :type mongo_collection: str
+ :param filter_doc: A query that matches the documents to update.
+ :type filter_doc: dict
+ :param update_doc: The modifications to apply.
+ :type update_doc: dict
+ :param mongo_db: The name of the database to use.
+ Can be omitted; then the database from the connection string is
used.
+ :type mongo_db: str
+
+ """
+ collection = self.get_collection(mongo_collection, mongo_db=mongo_db)
+
+ return collection.update_many(filter_doc, update_doc, **kwargs)
+
+ def replace_one(self, mongo_collection, doc, filter_doc=None,
+ mongo_db=None, **kwargs):
+ """
+ Replaces a single document in a mongo collection.
+
https://api.mongodb.com/python/current/api/pymongo/collection.html#pymongo.collection.Collection.replace_one
+
+ If no filter document is given, it is assumed that the replacement
+ document contains the _id field which is then used as filter.
+
+ :param mongo_collection: The name of the collection to update.
+ :type mongo_collection: str
+ :param doc: The new document.
+ :type doc: dict
+ :param filter_doc: A query that matches the documents to replace.
+ Can be omitted; then the _id field from doc will be used.
+ :type filter_doc: dict
+ :param mongo_db: The name of the database to use.
+ Can be omitted; then the database from the connection string is
used.
+ :type mongo_db: str
+ """
+ collection = self.get_collection(mongo_collection, mongo_db=mongo_db)
+
+ if not filter_doc:
+ filter_doc = {'_id': doc['_id']}
+
+ return collection.replace_one(filter_doc, doc, **kwargs)
+
+ def replace_many(self, mongo_collection, docs,
+ filter_docs=None, mongo_db=None, upsert=False,
collation=None,
+ **bulk_kwargs):
Review comment:
Maybe call this one just `kwargs` as well, to make it congruent with the
other calls.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
[email protected]
With regards,
Apache Git Services
