justinpakzad commented on code in PR #68259: URL: https://github.com/apache/airflow/pull/68259#discussion_r3439585795
########## providers/snowflake/docs/operators/snowpark_containers.rst: ########## @@ -0,0 +1,75 @@ + .. Licensed to the Apache Software Foundation (ASF) under one + or more contributor license agreements. See the NOTICE file + distributed with this work for additional information + regarding copyright ownership. The ASF licenses this file + to you under the Apache License, Version 2.0 (the + "License"); you may not use this file except in compliance + with the License. You may obtain a copy of the License at + + .. http://www.apache.org/licenses/LICENSE-2.0 + + .. Unless required by applicable law or agreed to in writing, + software distributed under the License is distributed on an + "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + KIND, either express or implied. See the License for the + specific language governing permissions and limitations + under the License. + +.. _howto/operator:SnowparkContainerJobOperator: + +SnowparkContainerJobOperator +============================ + +Use the :class:`SnowparkContainerJobOperator +<airflow.providers.snowflake.operators.snowpark_containers.SnowparkContainerJobOperator>` to +submit and monitor container jobs on +`Snowpark Container Services <https://docs.snowflake.com/en/developer-guide/snowpark-container-services/overview>`__. + +The operator wraps the Snowflake +`EXECUTE JOB SERVICE <https://docs.snowflake.com/en/sql-reference/sql/execute-job-service>`__ +SQL command. It submits the job asynchronously, optionally polls until the job +reaches a terminal state, retrieves container logs, and can drop the service on +completion. + +Prerequisite Tasks +^^^^^^^^^^^^^^^^^^ + +To use this operator, you must do a few things: + + * Install the provider package via **pip**. + + .. code-block:: bash + + pip install 'apache-airflow-providers-snowflake' + + Detailed information is available for :doc:`Installation <apache-airflow:installation/index>`. + + * :doc:`Setup a Snowflake Connection </connections/snowflake>`. + + * Create a `compute pool <https://docs.snowflake.com/en/sql-reference/sql/create-compute-pool>`__ + in your Snowflake account. + + * Push a container image to a + `Snowflake image repository <https://docs.snowflake.com/en/developer-guide/snowpark-container-services/working-with-registry-repository>`__ + and prepare a + `service specification <https://docs.snowflake.com/en/developer-guide/snowpark-container-services/specification-reference>`__ + uploaded to a stage. Review Comment: Good call, I simplified it. ########## providers/snowflake/src/airflow/providers/snowflake/operators/snowpark_containers.py: ########## @@ -0,0 +1,215 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +from __future__ import annotations + +import time +from collections.abc import Sequence +from functools import cached_property +from typing import TYPE_CHECKING, Any + +from airflow.providers.common.compat.standard.operators import BaseOperator +from airflow.providers.common.sql.hooks.handlers import fetch_one_handler +from airflow.providers.snowflake.hooks.snowflake import SnowflakeHook + +if TYPE_CHECKING: + from airflow.providers.common.compat.sdk import Context + + +class SnowparkContainerJobOperator(BaseOperator): + """ + Execute a job on Snowpark Container Services. + + Submits a container job to a compute pool via ``EXECUTE JOB SERVICE``, + optionally polls for completion, retrieves container logs, and + drops the job service on success. + + .. seealso:: + `Snowpark Container Services <https://docs.snowflake.com/en/developer-guide/snowpark-container-services/overview>`_ + + :param compute_pool: name of the compute pool to run the job on + :param container_name: container name as defined in the service specification file, + used for retrieving container logs + :param spec: spec filename on the stage (e.g. ``'spec.yaml'``). + Must be provided together with ``spec_stage`` + :param spec_stage: stage where the spec file is stored (e.g. ``'@my_stage'``). + Must be provided together with ``spec`` + :param spec_text: inline YAML spec text, as an alternative to ``spec``/``spec_stage``. + The text is wrapped in ``$$`` delimiters automatically + :param name: (Optional) job service name. If not provided, Snowflake + auto-generates a name + :param query_warehouse: (Optional) warehouse for SQL queries run inside the container. + This is separate from the ``warehouse`` parameter used by the operator's + own SQL commands + :param replicas: (Optional) number of job replicas to run. (default value: 1) + :param wait_for_completion: poll until the job reaches a terminal state. + When disabled, the job is submitted and the operator returns + immediately. (default value: True) + :param drop_on_completion: drop the job service after the job finishes + successfully. Failed jobs are not dropped, allowing inspection + in Snowflake. (default value: True) + :param poll_interval: the interval in seconds to poll the query status. + (default value: 10) + :param snowflake_conn_id: Reference to + :ref:`Snowflake connection id<howto/connection:snowflake>` + :param database: name of database (will overwrite database defined + in connection) + :param schema: name of schema (will overwrite schema defined in + connection) + :param role: name of role (will overwrite any role defined in + connection's extra JSON) + :param warehouse: name of warehouse (will overwrite any warehouse + defined in the connection's extra JSON). Used for the operator's + own SQL commands, not for the container's queries + """ + + template_fields: Sequence[str] = ( + "compute_pool", + "spec", + "spec_stage", + "container_name", + "spec_text", + "name", + "query_warehouse", + "snowflake_conn_id", + ) + + def __init__( + self, + *, + compute_pool: str, + container_name: str, + spec: str | None = None, + spec_stage: str | None = None, + spec_text: str | None = None, + name: str | None = None, + query_warehouse: str | None = None, + replicas: int = 1, + wait_for_completion: bool = True, + drop_on_completion: bool = True, + poll_interval: int = 10, + snowflake_conn_id: str = "snowflake_default", + database: str | None = None, + schema: str | None = None, + role: str | None = None, + warehouse: str | None = None, + **kwargs: Any, + ) -> None: + super().__init__(**kwargs) + self.compute_pool = compute_pool + self.container_name = container_name + self.spec = spec + self.spec_stage = spec_stage + self.spec_text = spec_text + self.name = name + self.query_warehouse = query_warehouse + self.replicas = replicas + self.wait_for_completion = wait_for_completion + self.drop_on_completion = drop_on_completion + self.poll_interval = poll_interval + self.snowflake_conn_id = snowflake_conn_id + self.database = database + self.schema = schema + self.role = role + self.warehouse = warehouse + self.job_name: str | None = None Review Comment: Added a comment. ########## providers/snowflake/src/airflow/providers/snowflake/operators/snowpark_containers.py: ########## @@ -0,0 +1,215 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +from __future__ import annotations + +import time +from collections.abc import Sequence +from functools import cached_property +from typing import TYPE_CHECKING, Any + +from airflow.providers.common.compat.standard.operators import BaseOperator +from airflow.providers.common.sql.hooks.handlers import fetch_one_handler +from airflow.providers.snowflake.hooks.snowflake import SnowflakeHook + +if TYPE_CHECKING: + from airflow.providers.common.compat.sdk import Context + + +class SnowparkContainerJobOperator(BaseOperator): + """ + Execute a job on Snowpark Container Services. + + Submits a container job to a compute pool via ``EXECUTE JOB SERVICE``, + optionally polls for completion, retrieves container logs, and + drops the job service on success. + + .. seealso:: + `Snowpark Container Services <https://docs.snowflake.com/en/developer-guide/snowpark-container-services/overview>`_ + + :param compute_pool: name of the compute pool to run the job on + :param container_name: container name as defined in the service specification file, + used for retrieving container logs + :param spec: spec filename on the stage (e.g. ``'spec.yaml'``). + Must be provided together with ``spec_stage`` + :param spec_stage: stage where the spec file is stored (e.g. ``'@my_stage'``). + Must be provided together with ``spec`` + :param spec_text: inline YAML spec text, as an alternative to ``spec``/``spec_stage``. + The text is wrapped in ``$$`` delimiters automatically + :param name: (Optional) job service name. If not provided, Snowflake + auto-generates a name + :param query_warehouse: (Optional) warehouse for SQL queries run inside the container. + This is separate from the ``warehouse`` parameter used by the operator's + own SQL commands + :param replicas: (Optional) number of job replicas to run. (default value: 1) + :param wait_for_completion: poll until the job reaches a terminal state. + When disabled, the job is submitted and the operator returns + immediately. (default value: True) + :param drop_on_completion: drop the job service after the job finishes + successfully. Failed jobs are not dropped, allowing inspection + in Snowflake. (default value: True) + :param poll_interval: the interval in seconds to poll the query status. + (default value: 10) + :param snowflake_conn_id: Reference to + :ref:`Snowflake connection id<howto/connection:snowflake>` + :param database: name of database (will overwrite database defined + in connection) + :param schema: name of schema (will overwrite schema defined in + connection) + :param role: name of role (will overwrite any role defined in + connection's extra JSON) + :param warehouse: name of warehouse (will overwrite any warehouse + defined in the connection's extra JSON). Used for the operator's + own SQL commands, not for the container's queries + """ + + template_fields: Sequence[str] = ( + "compute_pool", + "spec", + "spec_stage", + "container_name", + "spec_text", + "name", + "query_warehouse", + "snowflake_conn_id", + ) + + def __init__( + self, + *, + compute_pool: str, + container_name: str, + spec: str | None = None, + spec_stage: str | None = None, + spec_text: str | None = None, + name: str | None = None, + query_warehouse: str | None = None, + replicas: int = 1, + wait_for_completion: bool = True, + drop_on_completion: bool = True, + poll_interval: int = 10, + snowflake_conn_id: str = "snowflake_default", + database: str | None = None, + schema: str | None = None, + role: str | None = None, + warehouse: str | None = None, + **kwargs: Any, + ) -> None: + super().__init__(**kwargs) + self.compute_pool = compute_pool + self.container_name = container_name + self.spec = spec + self.spec_stage = spec_stage + self.spec_text = spec_text + self.name = name + self.query_warehouse = query_warehouse + self.replicas = replicas + self.wait_for_completion = wait_for_completion + self.drop_on_completion = drop_on_completion + self.poll_interval = poll_interval + self.snowflake_conn_id = snowflake_conn_id + self.database = database + self.schema = schema + self.role = role + self.warehouse = warehouse + self.job_name: str | None = None + + if self.spec_text and (self.spec or self.spec_stage): + raise ValueError("Cannot specify both 'spec_text' and 'spec'/'spec_stage'") + if not self.spec_text and not (self.spec and self.spec_stage): + raise ValueError("Must provide either 'spec_text' or both 'spec' and 'spec_stage'") + + @cached_property + def _hook(self) -> SnowflakeHook: + return SnowflakeHook( + snowflake_conn_id=self.snowflake_conn_id, + warehouse=self.warehouse, + database=self.database, + schema=self.schema, + role=self.role, + ) + + def _build_sql(self) -> str: + """Build the execute job SQL statement.""" + sql = f"EXECUTE JOB SERVICE IN COMPUTE POOL {self.compute_pool}" + if self.name: + sql += f" NAME = {self.name}" + sql += " ASYNC = TRUE" + if self.replicas > 1: + sql += f" REPLICAS = {self.replicas}" + if self.query_warehouse: + sql += f" QUERY_WAREHOUSE = {self.query_warehouse}" + if self.spec_text: + sql += f" FROM SPECIFICATION $${self.spec_text}$$" + else: + sql += f" FROM {self.spec_stage} SPEC = '{self.spec}'" + return sql + + def _run_one(self, sql: str, return_dictionaries: bool = False) -> Any: + """Run a single statement that returns one row via fetch_one_handler.""" + return self._hook.run(sql, handler=fetch_one_handler, return_dictionaries=return_dictionaries) + + def _submit_job(self) -> None: + """Submit the job and store the job name on the instance.""" + response = self._run_one(self._build_sql()) + self.job_name = response[0].split("'")[1] Review Comment: To answer your first question - the other option I found was querying the job history (https://docs.snowflake.com/en/sql-reference/functions/get_job_history). We could use the query id that gets stored in the hook, apply a created_time_start filter (to get the most recent jobs), and then fetch the name from the result of that query (directly from the name column). I felt the extra layer of complexity wasn't really worth it given we can get it directly from the initial response. If there are any errors with the query itself, the hook would raise and the user would be aware. The expectation is that if the job does submit correctly, the job name should always be present in the response. To your second point - you're correct, it's more clear inside the execute method. I made the change. -- 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]
