From 39601fce342edfdcc4a5482b6cee269a48a303a2 Mon Sep 17 00:00:00 2001 From: Parman Mohammadalizadeh Date: Thu, 23 Jul 2026 23:01:20 +0200 Subject: [PATCH] Fix MongoToS3Operator aggregate-pipeline detection before rendering mongo_query is a template field, rendered after __init__ runs. Caching is_pipeline = isinstance(self.mongo_query, list) in the constructor inspects the un-rendered value, so a query that resolves to a list only after templating would be sent through find() instead of aggregate(). Decide the aggregate-vs-find path in execute() from the rendered value instead. --- .../amazon/aws/transfers/mongo_to_s3.py | 8 +++--- .../amazon/aws/transfers/test_mongo_to_s3.py | 26 ++++++++++++++++++- .../validate_operators_init_exemptions.txt | 1 - 3 files changed, 30 insertions(+), 5 deletions(-) diff --git a/providers/amazon/src/airflow/providers/amazon/aws/transfers/mongo_to_s3.py b/providers/amazon/src/airflow/providers/amazon/aws/transfers/mongo_to_s3.py index e53668ad3ac08..d8e2929bbea93 100644 --- a/providers/amazon/src/airflow/providers/amazon/aws/transfers/mongo_to_s3.py +++ b/providers/amazon/src/airflow/providers/amazon/aws/transfers/mongo_to_s3.py @@ -84,9 +84,7 @@ def __init__( self.mongo_db = mongo_db self.mongo_collection = mongo_collection - # Grab query and determine if we need to run an aggregate pipeline self.mongo_query = mongo_query - self.is_pipeline = isinstance(self.mongo_query, list) self.mongo_projection = mongo_projection self.s3_bucket = s3_bucket @@ -99,8 +97,12 @@ def execute(self, context: Context): """Is written to depend on transform method.""" s3_conn = S3Hook(self.aws_conn_id) + # mongo_query is a template field; decide aggregate-vs-find from the rendered value + # here rather than in __init__, where it would inspect the un-rendered value. + is_pipeline = isinstance(self.mongo_query, list) + # Grab collection and execute query according to whether or not it is a pipeline - if self.is_pipeline: + if is_pipeline: results: CommandCursor[Any] | Cursor = MongoHook(self.mongo_conn_id).aggregate( mongo_collection=self.mongo_collection, aggregate_query=cast("list", self.mongo_query), diff --git a/providers/amazon/tests/unit/amazon/aws/transfers/test_mongo_to_s3.py b/providers/amazon/tests/unit/amazon/aws/transfers/test_mongo_to_s3.py index b2138d887691d..61876def7a67e 100644 --- a/providers/amazon/tests/unit/amazon/aws/transfers/test_mongo_to_s3.py +++ b/providers/amazon/tests/unit/amazon/aws/transfers/test_mongo_to_s3.py @@ -111,7 +111,7 @@ def test_render_template(self, session, clean_dags_dagruns_and_dagbundles, testi ti.dag_run = dag_run render_template_fields(ti, self.mock_operator) expected_rendered_template = {"$lt": "2017-01-01T00:00:00+00:00Z"} - assert expected_rendered_template == getattr(self.mock_operator, "mongo_query") + assert expected_rendered_template == self.mock_operator.mongo_query @mock.patch("airflow.providers.amazon.aws.transfers.mongo_to_s3.MongoHook") @mock.patch("airflow.providers.amazon.aws.transfers.mongo_to_s3.S3Hook") @@ -140,6 +140,30 @@ def test_execute(self, mock_s3_hook, mock_mongo_hook): string_data=s3_doc_str, key=S3_KEY, bucket_name=S3_BUCKET, replace=False, compression=COMPRESSION ) + @mock.patch("airflow.providers.amazon.aws.transfers.mongo_to_s3.MongoHook") + @mock.patch("airflow.providers.amazon.aws.transfers.mongo_to_s3.S3Hook") + def test_execute_runs_aggregate_when_query_renders_to_list(self, mock_s3_hook, mock_mongo_hook): + """ + mongo_query is a template field, so whether to run an aggregate pipeline is decided from + the rendered value in execute(), not from the un-rendered value in __init__. A query that + resolves to a list after templating must take the aggregate path. + """ + operator = self.mock_operator + # Simulate templating resolving mongo_query to a list (an aggregate pipeline) after __init__. + operator.mongo_query = [{"$match": {"foo": "bar"}}] + mock_mongo_hook.return_value.aggregate.return_value = iter(MOCK_MONGO_RETURN) + mock_s3_hook.return_value.load_string.return_value = True + + operator.execute(None) + + mock_mongo_hook.return_value.aggregate.assert_called_once_with( + mongo_collection=MONGO_COLLECTION, + aggregate_query=[{"$match": {"foo": "bar"}}], + mongo_db=None, + allowDiskUse=False, + ) + mock_mongo_hook.return_value.find.assert_not_called() + @mock.patch("airflow.providers.amazon.aws.transfers.mongo_to_s3.MongoHook") @mock.patch("airflow.providers.amazon.aws.transfers.mongo_to_s3.S3Hook") def test_execute_compress(self, mock_s3_hook, mock_mongo_hook): diff --git a/scripts/ci/prek/validate_operators_init_exemptions.txt b/scripts/ci/prek/validate_operators_init_exemptions.txt index 01fd5bc56dbb7..5cc0a481fefed 100644 --- a/scripts/ci/prek/validate_operators_init_exemptions.txt +++ b/scripts/ci/prek/validate_operators_init_exemptions.txt @@ -23,7 +23,6 @@ providers/amazon/src/airflow/providers/amazon/aws/operators/sagemaker.py::SageMa providers/amazon/src/airflow/providers/amazon/aws/operators/step_function.py::StepFunctionStartExecutionOperator providers/amazon/src/airflow/providers/amazon/aws/transfers/base.py::AwsToAwsBaseOperator providers/amazon/src/airflow/providers/amazon/aws/transfers/gcs_to_s3.py::GCSToS3Operator -providers/amazon/src/airflow/providers/amazon/aws/transfers/mongo_to_s3.py::MongoToS3Operator providers/amazon/src/airflow/providers/amazon/aws/transfers/s3_to_redshift.py::S3ToRedshiftOperator providers/anthropic/src/airflow/providers/anthropic/operators/agent.py::AnthropicAgentSessionOperator providers/apache/hive/src/airflow/providers/apache/hive/sensors/hive_partition.py::HivePartitionSensor