recommit without precommit changes

Author: Bowrna

.pre-commit-config.yaml

@@ -1519,4 +1519,12 @@ repos:
         additional_dependencies: ['rich>=12.4.4', 'ruff==0.11.13']
         files: ^providers/.*/src/airflow/providers/.*version_compat.*\.py$
         require_serial: true
+      - id: verify-signatures
+        name: Verify if decorator signatures added in __init__.pyi is not missing any valid params of the operator
+        language: python
+        entry: ./scripts/ci/pre_commit/verify_signature_consistency.py
+        pass_filenames: false
+        additional_dependencies: ['rich>=12.4.4']
+        always_run: true
+        files: ^(providers/.*/)?airflow/.*/(sensors|operators)/.*\.py$
         ## ONLY ADD PRE-COMMITS HERE THAT REQUIRE CI IMAGE

scripts/ci/pre_commit/verify_signature_consistency.py

@@ -0,0 +1,40 @@
+#!/usr/bin/env python
+#
+# 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 sys
+from pathlib import Path
+
+sys.path.insert(0, str(Path(__file__).parent.resolve()))
+from common_precommit_utils import (
+    initialize_breeze_precommit,
+    run_command_via_breeze_shell,
+    validate_cmd_result,
+)
+
+initialize_breeze_precommit(__name__, __file__)
+
+cmd_result = run_command_via_breeze_shell(
+    ["python3", "/opt/airflow/scripts/in_container/run_signature_consistency_verify.py"],
+    backend="sqlite",
+    skip_environment_initialization=False,
+)
+
+validate_cmd_result(cmd_result)

scripts/in_container/run_signature_consistency_verify.py

@@ -0,0 +1,114 @@
+#!/usr/bin/env python
+#
+# 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
+import sys
+
+import libcst as cst
+from in_container_utils import AIRFLOW_ROOT_PATH, console
+
+DECORATOR_OPERATOR_MAP = {
+    "kubernetes": "airflow.providers.cncf.kubernetes.operators.pod.KubernetesPodOperator",
+    "sensor": "airflow.sdk.bases.sensor.BaseSensorOperator",
+    "virtualenv": "airflow.providers.standard.operators.python.PythonVirtualenvOperator",
+    "branch_virtualenv": "airflow.providers.standard.operators.python.BranchPythonVirtualenvOperator",
+    # Add more here...
+}
+DECORATOR_PYI_PATH = (
+    AIRFLOW_ROOT_PATH / "task-sdk" / "src" / "airflow" / "sdk" / "definitions" / "decorators" / "__init__.pyi"
+)
+decorator_pyi_file_content = DECORATOR_PYI_PATH.read_text()
+
+
+def extract_function_params(code, function_name, return_type):
+    """Extracts parameters from a specific function definition in the given code.
+
+    Args:
+        code (str): The Python code to parse.
+        function_name (str): The name of the function to extract parameters from.
+        return_type (str): As the pyi file has multiple @overload decorator, extract function param based on return type.
+
+    Returns:
+        list: A list of parameter names, or None if the function is not found.
+    """
+    module = cst.parse_module(code)
+
+    class FunctionParamExtractor(cst.CSTVisitor):
+        def __init__(self, target_function_name, target_return_type):
+            self.target_function_name = target_function_name
+            self.target_return_type = target_return_type
+            self.params: list[str] = []
+
+        def visit_FunctionDef(self, node):
+            # Match function name
+            if node.name.value == self.target_function_name:
+                if node.returns:
+                    annotation = node.returns.annotation
+                    if isinstance(annotation, cst.Name) and annotation.value == self.target_return_type:
+                        parameters_node = node.params
+                        self.params.extend(param.name.value for param in parameters_node.params)
+                        self.params.extend(param.name.value for param in parameters_node.kwonly_params)
+                        self.params.extend(param.name.value for param in parameters_node.posonly_params)
+                        if parameters_node.star_kwarg:
+                            self.params.append(parameters_node.star_kwarg.name.value)
+                            return False  # Stop traversing after finding the real function
+            return True  # Keep traversing
+
+    extractor = FunctionParamExtractor(function_name, return_type)
+    module.visit(extractor)
+    return extractor.params
+
+
+def get_decorator_params(decorator_name: str):
+    params = extract_function_params(decorator_pyi_file_content, decorator_name, "TaskDecorator")
+    return set(params)
+
+
+def get_operator_params(operator_path: str):
+    console.print("Operator path:", operator_path)
+    module_path, class_name = operator_path.rsplit(".", 1)
+    module = __import__(module_path, fromlist=[class_name])
+    operator_cls = getattr(module, class_name)
+    sig = inspect.signature(operator_cls.__init__)
+    return set(p for p in sig.parameters.keys() if p not in ("self", "args", "kwargs"))
+
+
+def verify_signature_consistency():
+    failure = False
+    console.print("Verify signature consistency")
+    for decorator, operator_path in DECORATOR_OPERATOR_MAP.items():
+        decorator_params = get_decorator_params(decorator)
+        operator_params = get_operator_params(operator_path)
+        missing_in_decorator = operator_params - decorator_params
+
+        ignored = {"kwargs", "args", "self"}
+        missing_in_decorator -= ignored
+        if missing_in_decorator:
+            failure = True
+            console.print("Following list of params are missing in  __init__.py[/]:", missing_in_decorator)
+    if failure:
+        console.print("[red]Some of the decorator signatures are missing in __init__.py[/]")
+        sys.exit(1)
+    console.print("[green]All decorator signature matches[/]")
+    sys.exit(0)
+
+
+if __name__ == "__main__":
+    verify_signature_consistency()

