You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by po...@apache.org on 2026/05/10 03:46:37 UTC
(airflow) branch main updated: Added statement timeout support in SnowflakeSqlApiHook & Operator (#63575)
This is an automated email from the ASF dual-hosted git repository.
potiuk 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 f3796805f86 Added statement timeout support in SnowflakeSqlApiHook & Operator (#63575)
f3796805f86 is described below
commit f3796805f86868a0a35c658f783de7ffde351076
Author: Justin Pakzad <11...@users.noreply.github.com>
AuthorDate: Sat May 9 23:46:29 2026 -0400
Added statement timeout support in SnowflakeSqlApiHook & Operator (#63575)
* Add statement timeout support for Snowflake SQL API
* Added missing docstring to sql api operator
---
.../providers/snowflake/hooks/snowflake_sql_api.py | 13 +++++++++++-
.../providers/snowflake/operators/snowflake.py | 9 +++++---
.../unit/snowflake/hooks/test_snowflake_sql_api.py | 24 ++++++++++++++++++++++
3 files changed, 42 insertions(+), 4 deletions(-)
diff --git a/providers/snowflake/src/airflow/providers/snowflake/hooks/snowflake_sql_api.py b/providers/snowflake/src/airflow/providers/snowflake/hooks/snowflake_sql_api.py
index bd0ea6b6485..e61783f6ed7 100644
--- a/providers/snowflake/src/airflow/providers/snowflake/hooks/snowflake_sql_api.py
+++ b/providers/snowflake/src/airflow/providers/snowflake/hooks/snowflake_sql_api.py
@@ -135,7 +135,12 @@ class SnowflakeSqlApiHook(SnowflakeHook):
self.aiohttp_request_kwargs = aiohttp_request_kwargs or {}
def execute_query(
- self, sql: str, statement_count: int, query_tag: str = "", bindings: dict[str, Any] | None = None
+ self,
+ sql: str,
+ statement_count: int,
+ query_tag: str = "",
+ bindings: dict[str, Any] | None = None,
+ timeout: int | None = None,
) -> list[str]:
"""
Run the query in Snowflake using SnowflakeSQL API by making API request.
@@ -150,6 +155,9 @@ class SnowflakeSqlApiHook(SnowflakeHook):
:param bindings: (Optional) Values of bind variables in the SQL statement.
When executing the statement, Snowflake replaces placeholders (? and :name) in
the statement with these specified values.
+ :param timeout: (Optional) Timeout in seconds for statement execution.
+ If not set, the timeout specified by STATEMENT_TIMEOUT_IN_SECONDS is used.
+ To set the timeout to the maximum value (604800 seconds), set timeout to 0.
"""
self.query_ids = []
conn_config = self._get_conn_params()
@@ -183,6 +191,9 @@ class SnowflakeSqlApiHook(SnowflakeHook):
},
}
+ if timeout is not None:
+ data["timeout"] = timeout
+
_, json_response = self._make_api_call_with_retries("POST", url, headers, params, data)
self.log.info("Snowflake SQL POST API response: %s", json_response)
if "statementHandles" in json_response:
diff --git a/providers/snowflake/src/airflow/providers/snowflake/operators/snowflake.py b/providers/snowflake/src/airflow/providers/snowflake/operators/snowflake.py
index 2b02d259e71..8f36c84ac12 100644
--- a/providers/snowflake/src/airflow/providers/snowflake/operators/snowflake.py
+++ b/providers/snowflake/src/airflow/providers/snowflake/operators/snowflake.py
@@ -352,6 +352,9 @@ class SnowflakeSqlApiOperator(SQLExecuteQueryOperator):
:param bindings: (Optional) Values of bind variables in the SQL statement.
When executing the statement, Snowflake replaces placeholders (? and :name) in
the statement with these specified values.
+ :param timeout: (Optional) Timeout in seconds for statement execution.
+ If not set, the timeout specified by STATEMENT_TIMEOUT_IN_SECONDS is used.
+ To set the timeout to the maximum value (604800 seconds), set timeout to 0.
:param deferrable: Run operator in the deferrable mode.
:param snowflake_api_retry_args: An optional dictionary with arguments passed to ``tenacity.Retrying`` & ``tenacity.AsyncRetrying`` classes.
"""
@@ -379,6 +382,7 @@ class SnowflakeSqlApiOperator(SQLExecuteQueryOperator):
token_life_time: timedelta = LIFETIME,
token_renewal_delta: timedelta = RENEWAL_DELTA,
bindings: dict[str, Any] | None = None,
+ timeout: int | None = None,
deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
snowflake_api_retry_args: dict[str, Any] | None = None,
**kwargs: Any,
@@ -389,6 +393,7 @@ class SnowflakeSqlApiOperator(SQLExecuteQueryOperator):
self.token_life_time = token_life_time
self.token_renewal_delta = token_renewal_delta
self.bindings = bindings
+ self.timeout = timeout
self.execute_async = False
self.snowflake_api_retry_args = snowflake_api_retry_args or {}
self.deferrable = deferrable
@@ -425,9 +430,7 @@ class SnowflakeSqlApiOperator(SQLExecuteQueryOperator):
"""
self.log.info("Executing: %s", self.sql)
self.query_ids = self._hook.execute_query(
- self.sql,
- statement_count=self.statement_count,
- bindings=self.bindings,
+ self.sql, statement_count=self.statement_count, bindings=self.bindings, timeout=self.timeout
)
self.log.info("List of query ids %s", self.query_ids)
diff --git a/providers/snowflake/tests/unit/snowflake/hooks/test_snowflake_sql_api.py b/providers/snowflake/tests/unit/snowflake/hooks/test_snowflake_sql_api.py
index 7b2d397e8ab..fcbbb2287a8 100644
--- a/providers/snowflake/tests/unit/snowflake/hooks/test_snowflake_sql_api.py
+++ b/providers/snowflake/tests/unit/snowflake/hooks/test_snowflake_sql_api.py
@@ -431,6 +431,30 @@ class TestSnowflakeSqlApiHook:
"Bindings are not supported for multi-statement queries. Bindings will be ignored."
)
+ @mock.patch(f"{HOOK_PATH}._get_conn_params")
+ @mock.patch(f"{HOOK_PATH}.get_headers")
+ def test_execute_query_with_timeout(
+ self,
+ mock_get_headers,
+ mock_conn_params,
+ mock_requests,
+ ):
+ """Test execute_query method with timeout value set"""
+ mock_conn_params.return_value = CONN_PARAMS
+ mock_get_headers.return_value = HEADERS
+ mock_requests.codes.ok = 200
+ mock_requests.request.side_effect = [
+ create_successful_response_mock({"statementHandle": "uuid"}),
+ ]
+
+ hook = SnowflakeSqlApiHook(snowflake_conn_id="mock_conn_id")
+ hook.execute_query(sql=SINGLE_STMT, statement_count=1, timeout=120)
+
+ call_kwargs = mock_requests.request.call_args
+ payload = call_kwargs.kwargs["json"]
+ assert "timeout" in payload
+ assert payload["timeout"] == 120
+
@pytest.mark.parametrize(
"query_ids",
[