Migrate DagFileProcessorManager._deactivate_stale_dags to Internal API

airflow/api_internal/endpoints/rpc_api_endpoint.py

@@ -23,6 +23,7 @@
 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
 
@@ -36,6 +37,7 @@ def _build_methods_map(list) -> dict:
 METHODS_MAP = _build_methods_map(
     [
         DagFileProcessor.update_import_errors,
+        DagFileProcessorManager.deactivate_stale_dags,
     ]
 )
 
@@ -68,7 +70,10 @@ def internal_airflow_api(
 
     log.debug("Calling method %.", {method_name})
     try:
-        output = handler(**params)
+        output = handler(
+            **params,
+            log=logging.getLogger(f"airflow.internal_api.{method_name}"),
+        )
         output_json = BaseSerialization.serialize(output)
         log.debug("Returning response")
         return 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
@@ -485,51 +486,69 @@ 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, session=None):
+        """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,
+                session=session,
             )
-            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 +613,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

@@ -18,6 +18,7 @@
 
 import json
 from unittest import mock
+from unittest.mock import ANY
 
 import pytest
 from flask import Flask
@@ -87,10 +88,7 @@ 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()
+        mock_test_method.assert_called_once_with(**(method_params or {}), log=ANY)
 
     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