Repository: apache/airflow

Issue:
AIP-44 Migrate DagModel.get_paused_dag_ids to Internal API



Discussion:

Resolving patch:
diff --git a/airflow/api_internal/endpoints/rpc_api_endpoint.py b/airflow/api_internal/endpoints/rpc_api_endpoint.py
--- a/airflow/api_internal/endpoints/rpc_api_endpoint.py
+++ b/airflow/api_internal/endpoints/rpc_api_endpoint.py
@@ -33,9 +33,11 @@
 @functools.lru_cache()
 def _initialize_map() -> dict[str, Callable]:
     from airflow.dag_processing.processor import DagFileProcessor
+    from airflow.models.dag import DagModel
 
     functions: list[Callable] = [
         DagFileProcessor.update_import_errors,
+        DagModel.get_paused_dag_ids,
     ]
     return {f"{func.__module__}.{func.__name__}": func for func in functions}
 
diff --git a/airflow/api_internal/internal_api_call.py b/airflow/api_internal/internal_api_call.py
--- a/airflow/api_internal/internal_api_call.py
+++ b/airflow/api_internal/internal_api_call.py
@@ -26,7 +26,6 @@
 
 from airflow.configuration import conf
 from airflow.exceptions import AirflowConfigException, AirflowException
-from airflow.serialization.serialized_objects import BaseSerialization
 from airflow.typing_compat import ParamSpec
 
 PS = ParamSpec("PS")
@@ -67,7 +66,7 @@ def _init_values():
         InternalApiConfig._internal_api_endpoint = internal_api_endpoint
 
 
-def internal_api_call(func: Callable[PS, RT | None]) -> Callable[PS, RT | None]:
+def internal_api_call(func: Callable[PS, RT]) -> Callable[PS, RT]:
     """Decorator for methods which may be executed in database isolation mode.
 
     If [core]database_access_isolation is true then such method are not executed locally,
@@ -94,21 +93,20 @@ def make_jsonrpc_request(method_name: str, params_json: str) -> bytes:
         return response.content
 
     @wraps(func)
-    def wrapper(*args, **kwargs) -> RT | None:
+    def wrapper(*args, **kwargs) -> RT:
         use_internal_api = InternalApiConfig.get_use_internal_api()
         if not use_internal_api:
             return func(*args, **kwargs)
 
+        from airflow.serialization.serialized_objects import BaseSerialization  # avoid circular import
+
         bound = inspect.signature(func).bind(*args, **kwargs)
         arguments_dict = dict(bound.arguments)
         if "session" in arguments_dict:
             del arguments_dict["session"]
         args_json = json.dumps(BaseSerialization.serialize(arguments_dict))
-        method_name = f"{func.__module__}.{func.__name__}"
+        method_name = f"{func.__module__}.{func.__qualname__}"
         result = make_jsonrpc_request(method_name, args_json)
-        if result:
-            return BaseSerialization.deserialize(json.loads(result))
-        else:
-            return None
+        return BaseSerialization.deserialize(json.loads(result))
 
     return wrapper
diff --git a/airflow/models/dag.py b/airflow/models/dag.py
--- a/airflow/models/dag.py
+++ b/airflow/models/dag.py
@@ -62,6 +62,7 @@
 
 import airflow.templates
 from airflow import settings, utils
+from airflow.api_internal.internal_api_call import internal_api_call
 from airflow.compat.functools import cached_property
 from airflow.configuration import conf, secrets_backend_list
 from airflow.exceptions import (
@@ -3259,6 +3260,7 @@ def get_is_paused(self, *, session: Session | None = None) -> bool:
         return self.is_paused
 
     @staticmethod
+    @internal_api_call
     @provide_session
     def get_paused_dag_ids(dag_ids: list[str], session: Session = NEW_SESSION) -> set[str]:
         """


Test patch:
diff --git a/tests/api_internal/test_internal_api_call.py b/tests/api_internal/test_internal_api_call.py
--- a/tests/api_internal/test_internal_api_call.py
+++ b/tests/api_internal/test_internal_api_call.py
@@ -55,17 +55,17 @@ def test_get_use_internal_api_enabled(self):
         assert InternalApiConfig.get_internal_api_endpoint() == "http://localhost:8888/internal_api/v1/rpcapi"
 
 
-@internal_api_call
-def fake_method() -> str:
-    return "local-call"
-
-
-@internal_api_call
-def fake_method_with_params(dag_id: str, task_id: int) -> str:
-    return f"local-call-with-params-{dag_id}-{task_id}"
+class TestInternalApiCall:
+    @staticmethod
+    @internal_api_call
+    def fake_method() -> str:
+        return "local-call"
 
+    @staticmethod
+    @internal_api_call
+    def fake_method_with_params(dag_id: str, task_id: int) -> str:
+        return f"local-call-with-params-{dag_id}-{task_id}"
 
-class TestInternalApiCall:
     @conf_vars(
         {
             ("core", "database_access_isolation"): "false",
@@ -74,7 +74,7 @@ class TestInternalApiCall:
     )
     @mock.patch("airflow.api_internal.internal_api_call.requests")
     def test_local_call(self, mock_requests):
-        result = fake_method()
+        result = TestInternalApiCall.fake_method()
 
         assert result == "local-call"
         mock_requests.post.assert_not_called()
@@ -94,12 +94,12 @@ def test_remote_call(self, mock_requests):
 
         mock_requests.post.return_value = response
 
-        result = fake_method()
+        result = TestInternalApiCall.fake_method()
         assert result == "remote-call"
         expected_data = json.dumps(
             {
                 "jsonrpc": "2.0",
-                "method": "tests.api_internal.test_internal_api_call.fake_method",
+                "method": "tests.api_internal.test_internal_api_call.TestInternalApiCall.fake_method",
                 "params": json.dumps(BaseSerialization.serialize({})),
             }
         )
@@ -124,12 +124,13 @@ def test_remote_call_with_params(self, mock_requests):
 
         mock_requests.post.return_value = response
 
-        result = fake_method_with_params("fake-dag", task_id=123)
+        result = TestInternalApiCall.fake_method_with_params("fake-dag", task_id=123)
         assert result == "remote-call"
         expected_data = json.dumps(
             {
                 "jsonrpc": "2.0",
-                "method": "tests.api_internal.test_internal_api_call.fake_method_with_params",
+                "method": "tests.api_internal.test_internal_api_call.TestInternalApiCall."
+                "fake_method_with_params",
                 "params": json.dumps(
                     BaseSerialization.serialize(
                         {