Add Snowpark operator and decorator (apache#42457)

airflow/providers/snowflake/decorators/__init__.py

@@ -0,0 +1,16 @@
+# 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.

airflow/providers/snowflake/decorators/snowpark.py

@@ -0,0 +1,124 @@
+# 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 typing import TYPE_CHECKING, Callable, Sequence
+
+from airflow.decorators.base import DecoratedOperator, task_decorator_factory
+from airflow.providers.snowflake.operators.snowpark import SnowparkOperator
+from airflow.providers.snowflake.utils.snowpark import inject_session_into_op_kwargs
+
+if TYPE_CHECKING:
+    from airflow.decorators.base import TaskDecorator
+
+
+class _SnowparkDecoratedOperator(DecoratedOperator, SnowparkOperator):
+    """
+    Wraps a Python callable that contains Snowpark code and captures args/kwargs when called for execution.
+
+    :param snowflake_conn_id: Reference to
+        :ref:`Snowflake connection id<howto/connection:snowflake>`
+    :param python_callable: A reference to an object that is callable
+    :param op_args: a list of positional arguments that will get unpacked when
+        calling your callable
+    :param op_kwargs: a dictionary of keyword arguments that will get unpacked
+        in your function
+    :param warehouse: name of warehouse (will overwrite any warehouse
+        defined in the connection's extra JSON)
+    :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 authenticator: authenticator for Snowflake.
+        'snowflake' (default) to use the internal Snowflake authenticator
+        'externalbrowser' to authenticate using your web browser and
+        Okta, ADFS or any other SAML 2.0-compliant identify provider
+        (IdP) that has been defined for your account
+        'https://<your_okta_account_name>.okta.com' to authenticate
+        through native Okta.
+    :param session_parameters: You can set session-level parameters at
+        the time you connect to Snowflake
+    :param multiple_outputs: If set to True, the decorated function's return value will be unrolled to
+        multiple XCom values. Dict will unroll to XCom values with its keys as XCom keys. Defaults to False.
+    """
+
+    custom_operator_name = "@task.snowpark"
+
+    def __init__(
+        self,
+        *,
+        snowflake_conn_id: str = "snowflake_default",
+        python_callable: Callable,
+        op_args: Sequence | None = None,
+        op_kwargs: dict | None = None,
+        warehouse: str | None = None,
+        database: str | None = None,
+        role: str | None = None,
+        schema: str | None = None,
+        authenticator: str | None = None,
+        session_parameters: dict | None = None,
+        **kwargs,
+    ) -> None:
+        kwargs_to_upstream = {
+            "python_callable": python_callable,
+            "op_args": op_args,
+            "op_kwargs": op_kwargs,
+        }
+        super().__init__(
+            kwargs_to_upstream=kwargs_to_upstream,
+            snowflake_conn_id=snowflake_conn_id,
+            python_callable=python_callable,
+            op_args=op_args,
+            # airflow.decorators.base.DecoratedOperator checks if the functions are bindable, so we have to
+            # add an artificial value to pass the validation if there is a keyword argument named `session`
+            # in the signature of the python callable. The real value is determined at runtime.
+            op_kwargs=inject_session_into_op_kwargs(python_callable, op_kwargs, None)
+            if op_kwargs is not None
+            else op_kwargs,
+            warehouse=warehouse,
+            database=database,
+            role=role,
+            schema=schema,
+            authenticator=authenticator,
+            session_parameters=session_parameters,
+            **kwargs,
+        )
+
+
+def snowpark_task(
+    python_callable: Callable | None = None,
+    multiple_outputs: bool | None = None,
+    **kwargs,
+) -> TaskDecorator:
+    """
+    Wrap a function that contains Snowpark code into an Airflow operator.
+
+    Accepts kwargs for operator kwarg. Can be reused in a single DAG.
+
+    :param python_callable: Function to decorate
+    :param multiple_outputs: If set to True, the decorated function's return value will be unrolled to
+        multiple XCom values. Dict will unroll to XCom values with its keys as XCom keys. Defaults to False.
+    """
+    return task_decorator_factory(
+        python_callable=python_callable,
+        multiple_outputs=multiple_outputs,
+        decorated_operator_class=_SnowparkDecoratedOperator,
+        **kwargs,
+    )

airflow/providers/snowflake/hooks/snowflake.py

@@ -321,6 +321,28 @@ def get_sqlalchemy_engine(self, engine_kwargs=None):
                 engine_kwargs["connect_args"][key] = conn_params[key]
         return create_engine(self._conn_params_to_sqlalchemy_uri(conn_params), **engine_kwargs)
 
+    def get_snowpark_session(self):
+        """
+        Get a Snowpark session object.
+
+        :return: the created session.
+        """
+        from snowflake.snowpark import Session
+
+        from airflow import __version__ as airflow_version
+        from airflow.providers.snowflake import __version__ as provider_version
+
+        conn_config = self._get_conn_params
+        session = Session.builder.configs(conn_config).create()
+        # add query tag for observability
+        session.update_query_tag(
+            {
+                "airflow_version": airflow_version,
+                "airflow_provider_version": provider_version,
+            }
+        )
+        return session
+
     def set_autocommit(self, conn, autocommit: Any) -> None:
         conn.autocommit(autocommit)
         conn.autocommit_mode = autocommit

airflow/providers/snowflake/operators/snowpark.py

@@ -0,0 +1,133 @@
+# 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 typing import Any, Callable, Collection, Mapping, Sequence
+
+from airflow.operators.python import PythonOperator, get_current_context
+from airflow.providers.snowflake.hooks.snowflake import SnowflakeHook
+from airflow.providers.snowflake.utils.snowpark import inject_session_into_op_kwargs
+
+
+class SnowparkOperator(PythonOperator):
+    """
+    Executes a Python function with Snowpark Python code.
+
+    .. seealso::
+        For more information on how to use this operator, take a look at the guide:
+        :ref:`howto/operator:SnowparkOperator`
+
+    :param snowflake_conn_id: Reference to
+        :ref:`Snowflake connection id<howto/connection:snowflake>`
+    :param python_callable: A reference to an object that is callable
+    :param op_args: a list of positional arguments that will get unpacked when
+        calling your callable
+    :param op_kwargs: a dictionary of keyword arguments that will get unpacked
+        in your function
+    :param templates_dict: a dictionary where the values are templates that
+        will get templated by the Airflow engine sometime between
+        ``__init__`` and ``execute`` takes place and are made available
+        in your callable's context after the template has been applied. (templated)
+    :param templates_exts: a list of file extensions to resolve while
+        processing templated fields, for examples ``['.sql', '.hql']``
+    :param show_return_value_in_logs: a bool value whether to show return_value
+        logs. Defaults to True, which allows return value log output.
+        It can be set to False to prevent log output of return value when you return huge data
+        such as transmission a large amount of XCom to TaskAPI.
+    :param warehouse: name of warehouse (will overwrite any warehouse
+        defined in the connection's extra JSON)
+    :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 authenticator: authenticator for Snowflake.
+        'snowflake' (default) to use the internal Snowflake authenticator
+        'externalbrowser' to authenticate using your web browser and
+        Okta, ADFS or any other SAML 2.0-compliant identify provider
+        (IdP) that has been defined for your account
+        'https://<your_okta_account_name>.okta.com' to authenticate
+        through native Okta.
+    :param session_parameters: You can set session-level parameters at
+        the time you connect to Snowflake
+    """
+
+    def __init__(
+        self,
+        *,
+        snowflake_conn_id: str = "snowflake_default",
+        python_callable: Callable,
+        op_args: Collection[Any] | None = None,
+        op_kwargs: Mapping[str, Any] | None = None,
+        templates_dict: dict[str, Any] | None = None,
+        templates_exts: Sequence[str] | None = None,
+        show_return_value_in_logs: bool = True,
+        warehouse: str | None = None,
+        database: str | None = None,
+        schema: str | None = None,
+        role: str | None = None,
+        authenticator: str | None = None,
+        session_parameters: dict | None = None,
+        **kwargs,
+    ):
+        super().__init__(
+            python_callable=python_callable,
+            op_args=op_args,
+            op_kwargs=op_kwargs,
+            templates_dict=templates_dict,
+            templates_exts=templates_exts,
+            show_return_value_in_logs=show_return_value_in_logs,
+            **kwargs,
+        )
+        self.snowflake_conn_id = snowflake_conn_id
+        self.warehouse = warehouse
+        self.database = database
+        self.schema = schema
+        self.role = role
+        self.authenticator = authenticator
+        self.session_parameters = session_parameters
+
+    def execute_callable(self):
+        hook = SnowflakeHook(
+            snowflake_conn_id=self.snowflake_conn_id,
+            warehouse=self.warehouse,
+            database=self.database,
+            role=self.role,
+            schema=self.schema,
+            authenticator=self.authenticator,
+            session_parameters=self.session_parameters,
+        )
+        session = hook.get_snowpark_session()
+        context = get_current_context()
+        session.update_query_tag(
+            {
+                "dag_id": context["dag_run"].dag_id,
+                "dag_run_id": context["dag_run"].run_id,
+                "task_id": context["task_instance"].task_id,
+                "operator": self.__class__.__name__,
+            }
+        )
+        try:
+            # inject session object if the function has "session" keyword as an argument
+            self.op_kwargs = inject_session_into_op_kwargs(
+                self.python_callable, dict(self.op_kwargs), session
+            )
+            return super().execute_callable()
+        finally:
+            session.close()

airflow/providers/snowflake/provider.yaml

@@ -90,19 +90,26 @@ dependencies:
   - pyarrow>=14.0.1
   - snowflake-connector-python>=3.7.1
   - snowflake-sqlalchemy>=1.4.0
+  - snowflake-snowpark-python>=1.17.0;python_version<"3.12"
 
 integrations:
   - integration-name: Snowflake
     external-doc-url: https://snowflake.com/
     how-to-guide:
       - /docs/apache-airflow-providers-snowflake/operators/snowflake.rst
+      - /docs/apache-airflow-providers-snowflake/operators/snowpark.rst
     logo: /integration-logos/snowflake/Snowflake.png
     tags: [service]
 
 operators:
   - integration-name: Snowflake
     python-modules:
       - airflow.providers.snowflake.operators.snowflake
+      - airflow.providers.snowflake.operators.snowpark
+
+task-decorators:
+  - class-name: airflow.providers.snowflake.decorators.snowpark.snowpark_task
+    name: snowpark
 
 hooks:
   - integration-name: Snowflake

airflow/providers/snowflake/utils/snowpark.py

@@ -0,0 +1,44 @@
+# 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 inspect
+from typing import TYPE_CHECKING, Callable
+
+if TYPE_CHECKING:
+    from snowflake.snowpark import Session
+
+
+def inject_session_into_op_kwargs(
+    python_callable: Callable, op_kwargs: dict, session: Session | None
+) -> dict:
+    """
+    Inject Snowpark session into operator kwargs based on signature of python callable.
+
+    If there is a keyword argument named `session` in the signature of the python callable,
+    a Snowpark session object will be injected into kwargs.
+
+    :param python_callable: Python callable
+    :param op_kwargs: Operator kwargs
+    :param session: Snowpark session
+    """
+    signature = inspect.signature(python_callable)
+    if "session" in signature.parameters:
+        return {**op_kwargs, "session": session}
+    else:
+        return op_kwargs