Repository: apache/airflow
Pull request: MongoToS3Operator failed when running with a single query (not aggregate pipeline)
Remove invalid `allowDiskUse` argument when calling `find`. This method does not expect that argument according to official documentation [here](https://pymongo.readthedocs.io/en/stable/api/pymongo/collection.html?highlight=allowDiskUse#pymongo.collection.Collection.find).

closes #15679 

Problem statement:
MongoToS3Operator failed when running with a single query (not aggregate pipeline)
**Apache Airflow version**: 2.0.2

**What happened**:

`MongoToS3Operator` failed when running with a single query (not aggregate pipeline):

```sh
Traceback (most recent call last):
  File "/home/airflow//bin/airflow", line 8, in <module>
    sys.exit(main())
  File "/home/airflow//lib/python3.8/site-packages/airflow/__main__.py", line 40, in main
    args.func(args)
  File "/home/airflow//lib/python3.8/site-packages/airflow/cli/cli_parser.py", line 48, in command
    return func(*args, **kwargs)
  File "/home/airflow//lib/python3.8/site-packages/airflow/utils/cli.py", line 89, in wrapper
    return f(*args, **kwargs)
  File "/home/airflow//lib/python3.8/site-packages/airflow/cli/commands/task_command.py", line 385, in task_test
    ti.run(ignore_task_deps=True, ignore_ti_state=True, test_mode=True)
  File "/home/airflow//lib/python3.8/site-packages/airflow/utils/session.py", line 70, in wrapper
    return func(*args, session=session, **kwargs)
  File "/home/airflow//lib/python3.8/site-packages/airflow/models/taskinstance.py", line 1413, in run
    self._run_raw_task(
  File "/home/airflow//lib/python3.8/site-packages/airflow/utils/session.py", line 67, in wrapper
    return func(*args, **kwargs)
  File "/home/airflow//lib/python3.8/site-packages/airflow/models/taskinstance.py", line 1138, in _run_raw_task
    self._prepare_and_execute_task_with_callbacks(context, task)
  File "/home/airflow//lib/python3.8/site-packages/airflow/models/taskinstance.py", line 1311, in _prepare_and_execute_task_with_callbacks
    result = self._execute_task(context, task_copy)
  File "/home/airflow//lib/python3.8/site-packages/airflow/models/taskinstance.py", line 1341, in _execute_task
    result = task_copy.execute(context=context)
  File "/home/airflow//lib/python3.8/site-packages/airflow/providers/amazon/aws/transfers/mongo_to_s3.py", line 116, in execute
    results = MongoHook(self.mongo_conn_id).find(
  File "/home/airflow//lib/python3.8/site-packages/airflow/providers/mongo/hooks/mongo.py", line 144, in find
    return collection.find(query, **kwargs)
  File "/home/airflow//lib/python3.8/site-packages/pymongo/collection.py", line 1523, in find
    return Cursor(self, *args, **kwargs)
TypeError: __init__() got an unexpected keyword argument 'allowDiskUse'
```

**What you expected to happen**:

I expect the data from MongoDB to be exported to a file in S3 with no errors.

**How to reproduce it**:

Run the following operator with a single  `mongo_query` (no aggregate pipeline):

```python
export_to_s3 = MongoToS3Operator(
    task_id='export_to_s3',
    mongo_conn_id=Variable.get('mongo_conn_id'),
    s3_conn_id=Variable.get('aws_conn_id'),
    mongo_collection='my_mongo_collection',
    mongo_query={},
    s3_bucket=Variable.get('s3_bucket'),
    s3_key="my_data.json",
    replace=True,
    dag=dag,
)
```

Hints from issue discussion:
This seems related to https://github.com/apache/airflow/pull/12033

I'll take a look ASAP
There’s already a PR proposed in #15680 so please take a look at that as well.

Commit message: Remove allowDiskUse argument from mongo hook find

Patch to review:
diff --git a/airflow/providers/amazon/aws/transfers/mongo_to_s3.py b/airflow/providers/amazon/aws/transfers/mongo_to_s3.py
index f5dd7262e6ebf..9bd5a8f5fa115 100644
--- a/airflow/providers/amazon/aws/transfers/mongo_to_s3.py
+++ b/airflow/providers/amazon/aws/transfers/mongo_to_s3.py
@@ -49,8 +49,8 @@ class MongoToS3Operator(BaseOperator):
     :type mongo_db: str
     :param replace: whether or not to replace the file in S3 if it previously existed
     :type replace: bool
-    :param allow_disk_use: in the case you are retrieving a lot of data, you may have
-        to use the disk to save it instead of saving all in the RAM
+    :param allow_disk_use: enables writing to temporary files in the case you are handling large dataset
+        when running an aggregate pipeline.
     :type allow_disk_use: bool
     :param compression: type of compression to use for output file in S3. Currently only gzip is supported.
     :type compression: str
@@ -115,7 +115,6 @@ def execute(self, context) -> bool:
                 mongo_collection=self.mongo_collection,
                 query=cast(dict, self.mongo_query),
                 mongo_db=self.mongo_db,
-                allowDiskUse=self.allow_disk_use,
             )
 
         # Performs transform then stringifies the docs results into json format
diff --git a/tests/providers/amazon/aws/transfers/test_mongo_to_s3.py b/tests/providers/amazon/aws/transfers/test_mongo_to_s3.py
index 9e9811aa6302e..b0bfb9aab694c 100644
--- a/tests/providers/amazon/aws/transfers/test_mongo_to_s3.py
+++ b/tests/providers/amazon/aws/transfers/test_mongo_to_s3.py
@@ -94,7 +94,7 @@ def test_execute(self, mock_s3_hook, mock_mongo_hook):
         operator.execute(None)
 
         mock_mongo_hook.return_value.find.assert_called_once_with(
-            mongo_collection=MONGO_COLLECTION, query=MONGO_QUERY, mongo_db=None, allowDiskUse=False
+            mongo_collection=MONGO_COLLECTION, query=MONGO_QUERY, mongo_db=None
         )
 
         op_stringify = self.mock_operator._stringify
@@ -117,7 +117,7 @@ def test_execute_compress(self, mock_s3_hook, mock_mongo_hook):
         operator.execute(None)
 
         mock_mongo_hook.return_value.find.assert_called_once_with(
-            allowDiskUse=False, mongo_collection=MONGO_COLLECTION, query=MONGO_QUERY, mongo_db=None
+            mongo_collection=MONGO_COLLECTION, query=MONGO_QUERY, mongo_db=None
         )
 
         op_stringify = self.mock_operator._stringify


Code review comments:
diff_hunk: @@ -49,8 +49,8 @@ class MongoToS3Operator(BaseOperator):
     :type mongo_db: str
     :param replace: whether or not to replace the file in S3 if it previously existed
     :type replace: bool
-    :param allow_disk_use: in the case you are retrieving a lot of data, you may have
-        to use the disk to save it instead of saving all in the RAM
+    :param allow_disk_use: enables writing to temporary files in the case you are handling large dataset
+        when running an aggregate pipeline.
path: airflow/providers/amazon/aws/transfers/mongo_to_s3.py
text: @user1:
```suggestion
        This only takes effect when `mongo_query` is a list - running an aggregate pipeline
```

@user1:
@author  Can you also update line:43 to `:type mongo_query: Union[list, dict]`? Much appreciated!

@author:
@user1 Both changes done :heavy_check_mark: