justinpakzad commented on code in PR #68259: URL: https://github.com/apache/airflow/pull/68259#discussion_r3439586881
########## 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: Done. ########## 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: Correct, the idea was just to rely on execution_timeout for the moment. As for deferrable mode, I had the same thoughts, I was going to create a follow-up PR but wanted to get the normal implementation up and running first. That is where adding operator level timeout will probably be more useful imo. ########## 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: Added. ########## 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: Added. -- 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]
