This is an automated email from the ASF dual-hosted git repository.
dstandish 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 2e7b9f5504 Defer to hook setting for split_statements in
SQLExecuteQueryOperator (#28635)
2e7b9f5504 is described below
commit 2e7b9f550403cc6937b3210aaaf9e80e3e944445
Author: Daniel Standish <[email protected]>
AuthorDate: Thu Dec 29 22:04:13 2022 -0800
Defer to hook setting for split_statements in SQLExecuteQueryOperator
(#28635)
Some databases, such as snowflake, require you to split statements in order
to submit multi-statement sql. For such databases, splitting is the natural
default, and we should defer to the hook to control that.
---
airflow/providers/common/sql/operators/sql.py | 11 ++++++++---
airflow/providers/common/sql/operators/sql.pyi | 2 +-
tests/providers/amazon/aws/operators/test_redshift_sql.py | 1 -
tests/providers/common/sql/operators/test_sql.py | 2 --
tests/providers/exasol/operators/test_exasol.py | 2 --
tests/providers/jdbc/operators/test_jdbc.py | 2 --
tests/providers/oracle/operators/test_oracle.py | 1 -
tests/providers/trino/operators/test_trino.py | 1 -
tests/providers/vertica/operators/test_vertica.py | 1 -
9 files changed, 9 insertions(+), 14 deletions(-)
diff --git a/airflow/providers/common/sql/operators/sql.py
b/airflow/providers/common/sql/operators/sql.py
index cc2e9e4b57..c4705f77e5 100644
--- a/airflow/providers/common/sql/operators/sql.py
+++ b/airflow/providers/common/sql/operators/sql.py
@@ -198,7 +198,8 @@ class SQLExecuteQueryOperator(BaseSQLOperator):
:param autocommit: (optional) if True, each command is automatically
committed (default: False).
:param parameters: (optional) the parameters to render the SQL query with.
:param handler: (optional) the function that will be applied to the cursor
(default: fetch_all_handler).
- :param split_statements: (optional) if split single SQL string into
statements (default: False).
+ :param split_statements: (optional) if split single SQL string into
statements. By default, defers
+ to the default value in the ``run`` method of the configured hook.
:param return_last: (optional) return the result of only last statement
(default: True).
.. seealso::
@@ -218,7 +219,7 @@ class SQLExecuteQueryOperator(BaseSQLOperator):
autocommit: bool = False,
parameters: Mapping | Iterable | None = None,
handler: Callable[[Any], Any] = fetch_all_handler,
- split_statements: bool = False,
+ split_statements: bool | None = None,
return_last: bool = True,
**kwargs,
) -> None:
@@ -252,13 +253,17 @@ class SQLExecuteQueryOperator(BaseSQLOperator):
def execute(self, context):
self.log.info("Executing: %s", self.sql)
hook = self.get_db_hook()
+ if self.split_statements is not None:
+ extra_kwargs = {"split_statements": self.split_statements}
+ else:
+ extra_kwargs = {}
output = hook.run(
sql=self.sql,
autocommit=self.autocommit,
parameters=self.parameters,
handler=self.handler if self.do_xcom_push else None,
- split_statements=self.split_statements,
return_last=self.return_last,
+ **extra_kwargs,
)
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
diff --git a/airflow/providers/common/sql/operators/sql.pyi
b/airflow/providers/common/sql/operators/sql.pyi
index 5e95dd8c57..956d9ffa08 100644
--- a/airflow/providers/common/sql/operators/sql.pyi
+++ b/airflow/providers/common/sql/operators/sql.pyi
@@ -75,7 +75,7 @@ class SQLExecuteQueryOperator(BaseSQLOperator):
autocommit: bool = ...,
parameters: Union[Mapping, Iterable, None] = ...,
handler: Callable[[Any], Any] = ...,
- split_statements: bool = ...,
+ split_statements: Union[bool, None] = ...,
return_last: bool = ...,
**kwargs,
) -> None: ...
diff --git a/tests/providers/amazon/aws/operators/test_redshift_sql.py
b/tests/providers/amazon/aws/operators/test_redshift_sql.py
index e865be0c65..5908223654 100644
--- a/tests/providers/amazon/aws/operators/test_redshift_sql.py
+++ b/tests/providers/amazon/aws/operators/test_redshift_sql.py
@@ -43,5 +43,4 @@ class TestRedshiftSQLOperator:
parameters=test_parameters,
handler=fetch_all_handler,
return_last=True,
- split_statements=False,
)
diff --git a/tests/providers/common/sql/operators/test_sql.py
b/tests/providers/common/sql/operators/test_sql.py
index be9b18c4b3..6a6972835f 100644
--- a/tests/providers/common/sql/operators/test_sql.py
+++ b/tests/providers/common/sql/operators/test_sql.py
@@ -75,7 +75,6 @@ class TestSQLExecuteQueryOperator:
handler=fetch_all_handler,
parameters=None,
return_last=True,
- split_statements=False,
)
@mock.patch.object(SQLExecuteQueryOperator, "get_db_hook")
@@ -87,7 +86,6 @@ class TestSQLExecuteQueryOperator:
sql="SELECT 1;",
autocommit=False,
parameters=None,
- split_statements=False,
handler=None,
return_last=True,
)
diff --git a/tests/providers/exasol/operators/test_exasol.py
b/tests/providers/exasol/operators/test_exasol.py
index ffc06f6eb6..d893177f21 100644
--- a/tests/providers/exasol/operators/test_exasol.py
+++ b/tests/providers/exasol/operators/test_exasol.py
@@ -34,7 +34,6 @@ class TestExasol:
parameters=None,
handler=fetch_all_handler,
return_last=True,
- split_statements=False,
)
@mock.patch("airflow.providers.common.sql.operators.sql.SQLExecuteQueryOperator.get_db_hook")
@@ -47,7 +46,6 @@ class TestExasol:
parameters={"value": 1},
handler=fetch_all_handler,
return_last=True,
- split_statements=False,
)
@mock.patch("airflow.providers.common.sql.operators.sql.BaseSQLOperator.__init__")
diff --git a/tests/providers/jdbc/operators/test_jdbc.py
b/tests/providers/jdbc/operators/test_jdbc.py
index 1b149bbb1f..904f7df9d9 100644
--- a/tests/providers/jdbc/operators/test_jdbc.py
+++ b/tests/providers/jdbc/operators/test_jdbc.py
@@ -38,7 +38,6 @@ class TestJdbcOperator:
handler=fetch_all_handler,
parameters=jdbc_operator.parameters,
return_last=True,
- split_statements=False,
)
@patch("airflow.providers.common.sql.operators.sql.SQLExecuteQueryOperator.get_db_hook")
@@ -52,5 +51,4 @@ class TestJdbcOperator:
parameters=jdbc_operator.parameters,
handler=None,
return_last=True,
- split_statements=False,
)
diff --git a/tests/providers/oracle/operators/test_oracle.py
b/tests/providers/oracle/operators/test_oracle.py
index a48e980b6d..2950da2bef 100644
--- a/tests/providers/oracle/operators/test_oracle.py
+++ b/tests/providers/oracle/operators/test_oracle.py
@@ -54,7 +54,6 @@ class TestOracleOperator:
parameters=parameters,
handler=fetch_all_handler,
return_last=True,
- split_statements=False,
)
diff --git a/tests/providers/trino/operators/test_trino.py
b/tests/providers/trino/operators/test_trino.py
index 8fef756ecf..8fc1353145 100644
--- a/tests/providers/trino/operators/test_trino.py
+++ b/tests/providers/trino/operators/test_trino.py
@@ -47,5 +47,4 @@ class TestTrinoOperator:
handler=list,
parameters=None,
return_last=True,
- split_statements=False,
)
diff --git a/tests/providers/vertica/operators/test_vertica.py
b/tests/providers/vertica/operators/test_vertica.py
index 3fd6aa3e52..8c561add3b 100644
--- a/tests/providers/vertica/operators/test_vertica.py
+++ b/tests/providers/vertica/operators/test_vertica.py
@@ -35,5 +35,4 @@ class TestVerticaOperator:
handler=fetch_all_handler,
parameters=None,
return_last=True,
- split_statements=False,
)
