You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by sh...@apache.org on 2026/07/24 14:21:30 UTC
(airflow) branch main updated: Fix MongoToS3Operator aggregate-pipeline detection before rendering (#70330)
This is an automated email from the ASF dual-hosted git repository.
shahar1 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new c42540fca87 Fix MongoToS3Operator aggregate-pipeline detection before rendering (#70330)
c42540fca87 is described below
commit c42540fca87dbc4b366ee8b75487c82320b3a346
Author: Parman Mohammadalizadeh <pr...@gmail.com>
AuthorDate: Fri Jul 24 17:51:22 2026 +0330
Fix MongoToS3Operator aggregate-pipeline detection before rendering (#70330)
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.
---
.../providers/amazon/aws/transfers/mongo_to_s3.py | 8 ++++---
.../unit/amazon/aws/transfers/test_mongo_to_s3.py | 26 +++++++++++++++++++++-
.../ci/prek/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 e53668ad3ac..d8e2929bbea 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 @@ class MongoToS3Operator(BaseOperator):
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 @@ class MongoToS3Operator(BaseOperator):
"""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 b2138d88769..61876def7a6 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 @@ class TestMongoToS3Operator:
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 @@ class TestMongoToS3Operator:
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 a4846a7da2e..a1c334b20c2 100644
--- a/scripts/ci/prek/validate_operators_init_exemptions.txt
+++ b/scripts/ci/prek/validate_operators_init_exemptions.txt
@@ -19,7 +19,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