SameerMesiah97 commented on code in PR #68259: URL: https://github.com/apache/airflow/pull/68259#discussion_r3438611199
########## 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: I am just wondering whether lines 49-56 are going into Snowflake-platform specific implementation details. Perhaps, it would be better to just add a link to the Snowflake docs page where the steps to setup a Snowpark Container Job are described. ########## 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: nit: I would add a blank line above `self.job_name` to separate it from the arguments for the operator. You appear to be using it persist state within the operator. Maybe add a comment to make this clear, ########## 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: 1) Is there a structured way to retrieve the generated service name rather than parsing `response[0].split("'")[1]`? 2) Would it not be better for `submit_job` to return the job name rather than mutate shared state? I think you can assign the job name to `self.job_name` in the `execute` method instead so it looks less like a hidden side effect. ########## 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] + + def _poll_for_status(self) -> str: + """Poll until the job reaches a terminal state.""" + while True: + response = self._run_one(f"DESCRIBE SERVICE {self.job_name}", return_dictionaries=True) + status = response.get("status") + if status in ("DONE", "FAILED", "CANCELLED", "INTERNAL_ERROR"): + return status + if status not in ("PENDING", "RUNNING", "CANCELLING", "SUSPENDING", "DELETING"): + raise RuntimeError(f"Job {self.job_name} returned unexpected status: {status}") + time.sleep(self.poll_interval) + + def _log_container_output(self, status: str) -> None: + """Fetch and log container output for all replicas.""" + for instance_id in range(self.replicas): + sql = f"SELECT SYSTEM$GET_SERVICE_LOGS('{self.job_name}', {instance_id}, '{self.container_name}')" + response = self._run_one(sql)[0] + if not response: + continue + if status != "DONE": + self.log.error("Logs for instance_id %d:\n%s", instance_id, response) + else: + self.log.info("Logs for instance_id %d:\n%s", instance_id, response) + + def on_kill(self) -> None: + """Drop the running service on task kill.""" + if self.job_name: + try: + self._hook.run(f"DROP SERVICE IF EXISTS {self.job_name}") + except Exception as e: + self.log.error("Error dropping service %s: %s", self.job_name, e) + + def execute(self, context: Context) -> str: + """Submit and optionally wait for a Snowpark Container Services job.""" + self._submit_job() + if not self.job_name: + raise RuntimeError("Job name was not returned") + if not self.wait_for_completion: + return self.job_name + status = self._poll_for_status() + self._log_container_output(status) + if status != "DONE": Review Comment: Replace this with enum check. ########## 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, Review Comment: 1) Where is the operator-level timeout? Or are we solely going to rely on `execution_timeout` here?~ 2) This operator is a textbook candidate for implementing deferrable mode. But that can be a follow-up PR. ########## providers/snowflake/tests/unit/snowflake/operators/test_snowpark_containers.py: ########## @@ -0,0 +1,272 @@ +# 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 + +from unittest import mock + +import pytest + +from airflow.providers.snowflake.operators.snowpark_containers import SnowparkContainerJobOperator + +TASK_ID = "test_spcs_job" +COMPUTE_POOL = "test_pool" +CONTAINER_NAME = "main" +SPEC = "spec.yaml" +SPEC_STAGE = "@test_stage" +SPEC_TEXT = "spec:\n containers:\n - name: main\n image: /db/schema/repo/img:latest" +JOB_NAME = "TEST_JOB" +SNOWFLAKE_CONN_ID = "snowflake_default" +MOCK_HOOK_PATH = "airflow.providers.snowflake.operators.snowpark_containers.SnowflakeHook" + + +def _make_operator(**kwargs): + defaults = { + "task_id": TASK_ID, + "compute_pool": COMPUTE_POOL, + "container_name": CONTAINER_NAME, + "spec": SPEC, + "spec_stage": SPEC_STAGE, + } + defaults.update(kwargs) + return SnowparkContainerJobOperator(**defaults) + + +class TestSnowparkContainerJobOperator: + @pytest.mark.parametrize( + ("kwargs", "match"), + ( + pytest.param( + {"spec": None, "spec_stage": None, "spec_text": None}, + "Must provide either", + id="no_spec_provided", + ), + pytest.param( + {"spec": SPEC, "spec_stage": SPEC_STAGE, "spec_text": SPEC_TEXT}, + "Cannot specify both", + id="both_spec_and_spec_text", + ), + pytest.param( + {"spec": SPEC, "spec_stage": None, "spec_text": None}, + "Must provide either", + id="spec_without_stage", + ), + pytest.param( + {"spec": None, "spec_stage": SPEC_STAGE, "spec_text": None}, + "Must provide either", + id="stage_without_spec", + ), + ), + ) + def test_invalid_spec_combinations(self, kwargs, match): + with pytest.raises(ValueError, match=match): + _make_operator(**kwargs) + + def test_build_sql_with_spec_stage(self): + op = _make_operator() + sql = op._build_sql() + assert sql == ( + f"EXECUTE JOB SERVICE IN COMPUTE POOL {COMPUTE_POOL}" + " ASYNC = TRUE" + f" FROM {SPEC_STAGE} SPEC = '{SPEC}'" + ) + + def test_build_sql_with_spec_text(self): + op = _make_operator(spec=None, spec_stage=None, spec_text=SPEC_TEXT) + sql = op._build_sql() + assert sql == ( + f"EXECUTE JOB SERVICE IN COMPUTE POOL {COMPUTE_POOL}" + " ASYNC = TRUE" + f" FROM SPECIFICATION $${SPEC_TEXT}$$" + ) + + @pytest.mark.parametrize( + ("kwargs", "expected"), + ( + pytest.param({"name": "my_job"}, "NAME = my_job", id="name"), + pytest.param({"replicas": 5}, "REPLICAS = 5", id="replicas"), + pytest.param( + {"query_warehouse": "COMPUTE_WH"}, "QUERY_WAREHOUSE = COMPUTE_WH", id="query_warehouse" + ), + ), + ) + def test_build_sql_optional_params(self, kwargs, expected): + op = _make_operator(**kwargs) + assert expected in op._build_sql() + + @mock.patch(MOCK_HOOK_PATH) + def test_submit_job_parses_job_name(self, mock_hook_cls): + mock_hook = mock_hook_cls.return_value + mock_hook.run.return_value = ["Started Snowpark Container Services Job 'TEST_JOB'."] + op = _make_operator() + op._submit_job() + assert op.job_name == JOB_NAME Review Comment: Could you add another test for malformed job name response? I would suggest something like this: ``` @mock.patch(MOCK_HOOK_PATH) def test_submit_job_raises_on_unexpected_response(mock_hook_cls): mock_hook = mock_hook_cls.return_value mock_hook.run.return_value = ["unexpected response"] op = _make_operator() with pytest.raises(IndexError): op._submit_job() ``` ########## 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] + + def _poll_for_status(self) -> str: + """Poll until the job reaches a terminal state.""" + while True: + response = self._run_one(f"DESCRIBE SERVICE {self.job_name}", return_dictionaries=True) + status = response.get("status") + if status in ("DONE", "FAILED", "CANCELLED", "INTERNAL_ERROR"): + return status + if status not in ("PENDING", "RUNNING", "CANCELLING", "SUSPENDING", "DELETING"): + raise RuntimeError(f"Job {self.job_name} returned unexpected status: {status}") Review Comment: I think you should create an enum class in this operator file like the below: ``` class SnowparkContainerJobStatus(str, Enum): PENDING = "PENDING" RUNNING = "RUNNING" CANCELLING = "CANCELLING" SUSPENDING = "SUSPENDING" DELETING = "DELETING" DONE = "DONE" FAILED = "FAILED" CANCELLED = "CANCELLED" INTERNAL_ERROR = "INTERNAL_ERROR" NON_TERMINAL_STATUSES = ( PENDING, RUNNING, CANCELLING, SUSPENDING, DELETING, ) TERMINAL_STATUSES = ( DONE, FAILED, CANCELLED, INTERNAL_ERROR, ) ``` ########## 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] + + def _poll_for_status(self) -> str: + """Poll until the job reaches a terminal state.""" + while True: + response = self._run_one(f"DESCRIBE SERVICE {self.job_name}", return_dictionaries=True) + status = response.get("status") + if status in ("DONE", "FAILED", "CANCELLED", "INTERNAL_ERROR"): + return status + if status not in ("PENDING", "RUNNING", "CANCELLING", "SUSPENDING", "DELETING"): + raise RuntimeError(f"Job {self.job_name} returned unexpected status: {status}") + time.sleep(self.poll_interval) + + def _log_container_output(self, status: str) -> None: + """Fetch and log container output for all replicas.""" + for instance_id in range(self.replicas): + sql = f"SELECT SYSTEM$GET_SERVICE_LOGS('{self.job_name}', {instance_id}, '{self.container_name}')" + response = self._run_one(sql)[0] + if not response: + continue + if status != "DONE": + self.log.error("Logs for instance_id %d:\n%s", instance_id, response) + else: + self.log.info("Logs for instance_id %d:\n%s", instance_id, response) + + def on_kill(self) -> None: + """Drop the running service on task kill.""" + if self.job_name: + try: + self._hook.run(f"DROP SERVICE IF EXISTS {self.job_name}") + except Exception as e: + self.log.error("Error dropping service %s: %s", self.job_name, e) + + def execute(self, context: Context) -> str: + """Submit and optionally wait for a Snowpark Container Services job.""" + self._submit_job() Review Comment: Please see my coment above near `_submit_job`. I think it would be better to assign the job name to `self.job_name` here. ########## providers/snowflake/tests/unit/snowflake/operators/test_snowpark_containers.py: ########## @@ -0,0 +1,272 @@ +# 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 + +from unittest import mock + +import pytest + +from airflow.providers.snowflake.operators.snowpark_containers import SnowparkContainerJobOperator + +TASK_ID = "test_spcs_job" +COMPUTE_POOL = "test_pool" +CONTAINER_NAME = "main" +SPEC = "spec.yaml" +SPEC_STAGE = "@test_stage" +SPEC_TEXT = "spec:\n containers:\n - name: main\n image: /db/schema/repo/img:latest" +JOB_NAME = "TEST_JOB" +SNOWFLAKE_CONN_ID = "snowflake_default" +MOCK_HOOK_PATH = "airflow.providers.snowflake.operators.snowpark_containers.SnowflakeHook" + + +def _make_operator(**kwargs): + defaults = { + "task_id": TASK_ID, + "compute_pool": COMPUTE_POOL, + "container_name": CONTAINER_NAME, + "spec": SPEC, + "spec_stage": SPEC_STAGE, + } + defaults.update(kwargs) + return SnowparkContainerJobOperator(**defaults) + + +class TestSnowparkContainerJobOperator: + @pytest.mark.parametrize( + ("kwargs", "match"), + ( + pytest.param( + {"spec": None, "spec_stage": None, "spec_text": None}, + "Must provide either", + id="no_spec_provided", + ), + pytest.param( + {"spec": SPEC, "spec_stage": SPEC_STAGE, "spec_text": SPEC_TEXT}, + "Cannot specify both", + id="both_spec_and_spec_text", + ), + pytest.param( + {"spec": SPEC, "spec_stage": None, "spec_text": None}, + "Must provide either", + id="spec_without_stage", + ), + pytest.param( + {"spec": None, "spec_stage": SPEC_STAGE, "spec_text": None}, + "Must provide either", + id="stage_without_spec", + ), + ), + ) + def test_invalid_spec_combinations(self, kwargs, match): + with pytest.raises(ValueError, match=match): + _make_operator(**kwargs) + + def test_build_sql_with_spec_stage(self): + op = _make_operator() + sql = op._build_sql() + assert sql == ( + f"EXECUTE JOB SERVICE IN COMPUTE POOL {COMPUTE_POOL}" + " ASYNC = TRUE" + f" FROM {SPEC_STAGE} SPEC = '{SPEC}'" + ) + + def test_build_sql_with_spec_text(self): + op = _make_operator(spec=None, spec_stage=None, spec_text=SPEC_TEXT) + sql = op._build_sql() + assert sql == ( + f"EXECUTE JOB SERVICE IN COMPUTE POOL {COMPUTE_POOL}" + " ASYNC = TRUE" + f" FROM SPECIFICATION $${SPEC_TEXT}$$" + ) + + @pytest.mark.parametrize( + ("kwargs", "expected"), + ( + pytest.param({"name": "my_job"}, "NAME = my_job", id="name"), + pytest.param({"replicas": 5}, "REPLICAS = 5", id="replicas"), + pytest.param( + {"query_warehouse": "COMPUTE_WH"}, "QUERY_WAREHOUSE = COMPUTE_WH", id="query_warehouse" + ), + ), + ) + def test_build_sql_optional_params(self, kwargs, expected): + op = _make_operator(**kwargs) + assert expected in op._build_sql() + + @mock.patch(MOCK_HOOK_PATH) + def test_submit_job_parses_job_name(self, mock_hook_cls): + mock_hook = mock_hook_cls.return_value + mock_hook.run.return_value = ["Started Snowpark Container Services Job 'TEST_JOB'."] + op = _make_operator() + op._submit_job() + assert op.job_name == JOB_NAME + + @pytest.mark.parametrize( + "status", + ("DONE", "FAILED", "CANCELLED", "INTERNAL_ERROR"), + ) + @mock.patch(MOCK_HOOK_PATH) + def test_poll_returns_terminal_status(self, mock_hook_cls, status): + mock_hook = mock_hook_cls.return_value + mock_hook.run.return_value = {"status": status} + op = _make_operator(poll_interval=0) + op.job_name = JOB_NAME + assert op._poll_for_status() == status + + @mock.patch(MOCK_HOOK_PATH) + def test_poll_raises_on_unexpected_status(self, mock_hook_cls): + mock_hook = mock_hook_cls.return_value + mock_hook.run.return_value = {"status": "UNKNOWN"} + op = _make_operator(poll_interval=0) + op.job_name = JOB_NAME + with pytest.raises(RuntimeError, match="unexpected status"): + op._poll_for_status() + + @mock.patch("time.sleep") + @mock.patch(MOCK_HOOK_PATH) + def test_poll_waits_through_pending_then_done(self, mock_hook_cls, mock_sleep): + mock_hook = mock_hook_cls.return_value + mock_hook.run.side_effect = [ + {"status": "PENDING"}, + {"status": "RUNNING"}, + {"status": "DONE"}, + ] + op = _make_operator(poll_interval=5) + op.job_name = JOB_NAME + assert op._poll_for_status() == "DONE" + assert mock_sleep.call_count == 2 + + @mock.patch(MOCK_HOOK_PATH) + def test_log_container_output_uses_info_on_done(self, mock_hook_cls): + mock_hook = mock_hook_cls.return_value + mock_hook.run.return_value = ["container output here"] + op = _make_operator() + op.job_name = JOB_NAME + with mock.patch.object(op.log, "info") as mock_info: + op._log_container_output("DONE") + mock_info.assert_called_once_with("Logs for instance_id %d:\n%s", 0, "container output here") + + @mock.patch(MOCK_HOOK_PATH) + def test_log_container_output_uses_error_on_failure(self, mock_hook_cls): + mock_hook = mock_hook_cls.return_value + mock_hook.run.return_value = ["container output here"] + op = _make_operator() + op.job_name = JOB_NAME + with mock.patch.object(op.log, "error") as mock_error: + op._log_container_output("FAILED") + mock_error.assert_called_once_with("Logs for instance_id %d:\n%s", 0, "container output here") + + @mock.patch(MOCK_HOOK_PATH) + def test_log_container_output_no_logs_skips(self, mock_hook_cls): + mock_hook = mock_hook_cls.return_value + mock_hook.run.return_value = [""] + op = _make_operator() + op.job_name = JOB_NAME + with mock.patch.object(op.log, "info") as mock_info, mock.patch.object(op.log, "error") as mock_error: + op._log_container_output("DONE") + mock_info.assert_not_called() + mock_error.assert_not_called() + + @mock.patch(MOCK_HOOK_PATH) + def test_log_container_output_multiple_replicas(self, mock_hook_cls): + mock_hook = mock_hook_cls.return_value + mock_hook.run.return_value = ["logs"] + op = _make_operator(replicas=3) + op.job_name = JOB_NAME + with mock.patch.object(op.log, "info") as mock_info: + op._log_container_output("DONE") + assert mock_info.call_count == 3 + + @mock.patch(MOCK_HOOK_PATH) + def test_on_kill_no_job_name(self, mock_hook_cls): + mock_hook = mock_hook_cls.return_value + op = _make_operator() + op.on_kill() + mock_hook.run.assert_not_called() + + @mock.patch(MOCK_HOOK_PATH) + def test_on_kill_drops_service(self, mock_hook_cls): + mock_hook = mock_hook_cls.return_value + op = _make_operator() + op.job_name = JOB_NAME + op.on_kill() + mock_hook.run.assert_called_once_with(f"DROP SERVICE IF EXISTS {JOB_NAME}") + + @mock.patch(MOCK_HOOK_PATH) + def test_on_kill_logs_error_on_exception(self, mock_hook_cls): + mock_hook = mock_hook_cls.return_value + mock_hook.run.side_effect = Exception("drop failed") + op = _make_operator() + op.job_name = JOB_NAME + with mock.patch.object(op.log, "error") as mock_error: + op.on_kill() + mock_error.assert_called_once() + + @mock.patch(MOCK_HOOK_PATH) + def test_execute_no_wait(self, mock_hook_cls): + mock_hook = mock_hook_cls.return_value + mock_hook.run.return_value = ["Started Snowpark Container Services Job 'TEST_JOB'."] + op = _make_operator(wait_for_completion=False) + result = op.execute(context=None) + assert result == JOB_NAME + assert mock_hook.run.call_count == 1 + + @mock.patch(MOCK_HOOK_PATH) + @mock.patch.object(SnowparkContainerJobOperator, "_log_container_output") + @mock.patch.object(SnowparkContainerJobOperator, "_poll_for_status", return_value="DONE") + @mock.patch.object(SnowparkContainerJobOperator, "_submit_job") + def test_execute_wait_success(self, mock_submit, mock_poll, mock_log, mock_hook_cls): + op = _make_operator() + op.job_name = JOB_NAME + result = op.execute(context=None) + mock_submit.assert_called_once() + mock_poll.assert_called_once() + mock_log.assert_called_once_with("DONE") + assert result == JOB_NAME + + @mock.patch(MOCK_HOOK_PATH) + @mock.patch.object(SnowparkContainerJobOperator, "_log_container_output") + @mock.patch.object(SnowparkContainerJobOperator, "_poll_for_status", return_value="FAILED") + @mock.patch.object(SnowparkContainerJobOperator, "_submit_job") + def test_execute_wait_failure_raises(self, mock_submit, mock_poll, mock_log, mock_hook_cls): + op = _make_operator() + op.job_name = JOB_NAME + with pytest.raises(RuntimeError, match="FAILED"): + op.execute(context=None) + mock_log.assert_called_once_with("FAILED") + + @mock.patch.object(SnowparkContainerJobOperator, "_log_container_output") + @mock.patch.object(SnowparkContainerJobOperator, "_poll_for_status", return_value="DONE") + @mock.patch.object(SnowparkContainerJobOperator, "_submit_job") + @mock.patch(MOCK_HOOK_PATH) + def test_execute_drops_service_on_completion(self, mock_hook_cls, mock_submit, mock_poll, mock_log): + mock_hook = mock_hook_cls.return_value + op = _make_operator(drop_on_completion=True) + op.job_name = JOB_NAME + op.execute(context=None) + mock_hook.run.assert_called_once_with(f"DROP SERVICE IF EXISTS {JOB_NAME}") + + @mock.patch.object(SnowparkContainerJobOperator, "_log_container_output") + @mock.patch.object(SnowparkContainerJobOperator, "_poll_for_status", return_value="DONE") + @mock.patch.object(SnowparkContainerJobOperator, "_submit_job") + @mock.patch(MOCK_HOOK_PATH) + def test_execute_skips_drop_when_disabled(self, mock_hook_cls, mock_submit, mock_poll, mock_log): + mock_hook = mock_hook_cls.return_value + op = _make_operator(drop_on_completion=False) + op.job_name = JOB_NAME + op.execute(context=None) + mock_hook.run.assert_not_called() Review Comment: Non-blocking: I would add this test too: ``` @mock.patch.object(SnowparkContainerJobOperator, "_submit_job") def test_execute_raises_when_job_name_not_returned(mock_submit): op = _make_operator() with pytest.raises(RuntimeError, match="Job name was not returned"): op.execute(context=None) ``` Realistically, `self.job_name` is unlikely to remain `None` during the execution flow but this behavior is part of the contract of the operator. -- 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]
