From 0a71d8cf6a137674e71030d6f5c88999d36728ba Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Oct=C3=A1vio=20Lage?= Date: Tue, 19 Sep 2023 00:45:03 -0300 Subject: [PATCH 1/5] Add distinct function to apache-airflow-providers-mongo --- airflow/providers/mongo/hooks/mongo.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/airflow/providers/mongo/hooks/mongo.py b/airflow/providers/mongo/hooks/mongo.py index fca855a51d98d..eedfd1d50b311 100644 --- a/airflow/providers/mongo/hooks/mongo.py +++ b/airflow/providers/mongo/hooks/mongo.py @@ -366,3 +366,22 @@ def delete_many( collection = self.get_collection(mongo_collection, mongo_db=mongo_db) return collection.delete_many(filter_doc, **kwargs) + + def distinct( + self, mongo_collection: str, key: str, filter_doc: dict | None, mongo_db: str | None = None, **kwargs + ) -> list[Any]: + """ + Returns a list of distinct values for the given key across a collection. + + https://pymongo.readthedocs.io/en/stable/api/pymongo/collection.html#pymongo.collection.Collection.distinct + + :param mongo_collection: The name of the collection to perform distinct on. + :param key: The field to return distinct values from. + :param filter_doc: A query that matches the documents get distinct values from. + Optional. Defaults to {}. + :param mongo_db: The name of the database to use. + Can be omitted; then the database from the connection string is used. + """ + collection = self.get_collection(mongo_collection, mongo_db=mongo_db) + + return collection.distinct(key, filter=filter_doc, **kwargs) From dc048f7499ccb9b5a8855274fd061e7782fb39b9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Oct=C3=A1vio=20Lage?= Date: Tue, 19 Sep 2023 01:39:59 -0300 Subject: [PATCH 2/5] Add unit tests for distinct() function in MongoHook in apache-airflow-provider-mongo --- tests/providers/mongo/hooks/test_mongo.py | 26 +++++++++++++++++++++++ 1 file changed, 26 insertions(+) diff --git a/tests/providers/mongo/hooks/test_mongo.py b/tests/providers/mongo/hooks/test_mongo.py index 4d46a613e755d..19aa4928fc092 100644 --- a/tests/providers/mongo/hooks/test_mongo.py +++ b/tests/providers/mongo/hooks/test_mongo.py @@ -303,6 +303,32 @@ def test_aggregate(self): results = self.hook.aggregate(collection, aggregate_query) assert len(list(results)) == 2 + def test_distinct(self): + collection = mongomock.MongoClient().db.collection + objs = [ + {"test_id": "1", "test_status": "success"}, + {"test_id": "2", "test_status": "failure"}, + {"test_id": "3", "test_status": "success"}, + ] + + collection.insert_many(objs) + + results = self.hook.distinct(collection, "test_status") + assert len(results) == 2 + + def test_distinct_with_filter(self): + collection = mongomock.MongoClient().db.collection + objs = [ + {"test_id": "1", "test_status": "success"}, + {"test_id": "2", "test_status": "failure"}, + {"test_id": "3", "test_status": "success"}, + ] + + collection.insert_many(objs) + + results = self.hook.distinct(collection, "test_id", {"test_status": "failure"}) + assert len(results) == 1 + def test_context_manager(): with MongoHook(conn_id="mongo_default", mongo_db="default") as ctx_hook: From ae52226fd2128dd32a5073acd1899886144dd6f5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Oct=C3=A1vio=20Lage?= Date: Tue, 19 Sep 2023 11:59:02 -0300 Subject: [PATCH 3/5] Refactor parameter name to distinct_key for clarity and update docstring --- airflow/providers/mongo/hooks/mongo.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/airflow/providers/mongo/hooks/mongo.py b/airflow/providers/mongo/hooks/mongo.py index eedfd1d50b311..da58fbbad38d9 100644 --- a/airflow/providers/mongo/hooks/mongo.py +++ b/airflow/providers/mongo/hooks/mongo.py @@ -368,7 +368,7 @@ def delete_many( return collection.delete_many(filter_doc, **kwargs) def distinct( - self, mongo_collection: str, key: str, filter_doc: dict | None, mongo_db: str | None = None, **kwargs + self, mongo_collection: str, distinct_key: str, filter_doc: dict | None, mongo_db: str | None = None, **kwargs ) -> list[Any]: """ Returns a list of distinct values for the given key across a collection. @@ -376,12 +376,12 @@ def distinct( https://pymongo.readthedocs.io/en/stable/api/pymongo/collection.html#pymongo.collection.Collection.distinct :param mongo_collection: The name of the collection to perform distinct on. - :param key: The field to return distinct values from. + :param distinct_key: The field to return distinct values from. :param filter_doc: A query that matches the documents get distinct values from. - Optional. Defaults to {}. + Can be omitted; then will cover the entire collection. :param mongo_db: The name of the database to use. Can be omitted; then the database from the connection string is used. """ collection = self.get_collection(mongo_collection, mongo_db=mongo_db) - return collection.distinct(key, filter=filter_doc, **kwargs) + return collection.distinct(distinct_key, filter=filter_doc, **kwargs) From 8dda8fde6158d31e7e4bfa46f7edfdbb724e63fa Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Oct=C3=A1vio=20Lage?= Date: Tue, 26 Sep 2023 22:30:57 -0300 Subject: [PATCH 4/5] Add default value to filter_doc in distinct function in MongoHook --- airflow/providers/mongo/hooks/mongo.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/airflow/providers/mongo/hooks/mongo.py b/airflow/providers/mongo/hooks/mongo.py index da58fbbad38d9..2880683b498ce 100644 --- a/airflow/providers/mongo/hooks/mongo.py +++ b/airflow/providers/mongo/hooks/mongo.py @@ -368,7 +368,7 @@ def delete_many( return collection.delete_many(filter_doc, **kwargs) def distinct( - self, mongo_collection: str, distinct_key: str, filter_doc: dict | None, mongo_db: str | None = None, **kwargs + self, mongo_collection: str, distinct_key: str, filter_doc: dict | None = None, mongo_db: str | None = None, **kwargs ) -> list[Any]: """ Returns a list of distinct values for the given key across a collection. From 8db28f11cdb2d9b87936200fd8ea8c7054d42771 Mon Sep 17 00:00:00 2001 From: Andrey Anshin Date: Wed, 25 Oct 2023 20:14:10 +0400 Subject: [PATCH 5/5] Fix static checks --- airflow/providers/mongo/hooks/mongo.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/airflow/providers/mongo/hooks/mongo.py b/airflow/providers/mongo/hooks/mongo.py index 2880683b498ce..928618a9dc4c6 100644 --- a/airflow/providers/mongo/hooks/mongo.py +++ b/airflow/providers/mongo/hooks/mongo.py @@ -368,7 +368,12 @@ def delete_many( return collection.delete_many(filter_doc, **kwargs) def distinct( - self, mongo_collection: str, distinct_key: str, filter_doc: dict | None = None, mongo_db: str | None = None, **kwargs + self, + mongo_collection: str, + distinct_key: str, + filter_doc: dict | None = None, + mongo_db: str | None = None, + **kwargs, ) -> list[Any]: """ Returns a list of distinct values for the given key across a collection.