Add @task.analytics Decorator (#62648)

* Add @task.analytics decorator

* Fix publish-docs to respect not ready packages when requested to build

Author: gopidesupavan

dev/breeze/src/airflow_breeze/utils/packages.py

@@ -419,7 +419,10 @@ def find_matching_long_package_names(
     removed_packages: list[str] = [
         f"apache-airflow-providers-{provider.replace('.', '-')}" for provider in get_removed_provider_ids()
     ]
-    all_packages_including_removed: list[str] = available_doc_packages + removed_packages
+    not_ready_packages: list[str] = [
+        f"apache-airflow-providers-{provider.replace('.', '-')}" for provider in get_not_ready_provider_ids()
+    ]
+    all_packages_including_removed: list[str] = available_doc_packages + removed_packages + not_ready_packages
     invalid_filters = [
         f
         for f in processed_package_filters

providers/common/sql/docs/operators.rst

@@ -317,3 +317,14 @@ Local File System Storage
     :dedent: 4
     :start-after: [START howto_analytics_operator_with_local]
     :end-before: [END howto_analytics_operator_with_local]
+
+Analytics TaskFlow Decorator
+----------------------------
+
+The ``@task.analytics`` decorator lets you write a function that returns the
+analytics sql queries:
+
+.. exampleinclude:: /../../sql/src/airflow/providers/common/sql/example_dags/example_analytics.py
+    :language: python
+    :start-after: [START howto_analytics_decorator]
+    :end-before: [END howto_analytics_decorator]

providers/common/sql/provider.yaml

@@ -129,3 +129,5 @@ sensors:
 task-decorators:
   - class-name: airflow.providers.common.sql.decorators.sql.sql_task
     name: sql
+  - class-name: airflow.providers.common.sql.decorators.analytics.analytics_task
+    name: analytics

providers/common/sql/src/airflow/providers/common/sql/decorators/analytics.py

@@ -0,0 +1,137 @@
+# 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 collections.abc import Callable, Collection, Mapping, Sequence
+from typing import TYPE_CHECKING, Any, ClassVar
+
+from airflow.providers.common.compat.sdk import (
+    AIRFLOW_V_3_0_PLUS,
+    DecoratedOperator,
+    TaskDecorator,
+    context_merge,
+    task_decorator_factory,
+)
+from airflow.providers.common.sql.operators.analytics import AnalyticsOperator
+from airflow.utils.operator_helpers import determine_kwargs
+
+if AIRFLOW_V_3_0_PLUS:
+    from airflow.sdk.definitions._internal.types import SET_DURING_EXECUTION
+else:
+    from airflow.utils.types import NOTSET as SET_DURING_EXECUTION  # type: ignore[attr-defined,no-redef]
+
+
+if TYPE_CHECKING:
+    from airflow.providers.common.compat.sdk import Context
+
+
+class _AnalyticsDecoratedOperator(DecoratedOperator, AnalyticsOperator):
+    """
+    Wraps a Python callable and uses the callable return value as the SQL commands to be executed.
+
+    :param python_callable: A reference to an object that is callable.
+    :param op_kwargs: A dictionary of keyword arguments that will get unpacked (templated).
+    :param op_args: A list of positional arguments that will get unpacked (templated).
+    """
+
+    template_fields: Sequence[str] = (
+        *DecoratedOperator.template_fields,
+        *AnalyticsOperator.template_fields,
+    )
+    template_fields_renderers: ClassVar[dict[str, str]] = {
+        **DecoratedOperator.template_fields_renderers,
+        **AnalyticsOperator.template_fields_renderers,
+    }
+
+    overwrite_rtif_after_execution: bool = True
+
+    custom_operator_name: str = "@task.analytics"
+
+    def __init__(
+        self,
+        python_callable: Callable,
+        op_args: Collection[Any] | None = None,
+        op_kwargs: Mapping[str, Any] | None = None,
+        **kwargs,
+    ) -> None:
+        super().__init__(
+            python_callable=python_callable,
+            op_args=op_args,
+            op_kwargs=op_kwargs,
+            queries=SET_DURING_EXECUTION,
+            **kwargs,
+        )
+
+    @property
+    def xcom_push(self) -> bool:
+        """Compatibility property for BaseDecorator that expects xcom_push attribute."""
+        return self.do_xcom_push
+
+    @xcom_push.setter
+    def xcom_push(self, value: bool) -> None:
+        """Compatibility setter for BaseDecorator that expects xcom_push attribute."""
+        self.do_xcom_push = value
+
+    def execute(self, context: Context) -> Any:
+        """
+        Build the SQL and execute the generated query (or queries).
+
+        :param context: Airflow context.
+        :return: Any
+        """
+        context_merge(context, self.op_kwargs)
+        kwargs = determine_kwargs(self.python_callable, self.op_args, context)
+
+        # Set the queries using the Python callable
+        result = self.python_callable(*self.op_args, **kwargs)
+
+        # Only non-empty strings and non-empty lists of non-empty strings are acceptable return types
+        if (
+            not isinstance(result, (str, list))
+            or (isinstance(result, str) and not result.strip())
+            or (
+                isinstance(result, list)
+                and (not result or not all(isinstance(s, str) and s.strip() for s in result))
+            )
+        ):
+            raise TypeError(
+                "The returned value from the @task.analytics callable must be a non-empty string "
+                "or a non-empty list of non-empty strings."
+            )
+
+        # AnalyticsOperator expects queries as a list of strings
+        self.queries = [result] if isinstance(result, str) else result
+
+        self.render_template_fields(context)
+
+        return AnalyticsOperator.execute(self, context)
+
+
+def analytics_task(python_callable=None, **kwargs) -> TaskDecorator:
+    """
+    Wrap a Python function into a AnalyticsOperator.
+
+    :param python_callable: Function to decorate.
+
+    :meta private:
+    """
+    return task_decorator_factory(
+        python_callable=python_callable,
+        decorated_operator_class=_AnalyticsDecoratedOperator,
+        **kwargs,
+    )

providers/common/sql/src/airflow/providers/common/sql/example_dags/example_analytics.py

@@ -20,7 +20,7 @@
 
 from airflow.providers.common.sql.config import DataSourceConfig
 from airflow.providers.common.sql.operators.analytics import AnalyticsOperator
-from airflow.sdk import DAG
+from airflow.sdk import DAG, task
 
 datasource_config_s3 = DataSourceConfig(
     conn_id="aws_default", table_name="users_data", uri="s3://bucket/path/", format="parquet"
@@ -56,3 +56,12 @@
     )
     analytics_with_s3 >> analytics_with_local
     # [END howto_analytics_operator_with_local]
+
+    # [START howto_analytics_decorator]
+    @task.analytics(datasource_configs=[datasource_config_s3])
+    def get_user_summary_queries():
+        return ["SELECT * FROM users_data LIMIT 10", "SELECT count(*) FROM users_data"]
+
+    # [END howto_analytics_decorator]
+
+    analytics_with_local >> get_user_summary_queries()

providers/common/sql/src/airflow/providers/common/sql/get_provider_info.py

@@ -71,6 +71,10 @@ def get_provider_info():
             {"integration-name": "Common SQL", "python-modules": ["airflow.providers.common.sql.sensors.sql"]}
         ],
         "task-decorators": [
-            {"class-name": "airflow.providers.common.sql.decorators.sql.sql_task", "name": "sql"}
+            {"class-name": "airflow.providers.common.sql.decorators.sql.sql_task", "name": "sql"},
+            {
+                "class-name": "airflow.providers.common.sql.decorators.analytics.analytics_task",
+                "name": "analytics",
+            },
         ],
     }

providers/common/sql/tests/unit/common/sql/decorators/test_analytics.py

@@ -0,0 +1,122 @@
+# 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.mock import MagicMock, patch
+
+import pytest
+
+from airflow.providers.common.sql.config import DataSourceConfig
+from airflow.providers.common.sql.decorators.analytics import _AnalyticsDecoratedOperator
+
+DATASOURCE_CONFIGS = [
+    DataSourceConfig(conn_id="", table_name="users_data", uri="file:///path/to/", format="parquet")
+]
+
+
+class TestAnalyticsDecoratedOperator:
+    def test_custom_operator_name(self):
+        assert _AnalyticsDecoratedOperator.custom_operator_name == "@task.analytics"
+
+    @patch(
+        "airflow.providers.common.sql.operators.analytics.AnalyticsOperator.execute",
+        autospec=True,
+    )
+    def test_execute_calls_callable_and_sets_queries_from_list(self, mock_execute):
+        """The callable return value (list) becomes self.queries."""
+        mock_execute.return_value = "mocked output"
+
+        def get_user_queries():
+            return ["SELECT * FROM users_data", "SELECT count(*) FROM users_data"]
+
+        op = _AnalyticsDecoratedOperator(
+            task_id="test",
+            python_callable=get_user_queries,
+            datasource_configs=DATASOURCE_CONFIGS,
+        )
+        result = op.execute(context={})
+
+        assert result == "mocked output"
+        assert op.queries == ["SELECT * FROM users_data", "SELECT count(*) FROM users_data"]
+        mock_execute.assert_called_once()
+
+    @patch(
+        "airflow.providers.common.sql.operators.analytics.AnalyticsOperator.execute",
+        autospec=True,
+    )
+    def test_execute_wraps_single_string_into_list(self, mock_execute):
+        """A single string return value is wrapped into a list for self.queries."""
+        mock_execute.return_value = "mocked output"
+
+        def get_single_query():
+            return "SELECT 1"
+
+        op = _AnalyticsDecoratedOperator(
+            task_id="test",
+            python_callable=get_single_query,
+            datasource_configs=DATASOURCE_CONFIGS,
+        )
+        op.execute(context={})
+
+        assert op.queries == ["SELECT 1"]
+
+    @pytest.mark.parametrize(
+        "return_value",
+        [42, "", "   ", None, [], [""], ["SELECT 1", ""], ["SELECT 1", "   "], [42]],
+        ids=[
+            "non-string",
+            "empty-string",
+            "whitespace-string",
+            "none",
+            "empty-list",
+            "list-with-empty-string",
+            "list-with-one-valid-one-empty",
+            "list-with-one-valid-one-whitespace",
+            "list-with-non-string",
+        ],
+    )
+    def test_execute_raises_on_invalid_return_value(self, return_value):
+        """TypeError when the callable returns an invalid value."""
+        op = _AnalyticsDecoratedOperator(
+            task_id="test",
+            python_callable=lambda: return_value,
+            datasource_configs=DATASOURCE_CONFIGS,
+        )
+        with pytest.raises(TypeError, match="non-empty string"):
+            op.execute(context={})
+
+    @patch(
+        "airflow.providers.common.sql.operators.analytics.AnalyticsOperator.execute",
+        autospec=True,
+    )
+    def test_execute_merges_op_kwargs_into_callable(self, mock_execute):
+        """op_kwargs are forwarded to the callable to build queries."""
+        mock_execute.return_value = "mocked output"
+
+        def get_queries_for_table(table_name):
+            return [f"SELECT * FROM {table_name}", f"SELECT count(*) FROM {table_name}"]
+
+        op = _AnalyticsDecoratedOperator(
+            task_id="test",
+            python_callable=get_queries_for_table,
+            datasource_configs=DATASOURCE_CONFIGS,
+            op_kwargs={"table_name": "orders"},
+        )
+        op.execute(context={"task_instance": MagicMock()})
+
+        assert op.queries == ["SELECT * FROM orders", "SELECT count(*) FROM orders"]