kaxil commented on code in PR #74793:
URL: https://github.com/apache/airflow/pull/74793#discussion_r4401154351
##########
airflow-core/src/airflow/dag_processing/manager.py:
##########
@@ -390,13 +486,21 @@ def before_run(self) -> None:
gc.freeze()
def after_run(self) -> None:
- """Tear down state after the parsing loop exits. Default no-op;
override to customize."""
+ """Tear down state after the parsing loop exits. Default
implementation; override to customize."""
+ if self.api_client is None:
+ return
+ from airflow.dag_processing.api_client import DagProcessorSecretsComms
+ from airflow.sdk.execution_time import task_runner
Review Comment:
processor.py already imports from `task_runner` at module level and this
module imports processor, so this import (and the one at 534) can go at the
top. The `api_client` imports are lazy on purpose per
`test_direct_manager_does_not_import_api_client`, so a one-line comment saying
that would stop someone hoisting them.
##########
airflow-core/src/airflow/dag_processing/manager.py:
##########
@@ -848,6 +985,11 @@ def get_bundle_state(self, bundle_name: str, *, session:
Session = NEW_SESSION)
Returns ``None`` if the bundle has no database record.
"""
+ if self.api_client is not None:
Review Comment:
Is this branch reachable? The only caller (line 1145) calls
`get_bundle_state` only when `catalog is None`, and in API mode line 1129 has
already `continue`d in that case. The `except httpx2.HTTPError` around
`catalog.get` at 1147 can't fire either.
##########
airflow-core/src/airflow/dag_processing/manager.py:
##########
@@ -902,157 +1095,229 @@ def _refresh_dag_bundles(self, known_files: dict[str,
set[DagFileInfo]]):
if periodic_check_due:
self._bundles_last_refreshed = now_seconds
+ catalog: dict[str, BundleState] | None = None
+ if self.api_client is not None and self._dag_bundles:
+ try:
+ bundles = self.api_client.get_bundles()
+ catalog = {
+ bundle.name: BundleState(bundle.last_refreshed,
bundle.version, bundle.revision)
+ for bundle in bundles
+ }
+ self._bundle_name_to_team_name = {bundle.name:
bundle.team_name for bundle in bundles}
+ except httpx2.HTTPError:
+ self.log.warning("Unable to read the bundle catalog; retrying
on a later refresh")
+
any_refreshed = False
- active_bundles = {file.bundle_name for file in self._processors if
file.bundle_version is None}
+ active_bundles = {
+ file.bundle_name
+ for file in self._processors.keys() |
self._pending_publications.keys()
+ if file.bundle_version is None
+ }
for bundle in self._dag_bundles:
if not periodic_check_due and bundle.name not in urgent_bundles:
continue
- # TODO: AIP-66 handle errors in the case of incomplete cloning?
And test this.
- # What if the cloning/refreshing took too long(longer than the
dag processor timeout)
- if not bundle.is_initialized:
- self._invalidate_bundle_parse_source(bundle.name)
+ with self._use_bundle(bundle.name):
+ if catalog is not None and bundle.name not in catalog:
+ self._discard_unusable_bundle_source(bundle.name)
Review Comment:
When a bundle is missing from the catalog the API returns, this stops
parsing it with no log line, and does so again on every check. A one-time
warning naming the bundle would tell an operator why its Dags stopped updating.
##########
airflow-core/src/airflow/dag_processing/manager.py:
##########
@@ -902,157 +1095,229 @@ def _refresh_dag_bundles(self, known_files: dict[str,
set[DagFileInfo]]):
if periodic_check_due:
self._bundles_last_refreshed = now_seconds
+ catalog: dict[str, BundleState] | None = None
+ if self.api_client is not None and self._dag_bundles:
+ try:
+ bundles = self.api_client.get_bundles()
+ catalog = {
+ bundle.name: BundleState(bundle.last_refreshed,
bundle.version, bundle.revision)
+ for bundle in bundles
+ }
+ self._bundle_name_to_team_name = {bundle.name:
bundle.team_name for bundle in bundles}
+ except httpx2.HTTPError:
+ self.log.warning("Unable to read the bundle catalog; retrying
on a later refresh")
+
any_refreshed = False
- active_bundles = {file.bundle_name for file in self._processors if
file.bundle_version is None}
+ active_bundles = {
+ file.bundle_name
+ for file in self._processors.keys() |
self._pending_publications.keys()
+ if file.bundle_version is None
+ }
for bundle in self._dag_bundles:
if not periodic_check_due and bundle.name not in urgent_bundles:
continue
- # TODO: AIP-66 handle errors in the case of incomplete cloning?
And test this.
- # What if the cloning/refreshing took too long(longer than the
dag processor timeout)
- if not bundle.is_initialized:
- self._invalidate_bundle_parse_source(bundle.name)
+ with self._use_bundle(bundle.name):
+ if catalog is not None and bundle.name not in catalog:
+ self._discard_unusable_bundle_source(bundle.name)
+ continue
+ if bundle.name in self._pending_inventories:
+ # A listing still waiting for the API must be published
before the bundle is listed again.
+ if bundle.name in due_inventories:
+ self._publish_bundle_inventory(bundle.name,
known_files)
+ any_refreshed = True
+ continue
+ if self.api_client is not None and catalog is None:
+ continue
+ # TODO: AIP-66 handle errors in the case of incomplete
cloning? And test this.
+ # What if the cloning/refreshing took too long(longer than
the dag processor timeout)
+ if not bundle.is_initialized:
+ self._invalidate_bundle_parse_source(bundle.name)
+ try:
+ bundle.initialize()
+ any_refreshed = True
+ except Exception as e:
+ self.log.exception("Error initializing bundle %s: %s",
bundle.name, e)
+ continue
try:
- bundle.initialize()
- any_refreshed = True
- except Exception as e:
- self.log.exception("Error initializing bundle %s: %s",
bundle.name, e)
+ bundle_state = (
+ catalog.get(bundle.name)
+ if catalog is not None
+ else self.get_bundle_state(bundle.name)
+ )
+ except httpx2.HTTPError:
+ self.log.warning("Unable to read metadata for bundle %s;
retrying later", bundle.name)
+ continue
+ except Exception:
+ self.log.exception("Error fetching state for bundle %s",
bundle.name)
+ continue
+ if bundle_state is None:
+ self.log.warning("Bundle model not found for %s",
bundle.name)
+ continue
+ elapsed_time_since_refresh = (
+ now - (bundle_state.last_refreshed or utc_epoch())
+ ).total_seconds()
+ if bundle.supports_versioning:
+ # we will also check the version of the bundle to see if
another DAG processor has seen
+ # a new version
+ pre_refresh_version =
self._bundle_versions.get(bundle.name)
+ # Use `is None` (not falsy) so an empty-string version is
treated as a valid cached value.
+ if pre_refresh_version is None:
+ pre_refresh_version, _ =
unpack_bundle_version(bundle.get_current_version(), bundle)
+ current_version_matches_db = pre_refresh_version ==
bundle_state.version
+ else:
+ # With no versioning, it always "matches"
+ current_version_matches_db = True
+
+ previously_seen = bundle.name in self._bundle_versions
+ if (
+ bundle.name not in self._bundles_waiting_for_refresh
+ and bundle.name not in self._bundles_needing_refresh
+ and bundle.name in self._bundle_parse_sources
+ and self.should_skip_refresh(
+ bundle=bundle,
+ elapsed_time_since_refresh=elapsed_time_since_refresh,
+ current_version_matches_db=current_version_matches_db,
+ previously_seen=previously_seen,
+ )
+ ):
+ self.log.debug("Not time to refresh bundle %s",
bundle.name)
continue
- try:
- bundle_state = self.get_bundle_state(bundle.name)
- except Exception:
- self.log.exception("Error fetching state for bundle %s",
bundle.name)
- continue
- if bundle_state is None:
- self.log.warning("Bundle model not found for %s", bundle.name)
- continue
- elapsed_time_since_refresh = (now - (bundle_state.last_refreshed
or utc_epoch())).total_seconds()
- if bundle.supports_versioning:
- # we will also check the version of the bundle to see if
another DAG processor has seen
- # a new version
- pre_refresh_version = self._bundle_versions.get(bundle.name)
- # Use `is None` (not falsy) so an empty-string version is
treated as a valid cached value.
- if pre_refresh_version is None:
- pre_refresh_version, _ =
unpack_bundle_version(bundle.get_current_version(), bundle)
- current_version_matches_db = pre_refresh_version ==
bundle_state.version
- else:
- # With no versioning, it always "matches"
- current_version_matches_db = True
-
- previously_seen = bundle.name in self._bundle_versions
- if (
- bundle.name not in self._bundles_waiting_for_refresh
- and bundle.name not in self._bundles_needing_refresh
- and bundle.name in self._bundle_parse_sources
- and self.should_skip_refresh(
- bundle=bundle,
- elapsed_time_since_refresh=elapsed_time_since_refresh,
- current_version_matches_db=current_version_matches_db,
- previously_seen=previously_seen,
- )
- ):
- self.log.debug("Not time to refresh bundle %s", bundle.name)
- continue
- if bundle.refresh_changes_source and bundle.name in active_bundles:
- # Forced refreshes wait too: refreshing under an import would
discard its result.
- self._bundles_waiting_for_refresh.add(bundle.name)
- continue
+ if bundle.refresh_changes_source and bundle.name in
active_bundles:
+ # Forced refreshes wait too: refreshing under an import
would discard its result.
+ self._bundles_waiting_for_refresh.add(bundle.name)
+ continue
- self._bundles_waiting_for_refresh.discard(bundle.name)
- self.log.info("Refreshing bundle %s", bundle.name)
+ self._bundles_waiting_for_refresh.discard(bundle.name)
+ self.log.info("Refreshing bundle %s", bundle.name)
- previous_source = self._bundle_parse_sources.get(bundle.name)
- if bundle.refresh_changes_source:
- self._invalidate_bundle_parse_source(bundle.name)
- try:
- bundle.refresh()
- any_refreshed = True
- except Exception as error:
- self.log.exception("Error refreshing bundle %s", bundle.name)
- source_unchanged = not bundle.refresh_changes_source or (
- isinstance(error, DagBundleRefreshError) and
error.source_unchanged
- )
- if source_unchanged and previous_source is not None:
- self._bundle_parse_sources[bundle.name] = attrs.evolve(
- previous_source,
-
refresh_generation=self._bundle_refresh_generations[bundle.name],
+ previous_source = self._bundle_parse_sources.get(bundle.name)
+ if bundle.refresh_changes_source:
+ self._invalidate_bundle_parse_source(bundle.name)
+ try:
+ bundle.refresh()
+ any_refreshed = True
+ except Exception as error:
+ self.log.exception("Error refreshing bundle %s",
bundle.name)
+ source_unchanged = not bundle.refresh_changes_source or (
+ isinstance(error, DagBundleRefreshError) and
error.source_unchanged
)
- self._force_refresh_bundles.discard(bundle.name)
- if bundle.name not in self._bundles_needing_refresh:
+ if source_unchanged and previous_source is not None:
+ self._bundle_parse_sources[bundle.name] = attrs.evolve(
+ previous_source,
+
refresh_generation=self._bundle_refresh_generations[bundle.name],
+ )
+ self._force_refresh_bundles.discard(bundle.name)
+ if bundle.name not in self._bundles_needing_refresh:
+ self.log.warning(
+ "Refresh failed without changing the source;
continuing to parse bundle",
+ bundle_name=bundle.name,
+ )
+ elif previous_source is not None or bundle.name not in
self._bundles_needing_refresh:
self.log.warning(
- "Refresh failed without changing the source;
continuing to parse bundle",
- bundle_name=bundle.name,
+ "Parsing paused until bundle refresh succeeds",
bundle_name=bundle.name
)
- elif previous_source is not None or bundle.name not in
self._bundles_needing_refresh:
- self.log.warning("Parsing paused until bundle refresh
succeeds", bundle_name=bundle.name)
- self._bundles_needing_refresh.add(bundle.name)
- continue
+ self._bundles_needing_refresh.add(bundle.name)
+ continue
+
+ self._bundles_needing_refresh.discard(bundle.name)
+ self._force_refresh_bundles.discard(bundle.name)
- self._bundles_needing_refresh.discard(bundle.name)
- self._force_refresh_bundles.discard(bundle.name)
+ if bundle.supports_versioning:
+ version_after_refresh, version_data_after_refresh =
unpack_bundle_version(
+ bundle.get_current_version(), bundle
+ )
+ else:
+ version_after_refresh = None
+ version_data_after_refresh = None
- if bundle.supports_versioning:
- version_after_refresh, version_data_after_refresh =
unpack_bundle_version(
- bundle.get_current_version(), bundle
+ self._bundle_parse_sources[bundle.name] = DagParseSource(
+ bundle_version=version_after_refresh,
+ version_data=deepcopy(version_data_after_refresh),
+
refresh_generation=self._bundle_refresh_generations[bundle.name],
)
- else:
- version_after_refresh = None
- version_data_after_refresh = None
- self._bundle_parse_sources[bundle.name] = DagParseSource(
- bundle_version=version_after_refresh,
- version_data=deepcopy(version_data_after_refresh),
-
refresh_generation=self._bundle_refresh_generations[bundle.name],
- )
+ if bundle.supports_versioning:
+ # We can short-circuit the rest of this if (1) bundle was
seen before by
+ # this dag processor and (2) the version of the bundle did
not change
+ # after refreshing it
+ if (
+ self.api_client is None
+ and previously_seen
+ and pre_refresh_version == version_after_refresh
+ ):
+ self.log.debug(
+ "Bundle %s version not changed after refresh: %s",
+ bundle.name,
+ version_after_refresh,
+ )
+ try:
+ self.update_bundle_state(bundle.name,
last_refreshed=now, version=None)
+ except Exception:
+ self.log.exception("Error persisting state for
bundle %s", bundle.name)
+ continue
- if bundle.supports_versioning:
- # We can short-circuit the rest of this if (1) bundle was seen
before by
- # this dag processor and (2) the version of the bundle did not
change
- # after refreshing it
- if previously_seen and pre_refresh_version ==
version_after_refresh:
- self.log.debug(
- "Bundle %s version not changed after refresh: %s",
- bundle.name,
- version_after_refresh,
+ self.log.info(
+ "Version changed for %s, new version: %s",
bundle.name, version_after_refresh
)
- try:
- self.update_bundle_state(bundle.name,
last_refreshed=now, version=None)
- except Exception:
- self.log.exception("Error persisting state for bundle
%s", bundle.name)
+ try:
+ discovery = self._find_files_in_bundle(bundle)
+ except Exception:
+ # Keep the bundle's known files and Dags, and leave its
version unadvanced so the
+ # next refresh lists it again.
+ self.log.exception("Error listing Dag definitions in
bundle %s", bundle.name)
+ if self.api_client is not None:
+ self._bundle_parse_sources.pop(bundle.name, None)
continue
- self.log.info("Version changed for %s, new version: %s",
bundle.name, version_after_refresh)
- try:
- discovery = self._find_files_in_bundle(bundle)
- except Exception:
- # Keep the bundle's known files and Dags, and leave its
version unadvanced so the
- # next refresh lists it again.
- self.log.exception("Error listing Dag definitions in bundle
%s", bundle.name)
- continue
+ found_files = discovery.files
+ if self.api_client is not None:
+ self._dispatch_sequence += 1
+ try:
+ inventory = DagBundleInventoryBody(
+ attempt_id=uuid7(),
+ dispatch_sequence=self._dispatch_sequence,
+ expected_revision=bundle_state.revision,
+ version=version_after_refresh,
+
files=sorted(self._get_observed_filelocs(found_files)),
+ failed_sources=sorted(discovery.failed_sources),
+ )
+ except ValidationError:
+ # Handled like a failed listing, so the API never sees
a partial inventory.
+ self.log.exception("Bundle %s has a listing the API
cannot accept", bundle.name)
+ self._discard_unusable_bundle_source(bundle.name)
Review Comment:
One path the inventory model rejects (a backslash in a filename, which is
legal on Linux, or a path over 2000 characters) takes the whole bundle out
here, and so does a 413 at line 1049 from a proxy whose body limit is below 16
MiB. With the parse source popped, `should_skip_refresh` is never consulted, so
every `bundle_refresh_check_interval` (5s by default) re-runs
`bundle.refresh()` (a git fetch for Git bundles), lists again and fails again,
while every file in the bundle waits for a source. DB mode would just parse the
rest of the bundle. Could the unrepresentable paths be dropped and logged by
name instead, and a 413 back off through `next_attempt_time`? The 413 branch
also logs the same "rejected" line as a 409, so including the status there
would help.
##########
airflow-core/src/airflow/dag_processing/manager.py:
##########
@@ -448,6 +574,8 @@ def load_dag_bundles(self) -> None:
def _scan_stale_dags(self):
"""Scan and deactivate DAGs which are no longer present in files."""
+ if self.api_client is not None:
Review Comment:
With this early return, what ages out the Dags of a file whose parse now
times out or exits without a result? Neither path publishes anything in API
mode (test_manager_api_results.py asserts no publish for the `failed` case),
and the server only marks earlier definitions stale when it gets an import
error. A file that starts hanging keeps its last good Dags scheduled with no
import error shown. Could those two paths send an import error through
`_build_rejection_report`, so the server's cutoff applies like
`stale_dag_threshold` does in DB mode?
##########
airflow-core/tests/unit/dag_processing/test_manager.py:
##########
@@ -5673,6 +5742,275 @@ def
test_a_sweep_pays_the_fixed_cost_once_per_call(self, session, testing_dag_bu
f"({SWEEP_CALLS} x {fixed} fixed + {SWEEP_FILES} x {per_dag} per
Dag)."
)
+ @pytest.mark.parametrize("publication_fails", [False, True])
+ @mock.patch.object(DagFileProcessorManager, "persist_parsing_result",
autospec=True)
+ @mock.patch.object(DagFileProcessorManager, "_create_process",
autospec=True)
+ def test_api_publication_uses_dispatch_identity_without_local_persistence(
+ self, create, persist, session, publication_fails
+ ):
+ client = mock.create_autospec(DagProcessorAPIClient, instance=True)
+ manager = DagFileProcessorManager(max_runs=1, api_client=client)
+ manager._bundle_parse_sources["testing"] =
DagParseSource(bundle_version="v1")
+ file = _get_file_infos(["dag.py"])[0]
+ manager._file_queue[file] = None
+ proc, read_end = self.mock_processor()
+ proc.had_callbacks = False
+ proc.parsing_result = DagFileParsingResult(fileloc="dag.py",
serialized_dags=[])
+ proc.parse_source = manager._bundle_parse_sources["testing"]
+ create.return_value = proc
+ if publication_fails:
+ client.publish_parse_result.side_effect =
RuntimeError("publication failed")
+ try:
+ manager._start_new_processes()
+ create.assert_called_once_with(manager, file,
source=proc.parse_source)
+ manager.handle_parsing_result(file, proc, session=session)
+ client.publish_parse_result.assert_not_called()
+ manager._publish_pending_results()
+ body = client.publish_parse_result.call_args.args[0]
+ assert body.attempt_id == proc.id
+ assert body.dispatch_sequence == proc.dispatch_sequence == 1
+ assert body.bundle_version == "v1"
+ assert body.relative_fileloc == "dag.py"
+ assert body.serialized_dags == []
+ assert manager._file_stats[file].run_count == 1
+ persist.assert_not_called()
+ finally:
+ read_end.close()
+ proc.close()
+
+ def test_api_publication_reports_a_result_the_api_would_refuse(self,
session):
+ client = mock.create_autospec(DagProcessorAPIClient, instance=True)
+ manager = DagFileProcessorManager(max_runs=1, api_client=client)
+ file = _get_file_infos(["dag.py"])[0]
+ dag = SdkDAG("published", schedule=None)
+ dag.fileloc = "/worker/dag.py"
+ dag.relative_fileloc = "dag.py"
+ # Mutating the set skips the SDK's constructor check, so only the API
model rejects the tag.
+ dag.tags.add("x" * 101)
+ revision = uuid7()
+ proc, read_end = self.mock_processor()
+ proc.had_callbacks = False
+ proc.dispatch_sequence = 1
+ proc.parse_source = DagParseSource(bundle_version="v1",
bundle_revision=revision)
+ proc.parsing_result = DagFileParsingResult(
+ fileloc=dag.fileloc,
serialized_dags=[LazyDeserializedDAG.from_dag(dag)]
+ )
+ try:
+ manager.handle_parsing_result(file, proc, session=session)
+ manager._publish_pending_results()
+ finally:
+ read_end.close()
+ proc.close()
+
+ report = client.publish_parse_result.call_args.args[0]
+ assert report.serialized_dags == []
+ assert (report.bundle_version, report.bundle_revision) == ("v1",
revision)
+ assert report.import_errors == {
+ "dag.py": "The parsing result was not published: it is not a valid
Dag publication "
+ "(Value error, Dag tags must be strings of at most 100
characters)."
+ }
+ assert manager._file_stats[file].import_errors == 1
+ assert not manager._pending_publications
+
+ @pytest.mark.parametrize("source_available", [True, False])
+ def
test_api_publication_carries_diagnostics_and_explicit_source_availability(self,
source_available):
+ client = mock.create_autospec(DagProcessorAPIClient, instance=True)
+ manager = DagFileProcessorManager(max_runs=1, api_client=client)
+ file = _get_file_infos(["dag.py"])[0]
+ dag = SdkDAG("published", schedule=None)
+ dag.fileloc = "/worker/dag.py"
+ dag.relative_fileloc = "dag.py"
+ proc, read_end = self.mock_processor()
+ proc.dispatch_sequence = 1
+ warning = {"dag_id": dag.dag_id, "warning_type": "python:deprecated",
"message": "warning"}
+ unserialized_warning = {"dag_id": "failed", "warning_type":
"python:deprecated", "message": "warning"}
+ proc.parsing_result = DagFileParsingResult(
+ fileloc=dag.fileloc,
+ serialized_dags=[LazyDeserializedDAG.from_dag(dag)],
+ warnings=[warning, unserialized_warning],
+ parsed_definitions=["dag.py"],
+ import_errors={"dag.py": "import error"},
+ dag_source_codes={dag.dag_id:
DagSourceCode(source_code="captured", language="python")}
+ if source_available
+ else {},
+ )
+ try:
+ body = manager._build_parse_result(file, proc, 0.5)
+ assert [published.model_dump() for published in body.warnings] ==
[warning]
+ assert body.import_errors == {"dag.py": "import error"}
+ assert body.parsed_definitions == ["dag.py"]
+ assert body.source_codes[dag.dag_id].source_code == ("captured" if
source_available else None)
+ finally:
+ read_end.close()
+ proc.close()
+
+ @pytest.mark.usefixtures("testing_dag_bundle")
+ @conf_vars({("workers", "execution_api_retries"): "1"})
+ @pytest.mark.parametrize("failure", [413, 422, 403, 500, "timeout"])
+ def test_rejected_publication_preserves_stale_detection(self, session,
time_machine, failure):
+ time_machine.move_to("2026-10-05T12:00:00Z", tick=False)
+ client = mock.create_autospec(DagProcessorAPIClient, instance=True)
+ request = httpx2.Request("POST", "http://api/jobs/1/parse-results";)
+ client.publish_parse_result.side_effect = (
+ httpx2.ReadTimeout("Lost response", request=request)
+ if failure == "timeout"
+ else httpx2.HTTPStatusError("Rejected", request=request,
response=httpx2.Response(failure))
+ )
+ manager = DagFileProcessorManager(max_runs=1, api_client=client)
+ file = _get_file_infos(["dag.py"])[0]
+ previous = timezone.utcnow() -
timedelta(seconds=manager.stale_dag_threshold + 1)
+ manager._file_stats[file] = DagFileStat(num_dags=1,
last_finish_time=previous)
+ dag = SdkDAG("published", schedule=None)
+ dag.fileloc = "/worker/dag.py"
+ dag.relative_fileloc = "dag.py"
+ session.add(
+ DagModel(
+ dag_id=dag.dag_id,
+ bundle_name="testing",
+ relative_fileloc="dag.py",
+ fileloc=dag.fileloc,
+ last_parsed_time=previous,
+ is_stale=False,
+ )
+ )
+ session.commit()
+ proc, read_end = self.mock_processor()
+ proc.dispatch_sequence = 1
+ proc.parsing_result = DagFileParsingResult(
+ fileloc=dag.fileloc,
serialized_dags=[LazyDeserializedDAG.from_dag(dag)]
+ )
+ try:
+ manager.handle_parsing_result(file, proc)
+ manager._publish_pending_results()
+ manager._scan_stale_dags()
Review Comment:
`_scan_stale_dags()` returns immediately in API mode, so the `is_stale`
assertion below passes whatever the publication does. Either drop the call and
that assertion, or assert something the API path controls.
##########
airflow-core/src/airflow/dag_processing/manager.py:
##########
@@ -103,15 +114,36 @@
from sqlalchemy.sql import Select
from airflow.callbacks.callback_requests import CallbackRequest
+ from airflow.dag_processing.api_client import DagProcessorAPIClient
from airflow.dag_processing.bundles.base import BaseDagBundle
from airflow.sdk.api.client import Client
+# Each send has its own short timeout; this bounds how long one loop spends on
queued sends
Review Comment:
`_get_write_timeout` only caps connect and pool at 1s; read and write keep
the client's full timeout. So one slow publish can hold the loop well past this
1s budget. Maybe say the budget limits when a new send starts, not how long a
started one runs.
##########
airflow-core/tests/unit/dag_processing/test_manager.py:
##########
@@ -6069,3 +6407,859 @@ def
test_normalized_file_path_for_stats_does_not_warn(caplog):
assert result == "dags_test_test_dag.py"
assert caplog.entries == []
+
+
+API_CLIENT_BUNDLES = [
+ {
+ "name": name,
+ "classpath": "airflow.dag_processing.bundles.local.LocalDagBundle",
+ "kwargs": {"path": "/"},
+ }
+ for name in ("bundle_a", "bundle_b")
+]
+
+
+def _read_bundle_connection() -> None:
+ Connection.get("bundle_conn")
+
+
+class TestDagFileProcessorManagerWithAPIClient:
+ """With an API client, bundle code resolves its connections through it,
for the bundle it belongs to."""
+
+ @pytest.fixture
+ def api_requests(self) -> list[httpx2.Request]:
+ return []
+
+ @pytest.fixture
+ def api_client(self, tmp_path, api_requests):
+ token_file = tmp_path / "processor.jwt"
+ token_file.write_text("session-token")
+ job_token = jwt.encode({"iat": 0, "exp": 300, "scope":
"dag_processor", "job_id": 1}, "secret")
Review Comment:
test_api_client.py already has `make_job_token` and
`make_registration_response`, and test_processor.py builds the same job token
and `/jobs` handler again. Moving those to a `conftest.py` in this directory
would let all three share them.
##########
airflow-core/src/airflow/dag_processing/manager.py:
##########
@@ -370,15 +463,18 @@ def run(self):
By processing them in separate processes, we can get parallelism and
isolation
from potentially harmful user code.
"""
- self.before_run()
try:
+ self.before_run()
return self._run_parsing_loop()
finally:
self.after_run()
def before_run(self) -> None:
"""Set up state required before the parsing loop starts. Default
implementation; override to customize."""
- self.prepare_server_process_context()
+ if self.api_client is None:
+ self.prepare_server_process_context()
+ else:
+ self.prepare_api_secrets_context(self.api_client)
Review Comment:
The base layer documented override points for an API-backed manager:
`claim_priority_files` and `fetch_callbacks` say to override them to source
work from an API, and `prepare_server_process_context` says to make it a no-op
in API-backed subclasses. This layer puts the API path inside those defaults
behind `if self.api_client` instead, 21 branches here and 28 at the top of the
stack. Was a subclass that implements those hooks considered? It would keep the
DB path free of API checks. If the flag stays, those docstrings need updating,
since the defaults no longer read only from the metadata DB.
##########
airflow-core/src/airflow/dag_processing/manager.py:
##########
@@ -418,6 +522,28 @@ def prepare_server_process_context(self) -> None:
# Related: https://github.com/apache/airflow/pull/57459
os.environ["_AIRFLOW_PROCESS_CONTEXT"] = "server"
+ def prepare_api_secrets_context(self, api_client: DagProcessorAPIClient)
-> None:
+ """
+ Resolve this process's own connection and variable lookups through
``api_client``.
+
+ Bundle code runs in this process, for example a Git bundle reading its
connection. Its secret
+ lookups use the API. Parse processes replace
Review Comment:
Lang-SDK parse processes don't: `_start_runtime_entrypoint` in
lang_sdk_processor.py never resets `SUPERVISOR_COMMS`, and on Linux it's a bare
fork. So the `get_dagbag_import_timeout` policy and the coordinator lookup run
with this process's `DagProcessorSecretsComms` and a forked copy of its httpx
client. Should that entrypoint drop it too, or should this say only Python
parse processes replace it?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
