Pin lang-SDK k8s test's Go/Java SDKs to upstream main

dev/breeze/src/airflow_breeze/commands/kubernetes_commands.py

@@ -100,6 +100,7 @@
     check_if_image_exists,
     run_command,
 )
+from airflow_breeze.utils.shared_options import get_dry_run
 
 KUBERNETES_PYTEST_ARGS = [
     "--strict-markers",
@@ -2497,6 +2498,7 @@ def deploy_cluster(
 LANG_SDK_GO_BUILDER_IMAGE = os.environ.get("GO_BUILDER_IMAGE", "golang:1.25-alpine")
 LANG_SDK_JAVA_BUILDER_IMAGE = "eclipse-temurin:17-jdk"
 LANG_SDK_MAVEN_CACHE_PATH = AIRFLOW_ROOT_PATH / "files" / "m2"
+LANG_SDK_GRADLE_CACHE_PATH = AIRFLOW_ROOT_PATH / "files" / "gradle"
 # The Java queue needs a JRE the JavaCoordinator can exec; the Go queue runs on
 # the plain prod image. Building the Java worker image as a separate tag (prod +
 # JRE, see Dockerfile.java) lets each coordinator route its queue to a distinct
@@ -2507,37 +2509,124 @@ def deploy_cluster(
     "aws://test:test@/?region_name=us-east-1&"
     "endpoint_url=http%3A%2F%2Flocalstack.airflow.svc.cluster.local%3A4566"
 )
+# The Go/Java SDKs are always built from upstream main so branches with stale or missing
+# go-sdk/java-sdk copies still test current SDK sources. See kubernetes-tests/lang_sdk/README.md.
+LANG_SDK_UPSTREAM_GIT_URL = "https://github.com/apache/airflow.git"
+LANG_SDK_UPSTREAM_REF = "main"
 
 
-def _lang_sdk_build_go_bundle(staging: Path, output: Output | None, *, native: bool = False) -> None:
+def _lang_sdk_fetch_upstream_sdk_sources(staging: Path, output: Output | None) -> tuple[Path, Path]:
+    """Extract go-sdk/ and java-sdk/ from upstream main into a throwaway staging dir.
+
+    Prefers the ``upstream`` remote when configured, falling back to the canonical GitHub URL
+    (CI has no ``upstream`` and ``origin`` may be a fork). Shallow fetch + ``git archive`` never
+    touch the working tree or index.
+
+    The real, local task-sdk is symlinked alongside the extraction because java-sdk's
+    ``sdk/build.gradle.kts`` reads a sibling ``../task-sdk/.../schema.json``. The gradle wrapper
+    scripts and jar are ``export-ignore`` (ASF LEGAL-570) so ``git archive`` drops them;
+    ``git show`` restores them, re-marking ``gradlew`` executable.
+    """
+    remotes = run_command(
+        ["git", "remote"], cwd=AIRFLOW_ROOT_PATH, output=output, capture_output=True, text=True, check=True
+    ).stdout.split()
+    fetch_source = "upstream" if "upstream" in remotes else LANG_SDK_UPSTREAM_GIT_URL
+    get_console(output=output).print(
+        f"[info]Fetching {LANG_SDK_UPSTREAM_REF} from {fetch_source} for the lang-SDK Go/Java sources"
+    )
+    run_command(
+        ["git", "fetch", "--depth=1", fetch_source, LANG_SDK_UPSTREAM_REF],
+        cwd=AIRFLOW_ROOT_PATH,
+        output=output,
+        check=True,
+    )
+    sha = run_command(
+        ["git", "rev-parse", "FETCH_HEAD"],
+        cwd=AIRFLOW_ROOT_PATH,
+        output=output,
+        capture_output=True,
+        text=True,
+        check=True,
+    ).stdout.strip()
+    get_console(output=output).print(f"[info]lang-SDK Go/Java sources pinned to upstream main @ {sha}")
+    extracted = staging / "upstream_lang_sdk_sources"
+    extracted.mkdir(parents=True, exist_ok=True)
+    archive_path = staging / "upstream_lang_sdk_sources.tar"
+    run_command(
+        ["git", "archive", "--format=tar", f"--output={archive_path}", sha, "--", "go-sdk", "java-sdk"],
+        cwd=AIRFLOW_ROOT_PATH,
+        output=output,
+        check=True,
+    )
+    run_command(["tar", "-xf", str(archive_path), "-C", str(extracted)], output=output, check=True)
+    for rel_path, mode in (
+        ("java-sdk/gradlew", 0o755),
+        ("java-sdk/gradlew.bat", 0o644),
+        ("java-sdk/gradle/wrapper/gradle-wrapper.jar", 0o644),
+    ):
+        restored = extracted / rel_path
+        restored.parent.mkdir(parents=True, exist_ok=True)
+        restored.write_bytes(
+            run_command(
+                ["git", "show", f"{sha}:{rel_path}"],
+                cwd=AIRFLOW_ROOT_PATH,
+                output=output,
+                capture_output=True,
+                check=True,
+            ).stdout
+            or b""
+        )
+        restored.chmod(mode)
+    (extracted / "task-sdk").symlink_to(AIRFLOW_ROOT_PATH / "task-sdk")
+    return extracted / "go-sdk", extracted / "java-sdk"
+
+
+def _lang_sdk_build_go_bundle(
+    staging: Path, upstream_go_sdk: Path, output: Output | None, *, native: bool = False
+) -> None:
     """Build the Go bundle into ``staging/go-artifacts`` and copy the result into the staging dir.
 
     By default the build runs in an ephemeral Go toolchain container so the host needs no Go install,
     writing its caches into a gitignored dir under the repo. In ``native`` mode (used in CI, where the
     host already has a cached Go toolchain via ``actions/setup-go``) it invokes the host ``go`` directly,
     skipping the container image pull and reusing the runner's module/build cache.
+
+    go_example's go.mod ``replace``s go-sdk by relative path, so the build runs in a scratch
+    workspace mirroring the repo layout with ``upstream_go_sdk`` at ``<workspace>/go-sdk``,
+    letting the unmodified directive resolve against the upstream copy.
     """
     go_dir = staging / "go-artifacts"
     go_dir.mkdir(parents=True, exist_ok=True)
-    output_bin = LANG_SDK_GO_EXAMPLE_PATH / "bin" / LANG_SDK_GO_BUNDLE_NAME
+    example_rel = LANG_SDK_GO_EXAMPLE_PATH.relative_to(AIRFLOW_ROOT_PATH)
+    workspace = staging / "go_workspace"
+    example_path = workspace / example_rel
+    example_path.parent.mkdir(parents=True, exist_ok=True)
+    shutil.copytree(LANG_SDK_GO_EXAMPLE_PATH, example_path, ignore=shutil.ignore_patterns(".home"))
+    # In dry-run the fetch/extract commands are skipped, so the upstream copy and build outputs
+    # never materialize -- skip the filesystem work that depends on them.
+    if not get_dry_run():
+        shutil.copytree(upstream_go_sdk, workspace / "go-sdk")
+    output_bin = example_path / "bin" / LANG_SDK_GO_BUNDLE_NAME
+    output_bin.parent.mkdir(parents=True, exist_ok=True)
 
     # CGO_ENABLED=0 yields a fully static binary that runs on the stock worker. The package built is
     # the current dir (".") because go_example is its own module.
     if native:
         get_console(output=output).print("[info]Building Go bundle with the host Go toolchain")
         run_command(
             ["go", "tool", "airflow-go-pack", "--output", str(output_bin), "."],
-            cwd=LANG_SDK_GO_EXAMPLE_PATH,
+            cwd=example_path,
             env={**os.environ, "CGO_ENABLED": "0"},
             output=output,
             check=True,
         )
     else:
         uid_gid = f"{os.getuid()}:{os.getgid()}"
-        go_example_ctr = f"/repo/{LANG_SDK_GO_EXAMPLE_PATH.relative_to(AIRFLOW_ROOT_PATH).as_posix()}"
+        go_example_ctr = f"/repo/{example_rel.as_posix()}"
         # USER/HOME must be set because the SDK calls user.Current() at init; with cgo disabled Go's
-        # pure-Go resolver reads those env vars and panics if either is empty. HOME points at a
-        # writable, gitignored dir under go_example so the caches persist between runs on a dev host.
+        # pure-Go resolver reads those env vars and panics if either is empty. HOME is mounted from
+        # the real go_example's gitignored cache dir so the caches persist across scratch workspaces.
+        (LANG_SDK_GO_EXAMPLE_PATH / ".home").mkdir(parents=True, exist_ok=True)
         get_console(output=output).print(f"[info]Building Go bundle in {LANG_SDK_GO_BUILDER_IMAGE}")
         run_command(
             [
@@ -2553,7 +2642,9 @@ def _lang_sdk_build_go_bundle(staging: Path, output: Output | None, *, native: b
                 "-e",
                 "CGO_ENABLED=0",
                 "-v",
-                f"{AIRFLOW_ROOT_PATH}:/repo",
+                f"{workspace}:/repo",
+                "-v",
+                f"{LANG_SDK_GO_EXAMPLE_PATH / '.home'}:{go_example_ctr}/.home",
                 "-w",
                 go_example_ctr,
                 LANG_SDK_GO_BUILDER_IMAGE,
@@ -2567,10 +2658,13 @@ def _lang_sdk_build_go_bundle(staging: Path, output: Output | None, *, native: b
             output=output,
             check=True,
         )
-    shutil.copy(output_bin, go_dir / LANG_SDK_GO_BUNDLE_NAME)
+    if not get_dry_run():
+        shutil.copy(output_bin, go_dir / LANG_SDK_GO_BUNDLE_NAME)
 
 
-def _lang_sdk_build_java_jar(staging: Path, output: Output | None, *, native: bool = False) -> None:
+def _lang_sdk_build_java_jar(
+    staging: Path, upstream_java_sdk: Path, output: Output | None, *, native: bool = False
+) -> None:
     """Publish the Java SDK to mavenLocal then build the java_example jar into ``staging/java-artifacts``.
 
     By default the build runs in an ephemeral JDK container so the host needs no JDK, persisting the
@@ -2579,49 +2673,58 @@ def _lang_sdk_build_java_jar(staging: Path, output: Output | None, *, native: bo
     host ``./gradlew`` directly, skipping the container image pull and reusing the runner's ``~/.gradle``
     cache. ``java_example`` resolves the SDK from ``mavenLocal()``, so the SDK is published first, then
     the bundle is built with java-sdk's gradle wrapper pointed at the example project (``-p``).
+
+    Both gradle invocations run against ``upstream_java_sdk`` rather than the local ``java-sdk/``;
+    only ``-p`` stays pointed at the local ``java_example``, which is test-harness code that keeps
+    tracking the checked-out branch.
     """
     java_dir = staging / "java-artifacts"
     java_dir.mkdir(parents=True, exist_ok=True)
-    java_sdk_path = AIRFLOW_ROOT_PATH / "java-sdk"
 
     if native:
         get_console(output=output).print(
             "[info]Publishing Java SDK artifacts to local Maven repository with the host Gradle toolchain"
         )
         run_command(
             ["./gradlew", "publishToMavenLocal", "-PskipSigning=true", "--no-daemon", "--console=plain"],
-            cwd=java_sdk_path,
+            cwd=upstream_java_sdk,
             output=output,
             check=True,
         )
         get_console(output=output).print("[info]Building Java jar with the host Gradle toolchain")
         run_command(
             ["./gradlew", "-p", str(LANG_SDK_JAVA_EXAMPLE_PATH), "bundle", "--no-daemon", "--console=plain"],
-            cwd=java_sdk_path,
+            cwd=upstream_java_sdk,
             output=output,
             check=True,
         )
     else:
         uid_gid = f"{os.getuid()}:{os.getgid()}"
         java_example_ctr = f"/repo/{LANG_SDK_JAVA_EXAMPLE_PATH.relative_to(AIRFLOW_ROOT_PATH).as_posix()}"
         # --user keeps build outputs owned by the host user; HOME is set explicitly because that UID has
-        # no /etc/passwd entry; GRADLE_USER_HOME and the mounted ~/.m2 persist the Gradle distribution and
-        # dependency caches between runs on a dev host.
+        # no /etc/passwd entry. GRADLE_USER_HOME and the mounted ~/.m2 persist the Gradle distribution
+        # and dependency caches between runs -- outside /repo/java-sdk, which is remounted to the
+        # (per-run) upstream copy below.
         LANG_SDK_MAVEN_CACHE_PATH.mkdir(parents=True, exist_ok=True)
+        LANG_SDK_GRADLE_CACHE_PATH.mkdir(parents=True, exist_ok=True)
         java_docker_prefix = [
             "docker",
             "run",
             "--rm",
             "--user",
             uid_gid,
             "-e",
-            "GRADLE_USER_HOME=/repo/java-sdk/.gradle",
+            "GRADLE_USER_HOME=/workspace-home/.gradle",
             "-e",
             "HOME=/workspace-home",
             "-v",
             f"{LANG_SDK_MAVEN_CACHE_PATH}:/workspace-home/.m2",
             "-v",
+            f"{LANG_SDK_GRADLE_CACHE_PATH}:/workspace-home/.gradle",
+            "-v",
             f"{AIRFLOW_ROOT_PATH}:/repo",
+            "-v",
+            f"{upstream_java_sdk}:/repo/java-sdk",
         ]
         get_console(output=output).print("[info]Publishing Java SDK artifacts to local Maven repository")
         run_command(
@@ -2656,6 +2759,8 @@ def _lang_sdk_build_java_jar(staging: Path, output: Output | None, *, native: bo
             output=output,
             check=True,
         )
+    if get_dry_run():
+        return
     jars = list((LANG_SDK_JAVA_EXAMPLE_PATH / "build" / "bundle").glob("*.jar"))
     if not jars:
         get_console(output=output).print("[error]No jar produced by the Java bundle build")
@@ -2716,7 +2821,11 @@ def _lang_sdk_upload_artifacts(
     ).stdout.strip()
 
     go_bundle = staging / "go-artifacts" / "lang_sdk_combined"
-    java_jar = next((staging / "java-artifacts").glob("*.jar"))
+    if get_dry_run():
+        # The dry-run build steps produce no jar; use the placeholder name so the commands still print.
+        java_jar = staging / "java-artifacts" / "app.jar"
+    else:
+        java_jar = next((staging / "java-artifacts").glob("*.jar"))
     stub_dag = LANG_SDK_PATH / "dags" / "lang_sdk_combined.py"
 
     for src, dest in (
@@ -2921,9 +3030,9 @@ def _setup_lang_sdk_test(
 ) -> None:
     """Provision the lang-SDK coordinator env on an already-deployed KubernetesExecutor cluster.
 
-    Builds the Go/Java artifacts, the Java worker image and deploys localstack in parallel, then
-    serially uploads the artifacts, applies the config + secret, and helm-upgrades Airflow with the
-    lang-SDK values.
+    Fetches go-sdk/java-sdk from upstream main, then builds the Go/Java artifacts, the Java worker
+    image and deploys localstack in parallel, then serially uploads the artifacts, applies the
+    config + secret, and helm-upgrades Airflow with the lang-SDK values.
     """
     go_image = go_image or f"{BuildProdParams(python=python).airflow_image_kubernetes}:latest"
     build_java_image = java_image is None
@@ -2936,9 +3045,16 @@ def _setup_lang_sdk_test(
     native = os.environ.get("LANG_SDK_NATIVE_TOOLCHAIN", "").lower() == "true"
     with tempfile.TemporaryDirectory(prefix="lang_sdk_artifacts_") as tmp:
         staging = Path(tmp)
+        upstream_go_sdk, upstream_java_sdk = _lang_sdk_fetch_upstream_sdk_sources(staging, output)
         steps: list[tuple[str, Callable[[Output | None], Any]]] = [
-            ("Build Go bundle", lambda o: _lang_sdk_build_go_bundle(staging, o, native=native)),
-            ("Build Java jar", lambda o: _lang_sdk_build_java_jar(staging, o, native=native)),
+            (
+                "Build Go bundle",
+                lambda o: _lang_sdk_build_go_bundle(staging, upstream_go_sdk, o, native=native),
+            ),
+            (
+                "Build Java jar",
+                lambda o: _lang_sdk_build_java_jar(staging, upstream_java_sdk, o, native=native),
+            ),
             ("Deploy localstack", lambda o: _lang_sdk_deploy_localstack(python, kubernetes_version, o)),
         ]
         if build_java_image: