Migrate DagFileProcessorManager._deactivate_stale_dags to Internal API

airflow/api_internal/endpoints/rpc_api_endpoint.py

@@ -19,10 +19,13 @@
 
 import json
 import logging
+from inspect import signature
+from typing import Callable
 
 from flask import Response
 
 from airflow.api_connexion.types import APIResponse
+from airflow.dag_processing.manager import DagFileProcessorManager
 from airflow.dag_processing.processor import DagFileProcessor
 from airflow.serialization.serialized_objects import BaseSerialization
 
@@ -33,9 +36,21 @@ def _build_methods_map(list) -> dict:
     return {f"{func.__module__}.{func.__name__}": func for func in list}
 
 
+def _in_parameters(func: Callable, parameter_name: str) -> bool:
+    """True if a parameter exists for a given function, False otherwise."""
+    func_params = signature(func).parameters
+    try:
+        # func_params is an ordered dict -- this is the "recommended" way of getting the position
+        tuple(func_params).index(parameter_name)
+        return True
+    except ValueError:
+        return False
+
+
 METHODS_MAP = _build_methods_map(
     [
         DagFileProcessor.update_import_errors,
+        DagFileProcessorManager.deactivate_stale_dags,
     ]
 )
 
@@ -68,6 +83,8 @@ def internal_airflow_api(
 
     log.debug("Calling method %.", {method_name})
     try:
+        if _in_parameters(handler, "log"):
+            params["log"] = logging.getLogger(f"airflow.internal_api.{handler.__name__}")
         output = handler(**params)
         output_json = BaseSerialization.serialize(output)
         log.debug("Returning response")

airflow/api_internal/internal_api_call.py

@@ -103,6 +103,8 @@ def wrapper(*args, **kwargs) -> RT | None:
         arguments_dict = dict(bound.arguments)
         if "session" in arguments_dict:
             del arguments_dict["session"]
+        if "log" in arguments_dict:
+            del arguments_dict["log"]
         args_json = json.dumps(BaseSerialization.serialize(arguments_dict))
         method_name = f"{func.__module__}.{func.__name__}"
         result = make_jsonrpc_request(method_name, args_json)

airflow/dag_processing/manager.py

@@ -42,6 +42,7 @@
 from tabulate import tabulate
 
 import airflow.models
+from airflow.api_internal.internal_api_call import internal_api_call
 from airflow.callbacks.callback_requests import CallbackRequest, SlaCallbackRequest
 from airflow.configuration import conf
 from airflow.dag_processing.processor import DagFileProcessorProcess
@@ -484,52 +485,68 @@ def start(self):
 
         return self._run_parsing_loop()
 
-    @provide_session
-    def _deactivate_stale_dags(self, session=None):
-        """
-        Detects DAGs which are no longer present in files.
-
-        Deactivate them and remove them in the serialized_dag table
-        """
+    def _scan_stale_dags(self):
+        """Scan at fix internal DAGs which are no longer present in files."""
         now = timezone.utcnow()
         elapsed_time_since_refresh = (now - self.last_deactivate_stale_dags_time).total_seconds()
         if elapsed_time_since_refresh > self.parsing_cleanup_interval:
             last_parsed = {
                 fp: self.get_last_finish_time(fp) for fp in self.file_paths if self.get_last_finish_time(fp)
             }
-            to_deactivate = set()
-            query = session.query(DagModel.dag_id, DagModel.fileloc, DagModel.last_parsed_time).filter(
-                DagModel.is_active
+            DagFileProcessorManager.deactivate_stale_dags(
+                last_parsed=last_parsed,
+                dag_directory=self.get_dag_directory(),
+                processor_timeout=self._processor_timeout,
+                log=self.log,
             )
-            if self.standalone_dag_processor:
-                query = query.filter(DagModel.processor_subdir == self.get_dag_directory())
-            dags_parsed = query.all()
-
-            for dag in dags_parsed:
-                # The largest valid difference between a DagFileStat's last_finished_time and a DAG's
-                # last_parsed_time is _processor_timeout. Longer than that indicates that the DAG is
-                # no longer present in the file.
-                if (
-                    dag.fileloc in last_parsed
-                    and (dag.last_parsed_time + self._processor_timeout) < last_parsed[dag.fileloc]
-                ):
-                    self.log.info("DAG %s is missing and will be deactivated.", dag.dag_id)
-                    to_deactivate.add(dag.dag_id)
-
-            if to_deactivate:
-                deactivated = (
-                    session.query(DagModel)
-                    .filter(DagModel.dag_id.in_(to_deactivate))
-                    .update({DagModel.is_active: False}, synchronize_session="fetch")
-                )
-                if deactivated:
-                    self.log.info("Deactivated %i DAGs which are no longer present in file.", deactivated)
+            self.last_deactivate_stale_dags_time = timezone.utcnow()
 
-                for dag_id in to_deactivate:
-                    SerializedDagModel.remove_dag(dag_id)
-                    self.log.info("Deleted DAG %s in serialized_dag table", dag_id)
+    @staticmethod
+    @internal_api_call
+    @provide_session
+    def deactivate_stale_dags(
+        last_parsed: dict[str, datetime | None],
+        dag_directory: str,
+        processor_timeout: timedelta,
+        log: logging.Logger,
+        session: Session = NEW_SESSION,
+    ):
+        """
+        Detects DAGs which are no longer present in files.
+        Deactivate them and remove them in the serialized_dag table
+        """
+        to_deactivate = set()
+        query = session.query(DagModel.dag_id, DagModel.fileloc, DagModel.last_parsed_time).filter(
+            DagModel.is_active
+        )
+        standalone_dag_processor = conf.getboolean("scheduler", "standalone_dag_processor")
+        if standalone_dag_processor:
+            query = query.filter(DagModel.processor_subdir == dag_directory)
+        dags_parsed = query.all()
+
+        for dag in dags_parsed:
+            # The largest valid difference between a DagFileStat's last_finished_time and a DAG's
+            # last_parsed_time is _processor_timeout. Longer than that indicates that the DAG is
+            # no longer present in the file.
+            if (
+                dag.fileloc in last_parsed
+                and (dag.last_parsed_time + processor_timeout) < last_parsed[dag.fileloc]
+            ):
+                log.info("DAG %s is missing and will be deactivated.", dag.dag_id)
+                to_deactivate.add(dag.dag_id)
+
+        if to_deactivate:
+            deactivated = (
+                session.query(DagModel)
+                .filter(DagModel.dag_id.in_(to_deactivate))
+                .update({DagModel.is_active: False}, synchronize_session="fetch")
+            )
+            if deactivated:
+                log.info("Deactivated %i DAGs which are no longer present in file.", deactivated)
 
-            self.last_deactivate_stale_dags_time = timezone.utcnow()
+            for dag_id in to_deactivate:
+                SerializedDagModel.remove_dag(dag_id)
+                log.info("Deleted DAG %s in serialized_dag table", dag_id)
 
     def _run_parsing_loop(self):
         # In sync mode we want timeout=None -- wait forever until a message is received
@@ -594,7 +611,7 @@ def _run_parsing_loop(self):
 
             if self.standalone_dag_processor:
                 self._fetch_callbacks(max_callbacks_per_loop)
-            self._deactivate_stale_dags()
+            self._scan_stale_dags()
             DagWarning.purge_inactive_dag_warnings()
             refreshed_dag_dir = self._refresh_dag_dir()
 

tests/api_internal/endpoints/test_rpc_api_endpoint.py

@@ -17,6 +17,7 @@
 from __future__ import annotations
 
 import json
+import logging
 from unittest import mock
 
 import pytest
@@ -29,8 +30,14 @@
 from tests.test_utils.decorators import dont_initialize_flask_app_submodules
 
 TEST_METHOD_NAME = "test_method"
+TEST_METHOD_WITH_LOG_NAME = "test_method_with_log"
 
 mock_test_method = mock.MagicMock()
+mock_test_method_with_log = mock.MagicMock()
+
+
+def method_with_log(*args, log, **kwargs):
+    return mock_test_method_with_log(*args, **kwargs, log=log)
 
 
 @pytest.fixture(scope="session")
@@ -52,16 +59,33 @@ class TestRpcApiEndpoint:
     @pytest.fixture(autouse=True)
     def setup_attrs(self, minimal_app_for_internal_api: Flask) -> None:
         rpc_api_endpoint.METHODS_MAP[TEST_METHOD_NAME] = mock_test_method
+        rpc_api_endpoint.METHODS_MAP[TEST_METHOD_WITH_LOG_NAME] = method_with_log
         self.app = minimal_app_for_internal_api
         self.client = self.app.test_client()  # type:ignore
         mock_test_method.reset_mock()
         mock_test_method.side_effect = None
+        mock_test_method_with_log.reset_mock()
+        mock_test_method_with_log.side_effect = None
 
     @pytest.mark.parametrize(
-        "input_data, method_result, method_params, expected_code",
+        "input_data, method_result, method_params, expected_logger, expected_mock, expected_code",
         [
-            ({"jsonrpc": "2.0", "method": TEST_METHOD_NAME, "params": ""}, "test_me", None, 200),
-            ({"jsonrpc": "2.0", "method": TEST_METHOD_NAME, "params": ""}, None, None, 200),
+            (
+                {"jsonrpc": "2.0", "method": TEST_METHOD_NAME, "params": ""},
+                "test_me",
+                {},
+                False,
+                mock_test_method,
+                200,
+            ),
+            (
+                {"jsonrpc": "2.0", "method": TEST_METHOD_NAME, "params": ""},
+                None,
+                {},
+                False,
+                mock_test_method,
+                200,
+            ),
             (
                 {
                     "jsonrpc": "2.0",
@@ -70,13 +94,25 @@ def setup_attrs(self, minimal_app_for_internal_api: Flask) -> None:
                 },
                 ("dag_id_15", "fake-task", 1),
                 {"dag_id": 15, "task_id": "fake-task"},
+                False,
+                mock_test_method,
+                200,
+            ),
+            (
+                {"jsonrpc": "2.0", "method": TEST_METHOD_WITH_LOG_NAME, "params": ""},
+                "test_me",
+                {},
+                True,
+                mock_test_method_with_log,
                 200,
             ),
         ],
     )
-    def test_method(self, input_data, method_result, method_params, expected_code):
+    def test_method(
+        self, input_data, method_result, method_params, expected_logger, expected_mock, expected_code
+    ):
         if method_result:
-            mock_test_method.return_value = method_result
+            expected_mock.return_value = method_result
 
         response = self.client.post(
             "/internal_api/v1/rpcapi",
@@ -87,10 +123,15 @@ def test_method(self, input_data, method_result, method_params, expected_code):
         if method_result:
             response_data = BaseSerialization.deserialize(json.loads(response.data))
             assert response_data == method_result
-        if method_params:
-            mock_test_method.assert_called_once_with(**method_params)
-        else:
-            mock_test_method.assert_called_once()
+
+        if expected_logger:
+            args, kwargs = expected_mock.call_args
+            logger = kwargs["log"]
+            assert isinstance(logger, logging.Logger)
+            assert logger.name == "airflow.internal_api.method_with_log"
+            method_params["log"] = logger
+
+        expected_mock.assert_called_once_with(**method_params)
 
     def test_method_with_exception(self):
         mock_test_method.side_effect = ValueError("Error!!!")

tests/api_internal/test_internal_api_call.py

@@ -61,7 +61,7 @@ def fake_method() -> str:
 
 
 @internal_api_call
-def fake_method_with_params(dag_id: str, task_id: int) -> str:
+def fake_method_with_params(dag_id: str, task_id: int, session, log) -> str:
     return f"local-call-with-params-{dag_id}-{task_id}"
 
 
@@ -124,7 +124,7 @@ 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 = fake_method_with_params("fake-dag", task_id=123, session="session", log="log")
         assert result == "remote-call"
         expected_data = json.dumps(
             {

tests/dag_processing/test_manager.py

@@ -495,7 +495,7 @@ def test_recently_modified_file_is_parsed_with_mtime_mode(
                 > (freezed_base_time - manager.get_last_finish_time("file_1.py")).total_seconds()
             )
 
-    def test_deactivate_stale_dags(self):
+    def test_scan_stale_dags(self):
         """
         Ensure that DAGs are marked inactive when the file is parsed but the
         DagModel.last_parsed_time is not updated.
@@ -545,7 +545,7 @@ def test_deactivate_stale_dags(self):
             )
             assert serialized_dag_count == 1
 
-            manager._deactivate_stale_dags()
+            manager._scan_stale_dags()
 
             active_dag_count = (
                 session.query(func.count(DagModel.dag_id))
@@ -567,7 +567,7 @@ def test_deactivate_stale_dags(self):
             ("scheduler", "standalone_dag_processor"): "True",
         }
     )
-    def test_deactivate_stale_dags_standalone_mode(self):
+    def test_scan_stale_dags_standalone_mode(self):
         """
         Ensure only dags from current dag_directory are updated
         """
@@ -612,7 +612,7 @@ def test_deactivate_stale_dags_standalone_mode(self):
             active_dag_count = session.query(func.count(DagModel.dag_id)).filter(DagModel.is_active).scalar()
             assert active_dag_count == 2
 
-            manager._deactivate_stale_dags()
+            manager._scan_stale_dags()
 
             active_dag_count = session.query(func.count(DagModel.dag_id)).filter(DagModel.is_active).scalar()
             assert active_dag_count == 1