You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by "justinpakzad (via GitHub)" <gi...@apache.org> on 2026/06/09 00:54:28 UTC
[PR] Add SnowparkContainerJobOperator [airflow]
justinpakzad opened a new pull request, #68259:
URL: https://github.com/apache/airflow/pull/68259
<!-- SPDX-License-Identifier: Apache-2.0
https://www.apache.org/licenses/LICENSE-2.0 -->
<!--
Thank you for contributing!
Please provide above a brief description of the changes made in this pull request.
Write a good git commit message following this guide: http://chris.beams.io/posts/git-commit/
Please make sure that your code changes are covered with tests.
And in case of new features or big changes remember to adjust the documentation.
Feel free to ping (in general) for the review if you do not see reaction for a few days
(72 Hours is the minimum reaction time you can expect from volunteers) - we sometimes miss notifications.
In case of an existing issue, reference it using one of the following:
* closes: #ISSUE
* related: #ISSUE
-->
---
**Description**
Snowpark Container Services lets users run containerized workloads directly within their Snowflake account. Things like ML training jobs, data processing pipelines, or any custom application packaged as a Docker image. Jobs are submitted via EXECUTE JOB SERVICE, which runs containers on a compute pool and terminates once the code exits. Currently there's no operator for orchestrating these jobs. Users would have to fall back to raw SQL via SQLExecuteQueryOperator and manually handle polling, log retrieval, error handling, and service teardown.
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.)
**Testing:**
- Unit tests covering parameter validation, SQL generation, polling logic, log retrieval, error handling for unexpected statuses, on_kill cleanup, and execution paths.
- Manually tested end-to-end against a real Snowflake account with an active compute pool, image repository, and pushed container image verified the success path (job completes, logs retrieved) and failure scenarios (job fails, error logs surfaced, operator raises appropriately).
Added an example DAG and operator documentation. Running the example DAG requires SPCS infrastructure to be provisioned beforehand (compute pool, image repository, pushed container image, etc.) those steps are covered in the operator's how-to guide.
##### Was generative AI tooling used to co-author this PR?
<!--
If generative AI tooling has been used in the process of authoring this PR, please
change below checkbox to `[X]` followed by the name of the tool, uncomment the "Generated-by".
-->
- [x] Yes (Claude Opus 4.6)
<!--
Generated-by: [Tool Name] following [the guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#gen-ai-assisted-contributions)
-->
---
* Read the **[Pull Request Guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#pull-request-guidelines)** for more information. Note: commit author/co-author name and email in commits become permanently public when merged.
* For fundamental code changes, an Airflow Improvement Proposal ([AIP](https://cwiki.apache.org/confluence/display/AIRFLOW/Airflow+Improvement+Proposals)) is needed.
* When adding dependency, check compliance with the [ASF 3rd Party License Policy](https://www.apache.org/legal/resolved.html#category-x).
* For significant user-facing changes create newsfragment: `{pr_number}.significant.rst`, in [airflow-core/newsfragments](https://github.com/apache/airflow/tree/main/airflow-core/newsfragments). You can add this file in a follow-up commit after the PR is created so you know the PR number.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add SnowparkContainerJobOperator [airflow]
Posted by "justinpakzad (via GitHub)" <gi...@apache.org>.
justinpakzad commented on PR #68259:
URL: https://github.com/apache/airflow/pull/68259#issuecomment-4747185477
> Good PR. I have left some comments.
I appreciate the detailed review. I addressed all of your feedback. I left one comment unresolved so you can see my response to your first question but your second point was addressed.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add SnowparkContainerJobOperator [airflow]
Posted by "SameerMesiah97 (via GitHub)" <gi...@apache.org>.
SameerMesiah97 commented on code in PR #68259:
URL: https://github.com/apache/airflow/pull/68259#discussion_r3444376129
##########
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:
@justinpakzad
You can resolve this.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add SnowparkContainerJobOperator [airflow]
Posted by "justinpakzad (via GitHub)" <gi...@apache.org>.
justinpakzad commented on PR #68259:
URL: https://github.com/apache/airflow/pull/68259#issuecomment-4699764095
Here are a couple of screenshots that show a successful run and a failed run against my Snowflake instance. Happy to provide more if needed.
<img width="2030" height="1061" alt="Screenshot 2026-06-13 at 16 50 05" src="https://github.com/user-attachments/assets/9da3d0ce-23c7-4912-a8d9-657c8740e0d6" />
<img width="2031" height="942" alt="Screenshot 2026-06-13 at 16 51 42" src="https://github.com/user-attachments/assets/dc0c60f3-77fe-4b37-bff8-90dfb78032d9" />
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add SnowparkContainerJobOperator [airflow]
Posted by "SameerMesiah97 (via GitHub)" <gi...@apache.org>.
SameerMesiah97 commented on code in PR #68259:
URL: https://github.com/apache/airflow/pull/68259#discussion_r3444376129
##########
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:
@justinpakzad
I have no further concerns. You can resolve this thread now.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add SnowparkContainerJobOperator [airflow]
Posted by "vincbeck (via GitHub)" <gi...@apache.org>.
vincbeck merged PR #68259:
URL: https://github.com/apache/airflow/pull/68259
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add SnowparkContainerJobOperator [airflow]
Posted by "justinpakzad (via GitHub)" <gi...@apache.org>.
justinpakzad commented on PR #68259:
URL: https://github.com/apache/airflow/pull/68259#issuecomment-4680944446
Could I get a re-trigger of the failed check please? Seems there was an error building the image. Thanks.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add SnowparkContainerJobOperator [airflow]
Posted by "justinpakzad (via GitHub)" <gi...@apache.org>.
justinpakzad commented on PR #68259:
URL: https://github.com/apache/airflow/pull/68259#issuecomment-4860276009
@potiuk would you mind taking a look at this when you have a moment. Thanks.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add SnowparkContainerJobOperator [airflow]
Posted by "justinpakzad (via GitHub)" <gi...@apache.org>.
justinpakzad commented on code in PR #68259:
URL: https://github.com/apache/airflow/pull/68259#discussion_r3439586348
##########
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:
Yea, this feels like a text book use case for an enum. Don't know why I didn't create one. Added. Just a note on the approach - I declared the groups of statuses as frozen sets. There isn't a well established pattern on how enums are typically used across the providers (at least not that I observed) and this approach made sense to me.
##########
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:
Yea, made that change.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add SnowparkContainerJobOperator [airflow]
Posted by "justinpakzad (via GitHub)" <gi...@apache.org>.
justinpakzad commented on PR #68259:
URL: https://github.com/apache/airflow/pull/68259#issuecomment-4813399698
Would appreciate a review on this. Thanks.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add SnowparkContainerJobOperator [airflow]
Posted by "justinpakzad (via GitHub)" <gi...@apache.org>.
justinpakzad commented on code in PR #68259:
URL: https://github.com/apache/airflow/pull/68259#discussion_r3439586162
##########
providers/snowflake/src/airflow/providers/snowflake/operators/snowpark_containers.py:
##########
@@ -0,0 +1,215 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from __future__ import annotations
+
+import time
+from collections.abc import Sequence
+from functools import cached_property
+from typing import TYPE_CHECKING, Any
+
+from airflow.providers.common.compat.standard.operators import BaseOperator
+from airflow.providers.common.sql.hooks.handlers import fetch_one_handler
+from airflow.providers.snowflake.hooks.snowflake import SnowflakeHook
+
+if TYPE_CHECKING:
+ from airflow.providers.common.compat.sdk import Context
+
+
+class SnowparkContainerJobOperator(BaseOperator):
+ """
+ Execute a job on Snowpark Container Services.
+
+ Submits a container job to a compute pool via ``EXECUTE JOB SERVICE``,
+ optionally polls for completion, retrieves container logs, and
+ drops the job service on success.
+
+ .. seealso::
+ `Snowpark Container Services <https://docs.snowflake.com/en/developer-guide/snowpark-container-services/overview>`_
+
+ :param compute_pool: name of the compute pool to run the job on
+ :param container_name: container name as defined in the service specification file,
+ used for retrieving container logs
+ :param spec: spec filename on the stage (e.g. ``'spec.yaml'``).
+ Must be provided together with ``spec_stage``
+ :param spec_stage: stage where the spec file is stored (e.g. ``'@my_stage'``).
+ Must be provided together with ``spec``
+ :param spec_text: inline YAML spec text, as an alternative to ``spec``/``spec_stage``.
+ The text is wrapped in ``$$`` delimiters automatically
+ :param name: (Optional) job service name. If not provided, Snowflake
+ auto-generates a name
+ :param query_warehouse: (Optional) warehouse for SQL queries run inside the container.
+ This is separate from the ``warehouse`` parameter used by the operator's
+ own SQL commands
+ :param replicas: (Optional) number of job replicas to run. (default value: 1)
+ :param wait_for_completion: poll until the job reaches a terminal state.
+ When disabled, the job is submitted and the operator returns
+ immediately. (default value: True)
+ :param drop_on_completion: drop the job service after the job finishes
+ successfully. Failed jobs are not dropped, allowing inspection
+ in Snowflake. (default value: True)
+ :param poll_interval: the interval in seconds to poll the query status.
+ (default value: 10)
+ :param snowflake_conn_id: Reference to
+ :ref:`Snowflake connection id<howto/connection:snowflake>`
+ :param database: name of database (will overwrite database defined
+ in connection)
+ :param schema: name of schema (will overwrite schema defined in
+ connection)
+ :param role: name of role (will overwrite any role defined in
+ connection's extra JSON)
+ :param warehouse: name of warehouse (will overwrite any warehouse
+ defined in the connection's extra JSON). Used for the operator's
+ own SQL commands, not for the container's queries
+ """
+
+ template_fields: Sequence[str] = (
+ "compute_pool",
+ "spec",
+ "spec_stage",
+ "container_name",
+ "spec_text",
+ "name",
+ "query_warehouse",
+ "snowflake_conn_id",
+ )
+
+ def __init__(
+ self,
+ *,
+ compute_pool: str,
+ container_name: str,
+ spec: str | None = None,
+ spec_stage: str | None = None,
+ spec_text: str | None = None,
+ name: str | None = None,
+ query_warehouse: str | None = None,
+ replicas: int = 1,
+ wait_for_completion: bool = True,
+ drop_on_completion: bool = True,
+ poll_interval: int = 10,
+ snowflake_conn_id: str = "snowflake_default",
+ database: str | None = None,
+ schema: str | None = None,
+ role: str | None = None,
+ warehouse: str | None = None,
+ **kwargs: Any,
+ ) -> None:
+ super().__init__(**kwargs)
+ self.compute_pool = compute_pool
+ self.container_name = container_name
+ self.spec = spec
+ self.spec_stage = spec_stage
+ self.spec_text = spec_text
+ self.name = name
+ self.query_warehouse = query_warehouse
+ self.replicas = replicas
+ self.wait_for_completion = wait_for_completion
+ self.drop_on_completion = drop_on_completion
+ self.poll_interval = poll_interval
+ self.snowflake_conn_id = snowflake_conn_id
+ self.database = database
+ self.schema = schema
+ self.role = role
+ self.warehouse = warehouse
+ self.job_name: str | None = None
+
+ if self.spec_text and (self.spec or self.spec_stage):
+ raise ValueError("Cannot specify both 'spec_text' and 'spec'/'spec_stage'")
+ if not self.spec_text and not (self.spec and self.spec_stage):
+ raise ValueError("Must provide either 'spec_text' or both 'spec' and 'spec_stage'")
+
+ @cached_property
+ def _hook(self) -> SnowflakeHook:
+ return SnowflakeHook(
+ snowflake_conn_id=self.snowflake_conn_id,
+ warehouse=self.warehouse,
+ database=self.database,
+ schema=self.schema,
+ role=self.role,
+ )
+
+ def _build_sql(self) -> str:
+ """Build the execute job SQL statement."""
+ sql = f"EXECUTE JOB SERVICE IN COMPUTE POOL {self.compute_pool}"
+ if self.name:
+ sql += f" NAME = {self.name}"
+ sql += " ASYNC = TRUE"
+ if self.replicas > 1:
+ sql += f" REPLICAS = {self.replicas}"
+ if self.query_warehouse:
+ sql += f" QUERY_WAREHOUSE = {self.query_warehouse}"
+ if self.spec_text:
+ sql += f" FROM SPECIFICATION $${self.spec_text}$$"
+ else:
+ sql += f" FROM {self.spec_stage} SPEC = '{self.spec}'"
+ return sql
+
+ def _run_one(self, sql: str, return_dictionaries: bool = False) -> Any:
+ """Run a single statement that returns one row via fetch_one_handler."""
+ return self._hook.run(sql, handler=fetch_one_handler, return_dictionaries=return_dictionaries)
+
+ def _submit_job(self) -> None:
+ """Submit the job and store the job name on the instance."""
+ response = self._run_one(self._build_sql())
+ self.job_name = response[0].split("'")[1]
Review Comment:
To answer your first question - the other option I found was querying the job history (https://docs.snowflake.com/en/sql-reference/functions/get_job_history). We could use the query id that gets stored in the hook, apply a created_time_start filter (to get the most recent jobs), and then fetch the name from the result of that query (directly from the name column). That being said, there is no guarantee of immediate consistency in the job history table and it could require additional polling. I felt the extra layer of complexity wasn't really worth it given we can get it directly from the initial response.
If there are any errors with the query itself, the hook would raise and the user would be aware. The expectation is that if the job does submit correctly, the job name should always be present in the response.
To your second point - you're correct, it's more clear inside the execute method. I made the change.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add SnowparkContainerJobOperator [airflow]
Posted by "SameerMesiah97 (via GitHub)" <gi...@apache.org>.
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add SnowparkContainerJobOperator [airflow]
Posted by "justinpakzad (via GitHub)" <gi...@apache.org>.
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add SnowparkContainerJobOperator [airflow]
Posted by "justinpakzad (via GitHub)" <gi...@apache.org>.
justinpakzad commented on code in PR #68259:
URL: https://github.com/apache/airflow/pull/68259#discussion_r3439585795
##########
providers/snowflake/docs/operators/snowpark_containers.rst:
##########
@@ -0,0 +1,75 @@
+ .. Licensed to the Apache Software Foundation (ASF) under one
+ or more contributor license agreements. See the NOTICE file
+ distributed with this work for additional information
+ regarding copyright ownership. The ASF licenses this file
+ to you under the Apache License, Version 2.0 (the
+ "License"); you may not use this file except in compliance
+ with the License. You may obtain a copy of the License at
+
+ .. http://www.apache.org/licenses/LICENSE-2.0
+
+ .. Unless required by applicable law or agreed to in writing,
+ software distributed under the License is distributed on an
+ "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ KIND, either express or implied. See the License for the
+ specific language governing permissions and limitations
+ under the License.
+
+.. _howto/operator:SnowparkContainerJobOperator:
+
+SnowparkContainerJobOperator
+============================
+
+Use the :class:`SnowparkContainerJobOperator
+<airflow.providers.snowflake.operators.snowpark_containers.SnowparkContainerJobOperator>` to
+submit and monitor container jobs on
+`Snowpark Container Services <https://docs.snowflake.com/en/developer-guide/snowpark-container-services/overview>`__.
+
+The operator wraps the Snowflake
+`EXECUTE JOB SERVICE <https://docs.snowflake.com/en/sql-reference/sql/execute-job-service>`__
+SQL command. It submits the job asynchronously, optionally polls until the job
+reaches a terminal state, retrieves container logs, and can drop the service on
+completion.
+
+Prerequisite Tasks
+^^^^^^^^^^^^^^^^^^
+
+To use this operator, you must do a few things:
+
+ * Install the provider package via **pip**.
+
+ .. code-block:: bash
+
+ pip install 'apache-airflow-providers-snowflake'
+
+ Detailed information is available for :doc:`Installation <apache-airflow:installation/index>`.
+
+ * :doc:`Setup a Snowflake Connection </connections/snowflake>`.
+
+ * Create a `compute pool <https://docs.snowflake.com/en/sql-reference/sql/create-compute-pool>`__
+ in your Snowflake account.
+
+ * Push a container image to a
+ `Snowflake image repository <https://docs.snowflake.com/en/developer-guide/snowpark-container-services/working-with-registry-repository>`__
+ and prepare a
+ `service specification <https://docs.snowflake.com/en/developer-guide/snowpark-container-services/specification-reference>`__
+ uploaded to a stage.
Review Comment:
Good call, I simplified it.
##########
providers/snowflake/src/airflow/providers/snowflake/operators/snowpark_containers.py:
##########
@@ -0,0 +1,215 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from __future__ import annotations
+
+import time
+from collections.abc import Sequence
+from functools import cached_property
+from typing import TYPE_CHECKING, Any
+
+from airflow.providers.common.compat.standard.operators import BaseOperator
+from airflow.providers.common.sql.hooks.handlers import fetch_one_handler
+from airflow.providers.snowflake.hooks.snowflake import SnowflakeHook
+
+if TYPE_CHECKING:
+ from airflow.providers.common.compat.sdk import Context
+
+
+class SnowparkContainerJobOperator(BaseOperator):
+ """
+ Execute a job on Snowpark Container Services.
+
+ Submits a container job to a compute pool via ``EXECUTE JOB SERVICE``,
+ optionally polls for completion, retrieves container logs, and
+ drops the job service on success.
+
+ .. seealso::
+ `Snowpark Container Services <https://docs.snowflake.com/en/developer-guide/snowpark-container-services/overview>`_
+
+ :param compute_pool: name of the compute pool to run the job on
+ :param container_name: container name as defined in the service specification file,
+ used for retrieving container logs
+ :param spec: spec filename on the stage (e.g. ``'spec.yaml'``).
+ Must be provided together with ``spec_stage``
+ :param spec_stage: stage where the spec file is stored (e.g. ``'@my_stage'``).
+ Must be provided together with ``spec``
+ :param spec_text: inline YAML spec text, as an alternative to ``spec``/``spec_stage``.
+ The text is wrapped in ``$$`` delimiters automatically
+ :param name: (Optional) job service name. If not provided, Snowflake
+ auto-generates a name
+ :param query_warehouse: (Optional) warehouse for SQL queries run inside the container.
+ This is separate from the ``warehouse`` parameter used by the operator's
+ own SQL commands
+ :param replicas: (Optional) number of job replicas to run. (default value: 1)
+ :param wait_for_completion: poll until the job reaches a terminal state.
+ When disabled, the job is submitted and the operator returns
+ immediately. (default value: True)
+ :param drop_on_completion: drop the job service after the job finishes
+ successfully. Failed jobs are not dropped, allowing inspection
+ in Snowflake. (default value: True)
+ :param poll_interval: the interval in seconds to poll the query status.
+ (default value: 10)
+ :param snowflake_conn_id: Reference to
+ :ref:`Snowflake connection id<howto/connection:snowflake>`
+ :param database: name of database (will overwrite database defined
+ in connection)
+ :param schema: name of schema (will overwrite schema defined in
+ connection)
+ :param role: name of role (will overwrite any role defined in
+ connection's extra JSON)
+ :param warehouse: name of warehouse (will overwrite any warehouse
+ defined in the connection's extra JSON). Used for the operator's
+ own SQL commands, not for the container's queries
+ """
+
+ template_fields: Sequence[str] = (
+ "compute_pool",
+ "spec",
+ "spec_stage",
+ "container_name",
+ "spec_text",
+ "name",
+ "query_warehouse",
+ "snowflake_conn_id",
+ )
+
+ def __init__(
+ self,
+ *,
+ compute_pool: str,
+ container_name: str,
+ spec: str | None = None,
+ spec_stage: str | None = None,
+ spec_text: str | None = None,
+ name: str | None = None,
+ query_warehouse: str | None = None,
+ replicas: int = 1,
+ wait_for_completion: bool = True,
+ drop_on_completion: bool = True,
+ poll_interval: int = 10,
+ snowflake_conn_id: str = "snowflake_default",
+ database: str | None = None,
+ schema: str | None = None,
+ role: str | None = None,
+ warehouse: str | None = None,
+ **kwargs: Any,
+ ) -> None:
+ super().__init__(**kwargs)
+ self.compute_pool = compute_pool
+ self.container_name = container_name
+ self.spec = spec
+ self.spec_stage = spec_stage
+ self.spec_text = spec_text
+ self.name = name
+ self.query_warehouse = query_warehouse
+ self.replicas = replicas
+ self.wait_for_completion = wait_for_completion
+ self.drop_on_completion = drop_on_completion
+ self.poll_interval = poll_interval
+ self.snowflake_conn_id = snowflake_conn_id
+ self.database = database
+ self.schema = schema
+ self.role = role
+ self.warehouse = warehouse
+ self.job_name: str | None = None
Review Comment:
Added a comment.
##########
providers/snowflake/src/airflow/providers/snowflake/operators/snowpark_containers.py:
##########
@@ -0,0 +1,215 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from __future__ import annotations
+
+import time
+from collections.abc import Sequence
+from functools import cached_property
+from typing import TYPE_CHECKING, Any
+
+from airflow.providers.common.compat.standard.operators import BaseOperator
+from airflow.providers.common.sql.hooks.handlers import fetch_one_handler
+from airflow.providers.snowflake.hooks.snowflake import SnowflakeHook
+
+if TYPE_CHECKING:
+ from airflow.providers.common.compat.sdk import Context
+
+
+class SnowparkContainerJobOperator(BaseOperator):
+ """
+ Execute a job on Snowpark Container Services.
+
+ Submits a container job to a compute pool via ``EXECUTE JOB SERVICE``,
+ optionally polls for completion, retrieves container logs, and
+ drops the job service on success.
+
+ .. seealso::
+ `Snowpark Container Services <https://docs.snowflake.com/en/developer-guide/snowpark-container-services/overview>`_
+
+ :param compute_pool: name of the compute pool to run the job on
+ :param container_name: container name as defined in the service specification file,
+ used for retrieving container logs
+ :param spec: spec filename on the stage (e.g. ``'spec.yaml'``).
+ Must be provided together with ``spec_stage``
+ :param spec_stage: stage where the spec file is stored (e.g. ``'@my_stage'``).
+ Must be provided together with ``spec``
+ :param spec_text: inline YAML spec text, as an alternative to ``spec``/``spec_stage``.
+ The text is wrapped in ``$$`` delimiters automatically
+ :param name: (Optional) job service name. If not provided, Snowflake
+ auto-generates a name
+ :param query_warehouse: (Optional) warehouse for SQL queries run inside the container.
+ This is separate from the ``warehouse`` parameter used by the operator's
+ own SQL commands
+ :param replicas: (Optional) number of job replicas to run. (default value: 1)
+ :param wait_for_completion: poll until the job reaches a terminal state.
+ When disabled, the job is submitted and the operator returns
+ immediately. (default value: True)
+ :param drop_on_completion: drop the job service after the job finishes
+ successfully. Failed jobs are not dropped, allowing inspection
+ in Snowflake. (default value: True)
+ :param poll_interval: the interval in seconds to poll the query status.
+ (default value: 10)
+ :param snowflake_conn_id: Reference to
+ :ref:`Snowflake connection id<howto/connection:snowflake>`
+ :param database: name of database (will overwrite database defined
+ in connection)
+ :param schema: name of schema (will overwrite schema defined in
+ connection)
+ :param role: name of role (will overwrite any role defined in
+ connection's extra JSON)
+ :param warehouse: name of warehouse (will overwrite any warehouse
+ defined in the connection's extra JSON). Used for the operator's
+ own SQL commands, not for the container's queries
+ """
+
+ template_fields: Sequence[str] = (
+ "compute_pool",
+ "spec",
+ "spec_stage",
+ "container_name",
+ "spec_text",
+ "name",
+ "query_warehouse",
+ "snowflake_conn_id",
+ )
+
+ def __init__(
+ self,
+ *,
+ compute_pool: str,
+ container_name: str,
+ spec: str | None = None,
+ spec_stage: str | None = None,
+ spec_text: str | None = None,
+ name: str | None = None,
+ query_warehouse: str | None = None,
+ replicas: int = 1,
+ wait_for_completion: bool = True,
+ drop_on_completion: bool = True,
+ poll_interval: int = 10,
+ snowflake_conn_id: str = "snowflake_default",
+ database: str | None = None,
+ schema: str | None = None,
+ role: str | None = None,
+ warehouse: str | None = None,
+ **kwargs: Any,
+ ) -> None:
+ super().__init__(**kwargs)
+ self.compute_pool = compute_pool
+ self.container_name = container_name
+ self.spec = spec
+ self.spec_stage = spec_stage
+ self.spec_text = spec_text
+ self.name = name
+ self.query_warehouse = query_warehouse
+ self.replicas = replicas
+ self.wait_for_completion = wait_for_completion
+ self.drop_on_completion = drop_on_completion
+ self.poll_interval = poll_interval
+ self.snowflake_conn_id = snowflake_conn_id
+ self.database = database
+ self.schema = schema
+ self.role = role
+ self.warehouse = warehouse
+ self.job_name: str | None = None
+
+ if self.spec_text and (self.spec or self.spec_stage):
+ raise ValueError("Cannot specify both 'spec_text' and 'spec'/'spec_stage'")
+ if not self.spec_text and not (self.spec and self.spec_stage):
+ raise ValueError("Must provide either 'spec_text' or both 'spec' and 'spec_stage'")
+
+ @cached_property
+ def _hook(self) -> SnowflakeHook:
+ return SnowflakeHook(
+ snowflake_conn_id=self.snowflake_conn_id,
+ warehouse=self.warehouse,
+ database=self.database,
+ schema=self.schema,
+ role=self.role,
+ )
+
+ def _build_sql(self) -> str:
+ """Build the execute job SQL statement."""
+ sql = f"EXECUTE JOB SERVICE IN COMPUTE POOL {self.compute_pool}"
+ if self.name:
+ sql += f" NAME = {self.name}"
+ sql += " ASYNC = TRUE"
+ if self.replicas > 1:
+ sql += f" REPLICAS = {self.replicas}"
+ if self.query_warehouse:
+ sql += f" QUERY_WAREHOUSE = {self.query_warehouse}"
+ if self.spec_text:
+ sql += f" FROM SPECIFICATION $${self.spec_text}$$"
+ else:
+ sql += f" FROM {self.spec_stage} SPEC = '{self.spec}'"
+ return sql
+
+ def _run_one(self, sql: str, return_dictionaries: bool = False) -> Any:
+ """Run a single statement that returns one row via fetch_one_handler."""
+ return self._hook.run(sql, handler=fetch_one_handler, return_dictionaries=return_dictionaries)
+
+ def _submit_job(self) -> None:
+ """Submit the job and store the job name on the instance."""
+ response = self._run_one(self._build_sql())
+ self.job_name = response[0].split("'")[1]
Review Comment:
To answer your first question - the other option I found was querying the job history (https://docs.snowflake.com/en/sql-reference/functions/get_job_history). We could use the query id that gets stored in the hook, apply a created_time_start filter (to get the most recent jobs), and then fetch the name from the result of that query (directly from the name column). I felt the extra layer of complexity wasn't really worth it given we can get it directly from the initial response.
If there are any errors with the query itself, the hook would raise and the user would be aware. The expectation is that if the job does submit correctly, the job name should always be present in the response.
To your second point - you're correct, it's more clear inside the execute method. I made the change.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org