task-sdk/src/airflow/sdk/definitions/decorators/__init__.pyi

@@ -27,6 +27,7 @@ from typing import Any, Callable, TypeVar, overload
 from docker.types import Mount
 from kubernetes.client import models as k8s
 
+from airflow.providers.cncf.kubernetes.callbacks import KubernetesPodOperatorCallback
 from airflow.providers.cncf.kubernetes.secret import Secret
 from airflow.sdk.bases.decorator import FParams, FReturn, Task, TaskDecorator, _TaskDecorator
 from airflow.sdk.definitions.dag import dag
@@ -96,16 +97,21 @@ class TaskDecoratorCollection:
         multiple_outputs: bool | None = None,
         # 'python_callable', 'op_args' and 'op_kwargs' since they are filled by
         # _PythonVirtualenvDecoratedOperator.
+        python_callable: Callable,
+        op_args: Collection[Any] | None = None,
+        op_kwargs: Mapping[str, Any] | None = None,
         requirements: None | Iterable[str] | str = None,
         python_version: None | str | int | float = None,
         serializer: Literal["pickle", "cloudpickle", "dill"] | None = None,
         system_site_packages: bool = True,
         templates_dict: Mapping[str, Any] | None = None,
+        templates_exts: list[str] | None = None,
         pip_install_options: list[str] | None = None,
+        expect_airflow: bool = True,
         skip_on_exit_code: int | Container[int] | None = None,
         index_urls: None | Collection[str] | str = None,
         venv_cache_path: None | str = None,
-        show_return_value_in_logs: bool = True,
+        string_args: Iterable[str] | None = None,
         env_vars: dict[str, str] | None = None,
         inherit_env: bool = True,
         **kwargs,
@@ -218,18 +224,24 @@ class TaskDecoratorCollection:
         self,
         *,
         multiple_outputs: bool | None = None,
-        # 'python_callable', 'op_args' and 'op_kwargs' since they are filled by
-        # _PythonVirtualenvDecoratedOperator.
+        python_callable: Callable,
         requirements: None | Iterable[str] | str = None,
+        op_args: Collection[Any] | None = None,
+        op_kwargs: Mapping[str, Any] | None = None,
+        string_args: Iterable[str] | None = None,
         python_version: None | str | int | float = None,
         serializer: Literal["pickle", "cloudpickle", "dill"] | None = None,
         system_site_packages: bool = True,
         templates_dict: Mapping[str, Any] | None = None,
+        templates_exts: list[str] | None = None,
         pip_install_options: list[str] | None = None,
         skip_on_exit_code: int | Container[int] | None = None,
         index_urls: None | Collection[str] | str = None,
         venv_cache_path: None | str = None,
+        expect_airflow: bool = True,
         show_return_value_in_logs: bool = True,
+        env_vars: dict[str, str] | None = None,
+        inherit_env: bool = True,
         **kwargs,
     ) -> TaskDecorator:
         """Create a decorator to wrap the decorated callable into a BranchPythonVirtualenvOperator.
@@ -507,6 +519,7 @@ class TaskDecoratorCollection:
         image: str | None = None,
         name: str | None = None,
         random_name_suffix: bool = ...,
+        cmds: list[str] | None = None,
         arguments: list[str] | None = None,
         ports: list[k8s.V1ContainerPort] | None = None,
         volume_mounts: list[k8s.V1VolumeMount] | None = None,
@@ -520,6 +533,7 @@ class TaskDecoratorCollection:
         reattach_on_restart: bool = ...,
         startup_timeout_seconds: int = ...,
         startup_check_interval_seconds: int = ...,
+        schedule_timeout_seconds: int | None = None,
         get_logs: bool = True,
         container_logs: Iterable[str] | str | Literal[True] = ...,
         image_pull_policy: str | None = None,
@@ -530,6 +544,7 @@ class TaskDecoratorCollection:
         node_selector: dict | None = None,
         image_pull_secrets: list[k8s.V1LocalObjectReference] | None = None,
         service_account_name: str | None = None,
+        automount_service_account_token: bool | None = None,
         hostnetwork: bool = False,
         host_aliases: list[k8s.V1HostAlias] | None = None,
         tolerations: list[k8s.V1Toleration] | None = None,
@@ -553,13 +568,20 @@ class TaskDecoratorCollection:
         skip_on_exit_code: int | Container[int] | None = None,
         base_container_name: str | None = None,
         base_container_status_polling_interval: float = ...,
+        init_container_logs: Iterable[str] | str | Literal[True] | None = None,
         deferrable: bool = ...,
         poll_interval: float = ...,
         log_pod_spec_on_failure: bool = ...,
         on_finish_action: str = ...,
+        is_delete_operator_pod: None | bool = None,
         termination_message_policy: str = ...,
         active_deadline_seconds: int | None = None,
+        callbacks: (
+            list[type[KubernetesPodOperatorCallback]] | type[KubernetesPodOperatorCallback] | None
+        ) = None,
         progress_callback: Callable[[str], None] | None = None,
+        logging_interval: int | None = None,
+        trigger_kwargs: dict | None = None,
         **kwargs,
     ) -> TaskDecorator:
         """Create a decorator to convert a callable to a Kubernetes Pod task.
@@ -666,6 +688,7 @@ class TaskDecoratorCollection:
         :param active_deadline_seconds: The active_deadline_seconds which matches to active_deadline_seconds
             in V1PodSpec.
         :param progress_callback: Callback function for receiving k8s container logs.
+        :param trigger_kwargs: additional keyword parameters passed to the trigger
         """
     @overload
     def kubernetes(self, python_callable: Callable[FParams, FReturn]) -> Task[FParams, FReturn]: ...
@@ -849,6 +872,8 @@ class TaskDecoratorCollection:
         mode: str = ...,
         exponential_backoff: bool = False,
         max_wait: timedelta | float | None = None,
+        silent_fail: bool = False,
+        never_fail: bool = False,
         **kwargs,
     ) -> TaskDecorator:
         """
@@ -873,6 +898,13 @@ class TaskDecoratorCollection:
         :param exponential_backoff: allow progressive longer waits between
             pokes by using exponential backoff algorithm
         :param max_wait: maximum wait interval between pokes, can be ``timedelta`` or ``float`` seconds
+        :param silent_fail: If true, and poke method raises an exception different from
+            AirflowSensorTimeout, AirflowTaskTimeout, AirflowSkipException
+            and AirflowFailException, the sensor will log the error and continue
+            its execution. Otherwise, the sensor task fails, and it can be retried
+            based on the provided `retries` parameter.
+        :param never_fail: If true, and poke method raises an exception, sensor will be skipped.
+           Mutually exclusive with soft_fail.
         """
     @overload
     def sensor(self, python_callable: Callable[FParams, FReturn] | None = None) -> Task[FParams, FReturn]: ...