Fix MongoToS3Operator aggregate-pipeline detection before rendering

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),

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):

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