nevcohen commented on code in PR #46997:
URL: https://github.com/apache/airflow/pull/46997#discussion_r1967807683
##########
providers/common/sql/src/airflow/providers/common/sql/operators/sql.py:
##########
@@ -223,6 +223,8 @@ class SQLExecuteQueryOperator(BaseSQLOperator):
:param return_last: (optional) return the result of only last statement
(default: True).
:param show_return_value_in_logs: (optional) if true operator output will
be printed to the task log.
Use with caution. It's not recommended to dump large datasets to the
log. (default: False).
+ :param requires_result_fetch: (optional) if True, ensures that query
results are fetched before
Review Comment:
Also add something about the fact that if the `do_xcom_push` is true then
there is no need to put the new value you added as true as well.
##########
providers/common/sql/tests/unit/common/sql/operators/test_sql.py:
##########
@@ -109,12 +109,13 @@ def
test_when_provider_min_airflow_version_is_3_0_or_higher_remove_obsolete_get_
class TestSQLExecuteQueryOperator:
- def _construct_operator(self, sql, **kwargs):
+ def _construct_operator(self, sql, requires_result_fetch, **kwargs):
Review Comment:
No need to add it here, it could already be part of the `kwargs`.
##########
providers/common/sql/tests/unit/common/sql/operators/test_sql.py:
##########
@@ -149,6 +150,21 @@ def test_dont_xcom_push(self, mock_get_db_hook,
mock_process_output):
)
mock_process_output.assert_not_called()
+ @mock.patch.object(SQLExecuteQueryOperator, "_process_output")
+ @mock.patch.object(SQLExecuteQueryOperator, "get_db_hook")
+ def test_requires_result_fetch(self, mock_get_db_hook,
mock_process_output):
+ operator = self._construct_operator("SELECT 1;",
requires_result_fetch=True)
+ operator.execute(context=MagicMock())
+
+ mock_get_db_hook.return_value.run.assert_called_once_with(
+ sql="SELECT 1;",
+ autocommit=False,
+ handler=fetch_all_handler,
+ parameters=None,
+ return_last=True,
+ )
+ mock_process_output.assert_called()
Review Comment:
If `requires_result_fetch` is true and `do_xcom_push` is false, this
function should not be called.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
