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_r217894039
##########
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):
+ """
+ Replaces many documents in a mongo collection.
+
+ Uses bulk_write with multiple ReplaceOne operations
+
https://api.mongodb.com/python/current/api/pymongo/collection.html#pymongo.collection.Collection.bulk_write
+
+ If no filter documents are given, it is assumed that all replacement
+ documents contain the _id field which are then used as filters.
+
+ :param mongo_collection: The name of the collection to update.
+ :type mongo_collection: str
+ :param docs: The new documents.
+ :type docs: list(dict)
+ :param filter_docs: A list of queries that match the documents to
replace.
+ Can be omitted; then the _id fields from docs will be used.
+ :type filter_docs: list(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
Review comment:
The `upsert` and `collation` are missing from the docstring.
----------------------------------------------------------------
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
