This is an automated email from the ASF dual-hosted git repository.
vincbeck pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new c8d890852cb Add SnowparkContainerJobOperator (#68259)
c8d890852cb is described below
commit c8d890852cbe73db98024df2c3d27ba3aafee9eb
Author: Justin Pakzad <[email protected]>
AuthorDate: Mon Jul 6 11:31:07 2026 -0400
Add SnowparkContainerJobOperator (#68259)
This PR adds SnowparkContainerJobOperator which:
- Submits the job asynchronously via EXECUTE JOB SERVICE
- Polls DESCRIBE SERVICE until the job reaches a terminal state (DONE,
FAILED, etc.)
- Fetches container logs via SYSTEM$GET_SERVICE_LOGS for all replicas and
surfaces them in the Airflow task logs
- Supports both stage-based specs (spec + spec_stage) and inline YAML
(spec_text)
- Supports fire-and-forget mode (wait_for_completion=False) for
long-running jobs as well as other optional parameters (drop_on_completion,
replicas, query_warehouse, etc.)
---
.../docs/operators/snowpark_containers.rst | 69 +++++
providers/snowflake/provider.yaml | 2 +
.../providers/snowflake/get_provider_info.py | 2 +
.../snowflake/operators/snowpark_containers.py | 251 ++++++++++++++++++
.../snowflake/example_snowpark_container_job.py | 54 ++++
.../operators/test_snowpark_containers.py | 282 +++++++++++++++++++++
6 files changed, 660 insertions(+)
diff --git a/providers/snowflake/docs/operators/snowpark_containers.rst
b/providers/snowflake/docs/operators/snowpark_containers.rst
new file mode 100644
index 00000000000..8677d16207c
--- /dev/null
+++ b/providers/snowflake/docs/operators/snowpark_containers.rst
@@ -0,0 +1,69 @@
+ .. 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 the necessary resources in your Snowflake account. See the
+ `Snowpark Container Services documentation
<https://docs.snowflake.com/en/developer-guide/snowpark-container-services/overview>`__.
+
+Using the Operator
+^^^^^^^^^^^^^^^^^^
+
+Use the ``snowflake_conn_id`` argument to specify the connection used. If not
+specified, ``snowflake_default`` will be used.
+
+An example usage of the SnowparkContainerJobOperator is as follows:
+
+.. exampleinclude::
/../../snowflake/tests/system/snowflake/example_snowpark_container_job.py
+ :language: python
+ :start-after: [START howto_operator_snowpark_container_job]
+ :end-before: [END howto_operator_snowpark_container_job]
+ :dedent: 4
+
+.. note::
+
+ Parameters that can be passed onto the operator will be given priority over
the parameters already given
+ in the Airflow connection metadata (such as ``schema``, ``role``,
``database`` and so forth).
diff --git a/providers/snowflake/provider.yaml
b/providers/snowflake/provider.yaml
index 42a832fc34f..af2de596afd 100644
--- a/providers/snowflake/provider.yaml
+++ b/providers/snowflake/provider.yaml
@@ -118,6 +118,7 @@ integrations:
how-to-guide:
- /docs/apache-airflow-providers-snowflake/operators/snowflake.rst
- /docs/apache-airflow-providers-snowflake/operators/snowpark.rst
+ -
/docs/apache-airflow-providers-snowflake/operators/snowpark_containers.rst
logo: /docs/integration-logos/Snowflake.png
tags: [service]
@@ -126,6 +127,7 @@ operators:
python-modules:
- airflow.providers.snowflake.operators.snowflake
- airflow.providers.snowflake.operators.snowpark
+ - airflow.providers.snowflake.operators.snowpark_containers
task-decorators:
- class-name: airflow.providers.snowflake.decorators.snowpark.snowpark_task
diff --git
a/providers/snowflake/src/airflow/providers/snowflake/get_provider_info.py
b/providers/snowflake/src/airflow/providers/snowflake/get_provider_info.py
index c6c4ffd84d4..d60e1d9eee0 100644
--- a/providers/snowflake/src/airflow/providers/snowflake/get_provider_info.py
+++ b/providers/snowflake/src/airflow/providers/snowflake/get_provider_info.py
@@ -33,6 +33,7 @@ def get_provider_info():
"how-to-guide": [
"/docs/apache-airflow-providers-snowflake/operators/snowflake.rst",
"/docs/apache-airflow-providers-snowflake/operators/snowpark.rst",
+
"/docs/apache-airflow-providers-snowflake/operators/snowpark_containers.rst",
],
"logo": "/docs/integration-logos/Snowflake.png",
"tags": ["service"],
@@ -44,6 +45,7 @@ def get_provider_info():
"python-modules": [
"airflow.providers.snowflake.operators.snowflake",
"airflow.providers.snowflake.operators.snowpark",
+
"airflow.providers.snowflake.operators.snowpark_containers",
],
}
],
diff --git
a/providers/snowflake/src/airflow/providers/snowflake/operators/snowpark_containers.py
b/providers/snowflake/src/airflow/providers/snowflake/operators/snowpark_containers.py
new file mode 100644
index 00000000000..971be06543a
--- /dev/null
+++
b/providers/snowflake/src/airflow/providers/snowflake/operators/snowpark_containers.py
@@ -0,0 +1,251 @@
+# 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 enum import Enum
+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 SnowparkContainerJobStatus(str, Enum):
+ """Statuses of a Snowpark Container Services."""
+
+ PENDING = "PENDING"
+ RUNNING = "RUNNING"
+ CANCELLING = "CANCELLING"
+ SUSPENDING = "SUSPENDING"
+ DELETING = "DELETING"
+ DONE = "DONE"
+ FAILED = "FAILED"
+ CANCELLED = "CANCELLED"
+ INTERNAL_ERROR = "INTERNAL_ERROR"
+
+
+TERMINAL_STATUSES: frozenset[SnowparkContainerJobStatus] = frozenset(
+ {
+ SnowparkContainerJobStatus.DONE,
+ SnowparkContainerJobStatus.FAILED,
+ SnowparkContainerJobStatus.CANCELLED,
+ SnowparkContainerJobStatus.INTERNAL_ERROR,
+ }
+)
+NON_TERMINAL_STATUSES: frozenset[SnowparkContainerJobStatus] = frozenset(
+ {
+ SnowparkContainerJobStatus.PENDING,
+ SnowparkContainerJobStatus.RUNNING,
+ SnowparkContainerJobStatus.CANCELLING,
+ SnowparkContainerJobStatus.SUSPENDING,
+ SnowparkContainerJobStatus.DELETING,
+ }
+)
+
+
+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
+ # Set after the job is submitted, parsed from the job submission
response.
+ 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) -> str:
+ """Submit the job and return the name."""
+ response = self._run_one(self._build_sql())
+ job_name = response[0].split("'")[1]
+ return job_name
+
+ 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 TERMINAL_STATUSES:
+ return status
+ if status not in NON_TERMINAL_STATUSES:
+ 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 != SnowparkContainerJobStatus.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.job_name = 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 != SnowparkContainerJobStatus.DONE:
+ raise RuntimeError(f"Job '{self.job_name}' finished with status:
{status}")
+ if self.drop_on_completion:
+ self._hook.run(f"DROP SERVICE IF EXISTS {self.job_name}")
+ return self.job_name
diff --git
a/providers/snowflake/tests/system/snowflake/example_snowpark_container_job.py
b/providers/snowflake/tests/system/snowflake/example_snowpark_container_job.py
new file mode 100644
index 00000000000..8c9a1a72f1f
--- /dev/null
+++
b/providers/snowflake/tests/system/snowflake/example_snowpark_container_job.py
@@ -0,0 +1,54 @@
+#
+# 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.
+"""
+Example use of SnowparkContainerJobOperator.
+"""
+
+from __future__ import annotations
+
+from datetime import datetime
+
+from airflow import DAG
+from airflow.providers.snowflake.operators.snowpark_containers import
SnowparkContainerJobOperator
+
+SNOWFLAKE_CONN_ID = "my_snowflake_conn"
+DAG_ID = "example_snowpark_container_job"
+
+with DAG(
+ DAG_ID,
+ start_date=datetime(2024, 1, 1),
+ schedule="@once",
+ default_args={"snowflake_conn_id": SNOWFLAKE_CONN_ID},
+ tags=["example"],
+ catchup=False,
+) as dag:
+ # [START howto_operator_snowpark_container_job]
+ run_job = SnowparkContainerJobOperator(
+ task_id="run_job",
+ compute_pool="my_compute_pool",
+ container_name="main",
+ spec="spec.yaml",
+ spec_stage="@my_stage",
+ )
+ # [END howto_operator_snowpark_container_job]
+
+
+from tests_common.test_utils.system_tests import get_test_run # noqa: E402
+
+# Needed to run the example DAG with pytest (see:
contributing-docs/testing/system_tests.rst)
+test_run = get_test_run(dag)
diff --git
a/providers/snowflake/tests/unit/snowflake/operators/test_snowpark_containers.py
b/providers/snowflake/tests/unit/snowflake/operators/test_snowpark_containers.py
new file mode 100644
index 00000000000..3b10bfea06d
--- /dev/null
+++
b/providers/snowflake/tests/unit/snowflake/operators/test_snowpark_containers.py
@@ -0,0 +1,282 @@
+# 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()
+ result = op._submit_job()
+ assert result == 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_submit_job_raises_on_malformed_response(self, 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()
+
+ @mock.patch.object(SnowparkContainerJobOperator, "_submit_job",
return_value=None)
+ def test_execute_raises_when_job_name_not_returned(self, mock_submit):
+ op = _make_operator()
+ with pytest.raises(RuntimeError, match="Job name was not returned"):
+ op.execute(context=None)
+
+ @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",
return_value=JOB_NAME)
+ def test_execute_wait_success(self, mock_submit, mock_poll, mock_log,
mock_hook_cls):
+ op = _make_operator()
+ 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",
return_value=JOB_NAME)
+ def test_execute_wait_failure_raises(self, mock_submit, mock_poll,
mock_log, mock_hook_cls):
+ op = _make_operator()
+ 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",
return_value=JOB_NAME)
+ @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.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",
return_value=JOB_NAME)
+ @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.execute(context=None)
+ mock_hook.run.assert_not_called()
