Copilot commented on code in PR #63470:
URL: https://github.com/apache/airflow/pull/63470#discussion_r3066476616
##########
providers/snowflake/src/airflow/providers/snowflake/operators/snowflake.py:
##########
@@ -526,3 +526,48 @@ def on_kill(self) -> None:
self.log.info("Cancelling the query ids %s", self.query_ids)
self._hook.cancel_queries(self.query_ids)
self.log.info("Query ids %s cancelled successfully",
self.query_ids)
+
+
+class SnowflakeNotebookOperator(SnowflakeSqlApiOperator):
+ """
+ Execute a Snowflake Notebook via the Snowflake SQL API.
+
+ Builds an ``EXECUTE NOTEBOOK`` statement and delegates execution to
+
:class:`~airflow.providers.snowflake.operators.snowflake.SnowflakeSqlApiOperator`,
+ which handles query submission, polling, deferral, and cancellation.
+
+ .. seealso::
+ `Snowflake EXECUTE NOTEBOOK
+ <https://docs.snowflake.com/en/sql-reference/sql/execute-notebook>`_
+
+ :param notebook: Fully-qualified notebook name
+ (e.g. ``MY_DB.MY_SCHEMA.MY_NOTEBOOK``).
+ :param parameters: Optional list of parameter strings to pass to the
+ notebook. Only string values are supported by Snowflake; other
+ data types are interpreted as NULL. Parameters are accessible in
+ the notebook via ``sys.argv``.
+ """
+
+ template_fields: Sequence[str] = tuple(
+ set(SnowflakeSqlApiOperator.template_fields) | {"notebook",
"parameters"}
+ )
+
+ def __init__(
+ self,
+ *,
+ notebook: str,
+ parameters: list[str] | None = None,
+ **kwargs: Any,
+ ) -> None:
+ self.notebook = notebook
+ self.parameters = parameters
+ sql = self._build_execute_notebook_query()
+ super().__init__(sql=sql, statement_count=1, **kwargs)
+
Review Comment:
`sql` is built once in `__init__`, but `notebook`/`parameters` are declared
as `template_fields`. After templating, `self.notebook` / `self.parameters` may
change but `self.sql` will not, so the executed statement can be stale. Rebuild
`self.sql` after template rendering (e.g., at the start of `execute()` / before
deferral) or remove `notebook`/`parameters` from `template_fields` and rely
solely on templating `sql`.
##########
providers/snowflake/src/airflow/providers/snowflake/operators/snowflake.py:
##########
@@ -526,3 +526,48 @@ def on_kill(self) -> None:
self.log.info("Cancelling the query ids %s", self.query_ids)
self._hook.cancel_queries(self.query_ids)
self.log.info("Query ids %s cancelled successfully",
self.query_ids)
+
+
+class SnowflakeNotebookOperator(SnowflakeSqlApiOperator):
+ """
+ Execute a Snowflake Notebook via the Snowflake SQL API.
+
+ Builds an ``EXECUTE NOTEBOOK`` statement and delegates execution to
+
:class:`~airflow.providers.snowflake.operators.snowflake.SnowflakeSqlApiOperator`,
+ which handles query submission, polling, deferral, and cancellation.
+
+ .. seealso::
+ `Snowflake EXECUTE NOTEBOOK
+ <https://docs.snowflake.com/en/sql-reference/sql/execute-notebook>`_
+
+ :param notebook: Fully-qualified notebook name
+ (e.g. ``MY_DB.MY_SCHEMA.MY_NOTEBOOK``).
+ :param parameters: Optional list of parameter strings to pass to the
+ notebook. Only string values are supported by Snowflake; other
+ data types are interpreted as NULL. Parameters are accessible in
+ the notebook via ``sys.argv``.
Review Comment:
The docstring states non-string parameter types are interpreted as `NULL`,
but the implementation assumes every element supports `.replace(...)` and
always emits quoted string literals. Either (a) update the docs to match the
actual supported type (`list[str]` only) and raise a clear error for
non-strings, or (b) implement the documented behavior by converting non-strings
to `NULL` (unquoted) when building the SQL.
##########
providers/snowflake/src/airflow/providers/snowflake/operators/snowflake.py:
##########
@@ -526,3 +526,48 @@ def on_kill(self) -> None:
self.log.info("Cancelling the query ids %s", self.query_ids)
self._hook.cancel_queries(self.query_ids)
self.log.info("Query ids %s cancelled successfully",
self.query_ids)
+
+
+class SnowflakeNotebookOperator(SnowflakeSqlApiOperator):
+ """
+ Execute a Snowflake Notebook via the Snowflake SQL API.
+
+ Builds an ``EXECUTE NOTEBOOK`` statement and delegates execution to
+
:class:`~airflow.providers.snowflake.operators.snowflake.SnowflakeSqlApiOperator`,
+ which handles query submission, polling, deferral, and cancellation.
+
+ .. seealso::
+ `Snowflake EXECUTE NOTEBOOK
+ <https://docs.snowflake.com/en/sql-reference/sql/execute-notebook>`_
+
+ :param notebook: Fully-qualified notebook name
+ (e.g. ``MY_DB.MY_SCHEMA.MY_NOTEBOOK``).
+ :param parameters: Optional list of parameter strings to pass to the
+ notebook. Only string values are supported by Snowflake; other
+ data types are interpreted as NULL. Parameters are accessible in
+ the notebook via ``sys.argv``.
+ """
+
+ template_fields: Sequence[str] = tuple(
+ set(SnowflakeSqlApiOperator.template_fields) | {"notebook",
"parameters"}
+ )
+
+ def __init__(
+ self,
+ *,
+ notebook: str,
+ parameters: list[str] | None = None,
+ **kwargs: Any,
+ ) -> None:
+ self.notebook = notebook
+ self.parameters = parameters
+ sql = self._build_execute_notebook_query()
+ super().__init__(sql=sql, statement_count=1, **kwargs)
+
+ def _build_execute_notebook_query(self) -> str:
+ """Build the ``EXECUTE NOTEBOOK`` SQL statement."""
+ params_clause = ""
+ if self.parameters:
+ sanitized = [p.replace("'", "''") for p in self.parameters] #
escape single quotes for SQL
+ params_clause = ", ".join(f"'{p}'" for p in sanitized)
+ return f"EXECUTE NOTEBOOK {self.notebook}({params_clause})"
Review Comment:
`notebook` is interpolated directly into SQL without validation/quoting.
Since this operator takes an identifier-like input (not free-form SQL), it
should defensively validate it (e.g., allow only expected identifier patterns /
`db.schema.name`) or properly quote identifiers to reduce SQL-injection risk
when values are templated or user-provided.
##########
providers/snowflake/src/airflow/providers/snowflake/operators/snowflake.py:
##########
@@ -526,3 +526,48 @@ def on_kill(self) -> None:
self.log.info("Cancelling the query ids %s", self.query_ids)
self._hook.cancel_queries(self.query_ids)
self.log.info("Query ids %s cancelled successfully",
self.query_ids)
+
+
+class SnowflakeNotebookOperator(SnowflakeSqlApiOperator):
+ """
+ Execute a Snowflake Notebook via the Snowflake SQL API.
+
+ Builds an ``EXECUTE NOTEBOOK`` statement and delegates execution to
+
:class:`~airflow.providers.snowflake.operators.snowflake.SnowflakeSqlApiOperator`,
+ which handles query submission, polling, deferral, and cancellation.
+
+ .. seealso::
+ `Snowflake EXECUTE NOTEBOOK
+ <https://docs.snowflake.com/en/sql-reference/sql/execute-notebook>`_
+
+ :param notebook: Fully-qualified notebook name
+ (e.g. ``MY_DB.MY_SCHEMA.MY_NOTEBOOK``).
+ :param parameters: Optional list of parameter strings to pass to the
+ notebook. Only string values are supported by Snowflake; other
+ data types are interpreted as NULL. Parameters are accessible in
+ the notebook via ``sys.argv``.
+ """
+
+ template_fields: Sequence[str] = tuple(
+ set(SnowflakeSqlApiOperator.template_fields) | {"notebook",
"parameters"}
+ )
+
+ def __init__(
+ self,
+ *,
+ notebook: str,
+ parameters: list[str] | None = None,
+ **kwargs: Any,
+ ) -> None:
+ self.notebook = notebook
+ self.parameters = parameters
+ sql = self._build_execute_notebook_query()
+ super().__init__(sql=sql, statement_count=1, **kwargs)
+
+ def _build_execute_notebook_query(self) -> str:
+ """Build the ``EXECUTE NOTEBOOK`` SQL statement."""
+ params_clause = ""
+ if self.parameters:
+ sanitized = [p.replace("'", "''") for p in self.parameters] #
escape single quotes for SQL
Review Comment:
The docstring states non-string parameter types are interpreted as `NULL`,
but the implementation assumes every element supports `.replace(...)` and
always emits quoted string literals. Either (a) update the docs to match the
actual supported type (`list[str]` only) and raise a clear error for
non-strings, or (b) implement the documented behavior by converting non-strings
to `NULL` (unquoted) when building the SQL.
##########
providers/snowflake/src/airflow/providers/snowflake/operators/snowflake.py:
##########
@@ -526,3 +526,48 @@ def on_kill(self) -> None:
self.log.info("Cancelling the query ids %s", self.query_ids)
self._hook.cancel_queries(self.query_ids)
self.log.info("Query ids %s cancelled successfully",
self.query_ids)
+
+
+class SnowflakeNotebookOperator(SnowflakeSqlApiOperator):
+ """
+ Execute a Snowflake Notebook via the Snowflake SQL API.
+
+ Builds an ``EXECUTE NOTEBOOK`` statement and delegates execution to
+
:class:`~airflow.providers.snowflake.operators.snowflake.SnowflakeSqlApiOperator`,
+ which handles query submission, polling, deferral, and cancellation.
+
+ .. seealso::
+ `Snowflake EXECUTE NOTEBOOK
+ <https://docs.snowflake.com/en/sql-reference/sql/execute-notebook>`_
+
+ :param notebook: Fully-qualified notebook name
+ (e.g. ``MY_DB.MY_SCHEMA.MY_NOTEBOOK``).
+ :param parameters: Optional list of parameter strings to pass to the
+ notebook. Only string values are supported by Snowflake; other
+ data types are interpreted as NULL. Parameters are accessible in
+ the notebook via ``sys.argv``.
+ """
+
+ template_fields: Sequence[str] = tuple(
+ set(SnowflakeSqlApiOperator.template_fields) | {"notebook",
"parameters"}
Review Comment:
Building `template_fields` via a `set(...)` makes the resulting tuple order
non-deterministic across runs/Python versions, which can lead to inconsistent
behavior and harder-to-debug template rendering. Prefer constructing it
deterministically (e.g., `SnowflakeSqlApiOperator.template_fields +
(\"notebook\", \"parameters\")`, optionally guarding against duplicates).
##########
providers/snowflake/tests/unit/snowflake/operators/test_snowflake.py:
##########
@@ -605,3 +609,265 @@ def test_snowflake_sql_api_on_kill_no_queries(self,
mock_cancel_queries):
operator.on_kill()
mock_cancel_queries.assert_not_called()
+
+
+class TestSnowflakeNotebookOperatorSQL:
+ """Tests for SQL query building."""
+
+ def test_build_sql_no_params(self):
+ operator = SnowflakeNotebookOperator(
+ task_id=TASK_ID,
+ notebook=NOTEBOOK,
+ )
+ assert operator.sql == "EXECUTE NOTEBOOK MY_DB.MY_SCHEMA.MY_NOTEBOOK()"
+
+ def test_build_sql_with_params(self):
+ operator = SnowflakeNotebookOperator(
+ task_id=TASK_ID,
+ notebook=NOTEBOOK,
+ parameters=["param1", "target_db=PROD"],
+ )
+ assert operator.sql == "EXECUTE NOTEBOOK
MY_DB.MY_SCHEMA.MY_NOTEBOOK('param1', 'target_db=PROD')"
+
+ def test_build_sql_escapes_single_quotes(self):
+ operator = SnowflakeNotebookOperator(
+ task_id=TASK_ID,
+ notebook=NOTEBOOK,
+ parameters=["O'Brien", "it's"],
+ )
+ assert operator.sql == "EXECUTE NOTEBOOK
MY_DB.MY_SCHEMA.MY_NOTEBOOK('O''Brien', 'it''s')"
+
+ def test_build_sql_empty_params(self):
+ operator = SnowflakeNotebookOperator(
+ task_id=TASK_ID,
+ notebook=NOTEBOOK,
+ parameters=[],
+ )
+ assert operator.sql == "EXECUTE NOTEBOOK MY_DB.MY_SCHEMA.MY_NOTEBOOK()"
+
+ def test_template_fields(self):
+ assert "notebook" in SnowflakeNotebookOperator.template_fields
+ assert "parameters" in SnowflakeNotebookOperator.template_fields
+ assert "snowflake_conn_id" in SnowflakeNotebookOperator.template_fields
+
+ def test_statement_count_is_one(self):
+ operator = SnowflakeNotebookOperator(
+ task_id=TASK_ID,
+ notebook=NOTEBOOK,
+ )
+ assert operator.statement_count == 1
+
+ def test_is_subclass_of_snowflake_sql_api_operator(self):
+ assert issubclass(SnowflakeNotebookOperator, SnowflakeSqlApiOperator)
+
+
[email protected]_test
+class TestSnowflakeNotebookOperator:
+ @pytest.fixture(autouse=True)
+ def setup_tests(self):
Review Comment:
The entire test class is marked `db_test` and clears the DB before/after,
but most tests here use only mocks and `context=None`. Consider narrowing
`@pytest.mark.db_test` (and DB cleanup) to only the tests that truly require
DB-backed `create_context(...)` (e.g., the deferrable/context test) to keep the
unit test suite faster and less operationally heavy.
--
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]
