Name the JAR manifest attributes the Java coordinator reads

Author: jason810496

task-sdk/src/airflow/sdk/coordinators/java/_dag_importer.py

@@ -23,22 +23,19 @@
 from typing import TYPE_CHECKING, ClassVar, Final
 
 from airflow.sdk.coordinators._dag_importer import CoordinatorDagImporter
-from airflow.sdk.coordinators.java._jar_manifest import read_main_attributes
+from airflow.sdk.coordinators.java._jar_manifest import DAG_CODE, MAIN_CLASS, read_main_attributes
 from airflow.sdk.importers.base import DagSourceCode
 
 if TYPE_CHECKING:
     from airflow.sdk.coordinators.java.coordinator import JavaCoordinator
     from airflow.sdk.importers.base import DagDefinition
 
-_DAG_CODE_ATTRIBUTE: Final = "airflow-java-sdk-dag-code"
 _MAX_SOURCE_BYTES: Final = 1024 * 1024
 _NO_SOURCE: Final = (
     "// This JAR embeds no Dag source. Build it with the Airflow Java SDK Gradle plugin, or set\n"
     "// airflowBundle.dagSource, to show the source here.\n"
 )
-_SOURCE_TOO_LARGE: Final = (
-    f"// The Dag source this JAR embeds is over {_MAX_SOURCE_BYTES} bytes, so it is not shown.\n"
-)
+_SOURCE_TOO_LARGE: Final = "// The Dag source this JAR embeds is over 1 MiB, so it is not shown.\n"
 
 
 class JavaDagImporter(CoordinatorDagImporter):
@@ -64,7 +61,7 @@ def might_contain_dag(self, definition: DagDefinition, safe_mode: bool) -> bool:
                 attributes = read_main_attributes(zf) or {}
         except (OSError, zipfile.BadZipFile):
             return True
-        if not (main_class := attributes.get("main-class")):
+        if not (main_class := attributes.get(MAIN_CLASS)):
             return False
         return not self.coordinator.main_class or main_class == self.coordinator.main_class
 
@@ -80,7 +77,7 @@ def get_source_code(self, definition: DagDefinition) -> DagSourceCode:
 
 
 def _find_source_entry(zf: zipfile.ZipFile) -> zipfile.ZipInfo | None:
-    if not (entry := (read_main_attributes(zf) or {}).get(_DAG_CODE_ATTRIBUTE)):
+    if not (entry := (read_main_attributes(zf) or {}).get(DAG_CODE)):
         return None
     try:
         return zf.getinfo(entry)

task-sdk/src/airflow/sdk/coordinators/java/_jar_manifest.py

@@ -27,6 +27,11 @@
 
 MANIFEST_NAME: Final = "META-INF/MANIFEST.MF"
 
+# Attribute keys as parse_main_attributes returns them, lower-cased.
+MAIN_CLASS: Final = "main-class"
+SUPERVISOR_SCHEMA_VERSION: Final = "airflow-supervisor-schema-version"
+DAG_CODE: Final = "airflow-java-sdk-dag-code"
+
 _LINE_END = re.compile(rb"\r\n|\r|\n")
 
 

task-sdk/src/airflow/sdk/coordinators/java/coordinator.py

@@ -35,7 +35,11 @@
 )
 from airflow.sdk.coordinators._subprocess import SubprocessCoordinator
 from airflow.sdk.coordinators.java._dag_importer import JavaDagImporter
-from airflow.sdk.coordinators.java._jar_manifest import read_main_attributes
+from airflow.sdk.coordinators.java._jar_manifest import (
+    MAIN_CLASS,
+    SUPERVISOR_SCHEMA_VERSION,
+    read_main_attributes,
+)
 
 if TYPE_CHECKING:
     from collections.abc import Iterable, Iterator, Sequence
@@ -108,7 +112,7 @@ def from_jar(cls, path: pathlib.Path) -> Self | None:
         if attributes is None:
             log.debug("JAR does not contain META-INF/MANIFEST.MF; ignored", path=path)
             return None
-        return cls(attributes.get("main-class"), attributes.get("airflow-supervisor-schema-version"))
+        return cls(attributes.get(MAIN_CLASS), attributes.get(SUPERVISOR_SCHEMA_VERSION))
 
 
 @attrs.define
@@ -230,7 +234,8 @@ def _build_execute_task_command(self, *, what: TaskInstance) -> tuple[list[str],
         return self._build_command(roots, jar.main_class), jar.schema_version
 
     def _build_parse_dag_command(self, *, path: pathlib.Path) -> tuple[list[str], str | None]:
-        # The classpath and main class match execution's, so a parse runs the code a task runs.
+        # Same command shape as execution. With one executable JAR per bundle, or main_class set,
+        # a parse runs the class a task runs.
         meta = _JarMetadata.from_jar(path)
         if meta is None:
             raise ValueError(f"Cannot read the manifest of {path}")