# 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 datetime import timedelta
from functools import cached_property
from typing import TYPE_CHECKING, Any
from airflow.providers.common.compat.sdk import conf
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
from airflow.providers.snowflake.triggers.snowpark_containers import SnowparkContainerJobTrigger
from airflow.providers.snowflake.utils.snowpark_containers import (
NON_TERMINAL_STATUSES,
TERMINAL_STATUSES,
SnowparkContainerJobStatus,
)
if TYPE_CHECKING:
from airflow.providers.common.compat.sdk import Context
[docs]
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 external_access_integrations: (Optional) Names of the external access
integrations that allow your job to access external sites. Names are
case-sensitive (default value: None)
: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 or on a timeout. 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 deferrable: Run the operator in deferrable mode. Only effective when
``wait_for_completion`` is True. With ``wait_for_completion=False`` the
operator submits the job and returns immediately without deferring.
(default value: False)
:param timeout: Maximum seconds to wait for the job to reach a terminal
state. When it elapses the task fails. (default value: 86400)
: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
"""
[docs]
template_fields: Sequence[str] = (
"compute_pool",
"spec",
"spec_stage",
"container_name",
"spec_text",
"name",
"query_warehouse",
"snowflake_conn_id",
"external_access_integrations",
)
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,
external_access_integrations: list[str] | None = None,
wait_for_completion: bool = True,
drop_on_completion: bool = True,
poll_interval: int = 10,
snowflake_conn_id: str = "snowflake_default",
deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
timeout: int = 24 * 60 * 60,
database: str | None = None,
schema: str | None = None,
role: str | None = None,
warehouse: str | None = None,
**kwargs: Any,
) -> None:
super().__init__(**kwargs)
if spec_text is not None and (spec is not None or spec_stage is not None):
raise ValueError("Cannot specify both 'spec_text' and 'spec'/'spec_stage'")
[docs]
self.compute_pool = compute_pool
[docs]
self.container_name = container_name
[docs]
self.spec_stage = spec_stage
[docs]
self.spec_text = spec_text
[docs]
self.query_warehouse = query_warehouse
[docs]
self.replicas = replicas
[docs]
self.external_access_integrations = external_access_integrations
[docs]
self.wait_for_completion = wait_for_completion
[docs]
self.drop_on_completion = drop_on_completion
[docs]
self.poll_interval = poll_interval
[docs]
self.snowflake_conn_id = snowflake_conn_id
[docs]
self.deferrable = deferrable
[docs]
self.database = database
[docs]
self.warehouse = warehouse
# Set after the job is submitted, parsed from the job submission response.
[docs]
self.job_name: str | None = None
if self.deferrable and not self.wait_for_completion:
self.log.warning("deferrable has no effect when wait_for_completion is False.")
@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.external_access_integrations:
eais = ", ".join(self.external_access_integrations)
sql += f" EXTERNAL_ACCESS_INTEGRATIONS = ({eais})"
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."""
status = None
end_time = time.monotonic() + self.timeout
while True:
if time.monotonic() >= end_time:
self._log_container_output(status)
if self.drop_on_completion:
self._drop_service()
raise TimeoutError(f"Job {self.job_name} did not reach a terminal status before the timeout.")
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) -> None:
"""Fetch and log container output for all replicas. Best-effort so it never blocks cleanup."""
for instance_id in range(self.replicas):
sql = f"SELECT SYSTEM$GET_SERVICE_LOGS('{self.job_name}', {instance_id}, '{self.container_name}')"
try:
response = self._run_one(sql)[0]
except Exception as e:
self.log.warning("Could not retrieve logs for instance_id %d: %s", instance_id, e)
continue
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 _drop_service(self) -> None:
"""Best-effort drop of the job service."""
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)
[docs]
def on_kill(self) -> None:
"""Drop the running service on task kill."""
if self.job_name:
self._drop_service()
def _handle_final_status(self, status: str) -> None:
"""Log container output, fail unless the job is DONE, and optionally drop the service on success."""
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:
# Job already succeeded, so a cleanup failure is logged rather than raised
# to avoid marking a successful job as failed.
self._drop_service()
[docs]
def execute(self, context: Context) -> str:
"""Submit and optionally wait for a Snowpark Container Services job."""
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'")
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
if self.deferrable:
# timeout and execution_timeout give the trigger two separate deadlines. timeout caps
# how long the job is polled, and execution_timeout, when set, enforces the task-level
# limit. The trigger times out on whichever is reached first.
now = time.time()
poll_buffer = timedelta(seconds=self.poll_interval + 60)
execution_deadline = None
defer_timeout = timedelta(seconds=self.timeout) + poll_buffer
if self.execution_timeout is not None:
# Hand the execution deadline to the trigger so it emits a timeout event that drops the
# service. The framework's defer timeout would otherwise kill the task with no cleanup.
execution_deadline = (
context["ti"].start_date.timestamp() + self.execution_timeout.total_seconds()
)
# Pad the backstop past that deadline so the trigger fires first.
defer_timeout = self.execution_timeout + poll_buffer
self.defer(
trigger=SnowparkContainerJobTrigger(
job_name=self.job_name,
snowflake_conn_id=self.snowflake_conn_id,
poll_interval=self.poll_interval,
end_time=now + self.timeout,
execution_deadline=execution_deadline,
database=self.database,
schema=self.schema,
role=self.role,
warehouse=self.warehouse,
),
timeout=defer_timeout,
method_name="execute_complete",
)
status = self._poll_for_status()
self._handle_final_status(status)
return self.job_name
[docs]
def execute_complete(self, context: Context, event: dict[str, Any]) -> str:
"""Resume after the trigger fires."""
self.job_name = event["job_name"]
status = event["status"]
if status == "timeout":
self._log_container_output(status)
if self.drop_on_completion:
self._drop_service()
raise TimeoutError(event.get("message", f"Job '{self.job_name}' did not complete: {status}"))
if status == "error":
raise RuntimeError(event.get("message", f"Job '{self.job_name}' did not complete: {status}"))
self._handle_final_status(status)
return self.job_name