[airflow] branch main updated: Do not process output when do_xcom_push=False (#29599)

[bo...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

bolke 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 19f1e7c27b Do not process output when do_xcom_push=False  (#29599)
19f1e7c27b is described below

commit 19f1e7c27b85e297497842c73f13533767ebd6ba
Author: fritz-astronomer <80...@users.noreply.github.com>
AuthorDate: Wed Feb 22 04:33:08 2023 -0500

    Do not process output when do_xcom_push=False  (#29599)
    
    Adds a guard check to short-circuit out of _process_output in the event of no output (such as do_xcom_push=False)
    
    Co-authored-by: Tzu-ping Chung <ur...@gmail.com>
---
 airflow/providers/common/sql/operators/sql.py    | 2 ++
 tests/providers/common/sql/operators/test_sql.py | 8 ++++++--
 2 files changed, 8 insertions(+), 2 deletions(-)

diff --git a/airflow/providers/common/sql/operators/sql.py b/airflow/providers/common/sql/operators/sql.py
index 86c4cd710b..135f0d3594 100644
--- a/airflow/providers/common/sql/operators/sql.py
+++ b/airflow/providers/common/sql/operators/sql.py
@@ -265,6 +265,8 @@ class SQLExecuteQueryOperator(BaseSQLOperator):
             return_last=self.return_last,
             **extra_kwargs,
         )
+        if not self.do_xcom_push:
+            return None
         if return_single_query_results(self.sql, self.return_last, self.split_statements):
             # For simplicity, we pass always list as input to _process_output, regardless if
             # single query results are going to be returned, and we return the first element
diff --git a/tests/providers/common/sql/operators/test_sql.py b/tests/providers/common/sql/operators/test_sql.py
index 6a6972835f..86608e7440 100644
--- a/tests/providers/common/sql/operators/test_sql.py
+++ b/tests/providers/common/sql/operators/test_sql.py
@@ -64,8 +64,9 @@ class TestSQLExecuteQueryOperator:
             dag=dag,
         )
 
+    @mock.patch.object(SQLExecuteQueryOperator, "_process_output")
     @mock.patch.object(SQLExecuteQueryOperator, "get_db_hook")
-    def test_do_xcom_push(self, mock_get_db_hook):
+    def test_do_xcom_push(self, mock_get_db_hook, mock_process_output):
         operator = self._construct_operator("SELECT 1;", do_xcom_push=True)
         operator.execute(context=MagicMock())
 
@@ -76,9 +77,11 @@ class TestSQLExecuteQueryOperator:
             parameters=None,
             return_last=True,
         )
+        mock_process_output.assert_called()
 
+    @mock.patch.object(SQLExecuteQueryOperator, "_process_output")
     @mock.patch.object(SQLExecuteQueryOperator, "get_db_hook")
-    def test_dont_xcom_push(self, mock_get_db_hook):
+    def test_dont_xcom_push(self, mock_get_db_hook, mock_process_output):
         operator = self._construct_operator("SELECT 1;", do_xcom_push=False)
         operator.execute(context=MagicMock())
 
@@ -89,6 +92,7 @@ class TestSQLExecuteQueryOperator:
             handler=None,
             return_last=True,
         )
+        mock_process_output.assert_not_called()
 
 
 class TestColumnCheckOperator: