PoC of AIP-69 - A first task can be called remote

.github/ISSUE_TEMPLATE/airflow_providers_bug_report.yml

@@ -94,6 +94,7 @@ body:
         - presto
         - qdrant
         - redis
+        - remote
         - salesforce
         - samba
         - segment

INSTALL

@@ -277,9 +277,9 @@ cncf.kubernetes, cohere, common.compat, common.io, common.sql, databricks, datad
 dingding, discord, docker, elasticsearch, exasol, fab, facebook, ftp, github, google, grpc,
 hashicorp, http, imap, influxdb, jdbc, jenkins, microsoft.azure, microsoft.mssql, microsoft.psrp,
 microsoft.winrm, mongo, mysql, neo4j, odbc, openai, openfaas, openlineage, opensearch, opsgenie,
-oracle, pagerduty, papermill, pgvector, pinecone, postgres, presto, qdrant, redis, salesforce,
-samba, segment, sendgrid, sftp, singularity, slack, smtp, snowflake, sqlite, ssh, tableau, tabular,
-telegram, teradata, trino, vertica, weaviate, yandex, ydb, zendesk
+oracle, pagerduty, papermill, pgvector, pinecone, postgres, presto, qdrant, redis, remote,
+salesforce, samba, segment, sendgrid, sftp, singularity, slack, smtp, snowflake, sqlite, ssh,
+tableau, tabular, telegram, teradata, trino, vertica, weaviate, yandex, ydb, zendesk
 
 # END PROVIDER EXTRAS HERE
 

airflow/api_internal/endpoints/rpc_api_endpoint.py

@@ -39,7 +39,7 @@
 
 
 @functools.lru_cache
-def _initialize_map() -> dict[str, Callable]:
+def initialize_method_map() -> dict[str, Callable]:
     from airflow.cli.commands.task_command import _get_ti_db_access
     from airflow.dag_processing.manager import DagFileProcessorManager
     from airflow.dag_processing.processor import DagFileProcessor
@@ -148,7 +148,7 @@ def internal_airflow_api(body: dict[str, Any]) -> APIResponse:
     if json_rpc != "2.0":
         return log_and_build_error_response(message="Expected jsonrpc 2.0 request.", status=400)
 
-    methods_map = _initialize_map()
+    methods_map = initialize_method_map()
     method_name = body.get("method")
     if method_name not in methods_map:
         return log_and_build_error_response(message=f"Unrecognized method: {method_name}.", status=400)

airflow/api_internal/internal_api_call.py

@@ -21,17 +21,21 @@
 import json
 import logging
 from functools import wraps
-from typing import Callable, TypeVar
+from typing import TYPE_CHECKING, Callable, TypeVar
 
 import requests
 import tenacity
+from requests.auth import HTTPBasicAuth
 from urllib3.exceptions import NewConnectionError
 
 from airflow.configuration import conf
 from airflow.exceptions import AirflowConfigException, AirflowException
 from airflow.settings import _ENABLE_AIP_44
 from airflow.typing_compat import ParamSpec
 
+if TYPE_CHECKING:
+    from requests.auth import AuthBase
+
 PS = ParamSpec("PS")
 RT = TypeVar("RT")
 
@@ -44,6 +48,7 @@ class InternalApiConfig:
     _initialized = False
     _use_internal_api = False
     _internal_api_endpoint = ""
+    _internal_api_auth: AuthBase | None = None
 
     @staticmethod
     def force_database_direct_access():
@@ -56,6 +61,19 @@ def force_database_direct_access():
         InternalApiConfig._initialized = True
         InternalApiConfig._use_internal_api = False
 
+    @staticmethod
+    def force_api_access(api_endpoint: str, auth: AuthBase):
+        """
+        Force using Internal API with provided endpoint.
+
+        All methods decorated with internal_api_call will always be executed remote/via API.
+        This mode is needed for remote setups/remote executor.
+        """
+        InternalApiConfig._initialized = True
+        InternalApiConfig._use_internal_api = True
+        InternalApiConfig._internal_api_endpoint = api_endpoint
+        InternalApiConfig._internal_api_auth = auth
+
     @staticmethod
     def get_use_internal_api():
         if not InternalApiConfig._initialized:
@@ -68,21 +86,31 @@ def get_internal_api_endpoint():
             InternalApiConfig._init_values()
         return InternalApiConfig._internal_api_endpoint
 
+    @staticmethod
+    def get_auth() -> AuthBase | None:
+        return InternalApiConfig._internal_api_auth
+
     @staticmethod
     def _init_values():
         use_internal_api = conf.getboolean("core", "database_access_isolation", fallback=False)
         if use_internal_api and not _ENABLE_AIP_44:
             raise RuntimeError("The AIP_44 is not enabled so you cannot use it.")
-        internal_api_endpoint = ""
         if use_internal_api:
-            internal_api_url = conf.get("core", "internal_api_url")
-            internal_api_endpoint = internal_api_url + "/internal_api/v1/rpcapi"
-            if not internal_api_endpoint.startswith("http://"):
-                raise AirflowConfigException("[core]internal_api_url must start with http://")
+            internal_api_endpoint = conf.get("core", "internal_api_url")
+            if internal_api_endpoint.find("/", 8) == -1:
+                internal_api_endpoint = internal_api_endpoint + "/internal_api/v1/rpcapi"
+            if not internal_api_endpoint.startswith("http://") and not internal_api_endpoint.startswith(
+                "https://"
+            ):
+                raise AirflowConfigException("[core]internal_api_url must start with http:// or https://")
+            InternalApiConfig._internal_api_endpoint = internal_api_endpoint
+            internal_api_user = conf.get("core", "internal_api_user")
+            internal_api_password = conf.get("core", "internal_api_password")
+            if internal_api_user and internal_api_password:
+                InternalApiConfig._internal_api_auth = HTTPBasicAuth(internal_api_user, internal_api_password)
 
         InternalApiConfig._initialized = True
         InternalApiConfig._use_internal_api = use_internal_api
-        InternalApiConfig._internal_api_endpoint = internal_api_endpoint
 
 
 def internal_api_call(func: Callable[PS, RT]) -> Callable[PS, RT]:
@@ -112,7 +140,8 @@ def internal_api_call(func: Callable[PS, RT]) -> Callable[PS, RT]:
     def make_jsonrpc_request(method_name: str, params_json: str) -> bytes:
         data = {"jsonrpc": "2.0", "method": method_name, "params": params_json}
         internal_api_endpoint = InternalApiConfig.get_internal_api_endpoint()
-        response = requests.post(url=internal_api_endpoint, data=json.dumps(data), headers=headers)
+        auth = InternalApiConfig.get_auth()
+        response = requests.post(url=internal_api_endpoint, data=json.dumps(data), headers=headers, auth=auth)
         if response.status_code != 200:
             raise AirflowException(
                 f"Got {response.status_code}:{response.reason} when sending "

airflow/example_dags/integration_test.py

@@ -0,0 +1,126 @@
+#
+# 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.
+"""
+In this DAG all critical functions as integration test are contained.
+
+The DAG should work in all standard setups without error.
+"""
+
+from __future__ import annotations
+
+from datetime import datetime
+
+from airflow.decorators import task, task_group
+from airflow.exceptions import AirflowNotFoundException
+from airflow.hooks.base import BaseHook
+from airflow.models.dag import DAG
+from airflow.models.param import Param
+from airflow.models.variable import Variable
+from airflow.operators.bash import BashOperator
+from airflow.operators.empty import EmptyOperator
+from airflow.operators.python import PythonOperator
+
+with DAG(
+    dag_id="integration_test",
+    dag_display_name="Integration Test",
+    description=__doc__.partition(".")[0],
+    doc_md=__doc__,
+    schedule=None,
+    start_date=datetime(2024, 7, 1),
+    tags=["example", "params", "integration test"],
+    params={
+        "mapping_count": Param(
+            4,
+            type="integer",
+            title="Mapping Count",
+            description="Amount of tasks that should be mapped",
+        ),
+    },
+) as dag:
+
+    @task
+    def my_setup():
+        print("Assume this is a setup task")
+
+    @task
+    def mapping_from_params(**context) -> list[int]:
+        mapping_count: int = context["params"]["mapping_count"]
+        return list(range(1, mapping_count + 1))
+
+    @task
+    def add_one(x: int):
+        return x + 1
+
+    @task
+    def sum_it(values):
+        total = sum(values)
+        print(f"Total was {total}")
+
+    @task_group(prefix_group_id=False)
+    def mapping_task_group():
+        added_values = add_one.expand(x=mapping_from_params())
+        sum_it(added_values)
+
+    @task.branch
+    def branching():
+        return ["bash", "virtualenv", "variable", "connection", "classic_bash", "classic_python"]
+
+    @task.bash
+    def bash():
+        return "echo hello world"
+
+    @task.virtualenv(requirements="numpy")
+    def virtualenv():
+        import numpy
+
+        print(f"Welcome to virtualenv with numpy version {numpy.__version__}.")
+
+    @task
+    def variable():
+        Variable.set("integration_test_key", "value")
+        assert Variable.get("integration_test_key") == "value"  # noqa: S101
+        Variable.delete("integration_test_key")
+
+    @task
+    def connection():
+        try:
+            conn = BaseHook.get_connection("integration_test")
+            print(f"Got connection {conn}")
+        except AirflowNotFoundException:
+            print("Connection not found... but also OK.")
+
+    @task_group(prefix_group_id=False)
+    def standard_tasks_group():
+        classic_bash = BashOperator(
+            task_id="classic_bash", bash_command="echo Parameter is {{ params.mapping_count }}"
+        )
+
+        empty = EmptyOperator(task_id="not_executed")
+
+        def python_call():
+            print("Hello world")
+
+        classic_py = PythonOperator(task_id="classic_python", python_callable=python_call)
+
+        branching() >> [bash(), virtualenv(), variable(), connection(), classic_bash, classic_py, empty]
+
+    @task
+    def my_teardown():
+        print("Assume this is a teardown task")
+
+    my_setup().as_setup() >> [mapping_task_group(), standard_tasks_group()] >> my_teardown().as_teardown()

airflow/executors/executor_constants.py

@@ -34,6 +34,7 @@ class ConnectorSource(Enum):
 CELERY_EXECUTOR = "CeleryExecutor"
 CELERY_KUBERNETES_EXECUTOR = "CeleryKubernetesExecutor"
 KUBERNETES_EXECUTOR = "KubernetesExecutor"
+REMOTE_EXECUTOR = "RemoteExecutor"
 DEBUG_EXECUTOR = "DebugExecutor"
 MOCK_EXECUTOR = "MockExecutor"
 CORE_EXECUTOR_NAMES = {
@@ -43,6 +44,7 @@ class ConnectorSource(Enum):
     CELERY_EXECUTOR,
     CELERY_KUBERNETES_EXECUTOR,
     KUBERNETES_EXECUTOR,
+    REMOTE_EXECUTOR,
     DEBUG_EXECUTOR,
     MOCK_EXECUTOR,
 }

airflow/executors/executor_loader.py

@@ -34,6 +34,7 @@
     KUBERNETES_EXECUTOR,
     LOCAL_EXECUTOR,
     LOCAL_KUBERNETES_EXECUTOR,
+    REMOTE_EXECUTOR,
     SEQUENTIAL_EXECUTOR,
     ConnectorSource,
 )
@@ -70,6 +71,7 @@ class ExecutorLoader:
         "executors.celery_kubernetes_executor.CeleryKubernetesExecutor",
         KUBERNETES_EXECUTOR: "airflow.providers.cncf.kubernetes."
         "executors.kubernetes_executor.KubernetesExecutor",
+        REMOTE_EXECUTOR: "airflow.providers.remote.executors.RemoteExecutor",
         DEBUG_EXECUTOR: "airflow.executors.debug_executor.DebugExecutor",
     }
 
@@ -334,6 +336,9 @@ def validate_database_executor_compatibility(cls, executor: type[BaseExecutor])
         if InternalApiConfig.get_use_internal_api():
             return
 
+        if executor.__name__ == REMOTE_EXECUTOR:
+            return
+
         from airflow.settings import engine
 
         # SQLite only works with single threaded executors

airflow/providers/remote/CHANGELOG.rst

@@ -0,0 +1,38 @@
+ .. 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.
+
+
+.. NOTE TO CONTRIBUTORS:
+   Please, only add notes to the Changelog just below the "Changelog" header when there are some breaking changes
+   and you want to add an explanation to the users on how they are supposed to deal with them.
+   The changelog is updated and maintained semi-automatically by release manager.
+
+``apache-airflow-providers-remote-executor``
+
+
+Changelog
+---------
+
+0.1.0
+.....
+
+|experimental|
+
+Initial version of the provider.
+
+.. note::
+  This provider is currently experimental