jacobcbeaudin commented on code in PR #63470:
URL: https://github.com/apache/airflow/pull/63470#discussion_r3906513723
##########
providers/snowflake/src/airflow/providers/snowflake/operators/snowflake.py:
##########
@@ -524,3 +524,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:
+ escaped = ", ".join(f"'{p.replace(chr(39), chr(39) + chr(39))}'"
for p in self.parameters)
+ params_clause = escaped
Review Comment:
I tested this against a live Snowflake account. Snowflake accepts the
escaped SQL. The field has the name `notebook_parameters` now, and the escape
logic is unchanged. See the test results in my comment above.
##########
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:
@jroachgolf84 I answered this in my comment above. Airflow writes the
rendered values back into the template fields in place. Thus `execute()` must
build the SQL again. The comment has a test you can run without a Snowflake
account.
--
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]
