Java SDK: Check native Dag serialization against Airflow's

Author: jason810496

.pre-commit-config.yaml

@@ -403,6 +403,23 @@ repos:
           ^airflow-core/src/airflow/serialization/serialized_objects\.py$|
           ^airflow-core/src/airflow/serialization/schema\.json$|
           ^scripts/ci/lang_sdk_serialization/.*$
+      - id: check-java-sdk-serialization-conformance
+        name: Check the Java SDK serializes Dags the way Airflow does
+        description: "Serialize the shared test Dags with the Java SDK and with Airflow, and compare the two"
+        entry: ./java-sdk/scripts/ci/prek/check_serialization_conformance.py
+        language: system
+        pass_filenames: false
+        require_serial: true
+        files: >
+          (?x)
+          ^airflow-core/src/airflow/serialization/serialized_objects\.py$|
+          ^airflow-core/src/airflow/serialization/schema\.json$|
+          ^java-sdk/scripts/ci/prek/check_serialization_conformance\.py$|
+          ^java-sdk/sdk/schema/dag-schema\.json$|
+          ^java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/(DagDef|Deps|Arg|TaskGroupRef)\.kt$|
+          ^java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Serde\.kt$|
+          ^java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/conformance/.*$|
+          ^scripts/ci/lang_sdk_serialization/.*$
       - id: check-go-version-in-sync
         name: Check Go toolchain version is consistent across build files
         entry: ./scripts/ci/prek/check_go_version_in_sync.py

java-sdk/scripts/ci/prek/check_serialization_conformance.py

@@ -0,0 +1,51 @@
+#!/usr/bin/env python3
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Check the Java SDK serializes Dags as Airflow does, with scripts/ci/lang_sdk_serialization/compare.py."""
+
+from __future__ import annotations
+
+import subprocess
+import sys
+from pathlib import Path
+
+sys.path.insert(0, str(Path(__file__).resolve().parents[4] / "scripts" / "ci" / "prek"))
+
+from common_prek_utils import AIRFLOW_ROOT_PATH
+
+if __name__ not in ("__main__", "__mp_main__"):
+    raise SystemExit(
+        "This file is intended to be executed as an executable program. You cannot use it as a module."
+        f"To run this script, run the ./{__file__} command"
+    )
+
+if __name__ == "__main__":
+    java_sdk = AIRFLOW_ROOT_PATH / "java-sdk"
+    classpath = (
+        subprocess.run(
+            [str(java_sdk / "gradlew"), "-p", str(java_sdk), "-q", ":sdk:printConformanceClasspath"],
+            check=True,
+            capture_output=True,
+            text=True,
+        )
+        .stdout.strip()
+        .splitlines()[-1]
+    )
+    compare = AIRFLOW_ROOT_PATH / "scripts" / "ci" / "lang_sdk_serialization" / "compare.py"
+    serializer = ["java", "-cp", classpath, "org.apache.airflow.sdk.conformance.SerializeJavaKt"]
+    command = [sys.executable, str(compare), "--sdk", "java", "--", *serializer]
+    sys.exit(subprocess.run(command, check=False).returncode)

java-sdk/sdk/build.gradle.kts

@@ -815,3 +815,12 @@ publishing {
         }
     }
 }
+
+// Prints the classpath that runs the conformance serializer, for
+// java-sdk/scripts/ci/prek/check_serialization_conformance.py.
+tasks.register("printConformanceClasspath") {
+    dependsOn("testClasses")
+    // Capture early to keep compatibility to the Gradle configuration cache.
+    val classpath = sourceSets.test.get().runtimeClasspath
+    doLast { println(classpath.asPath) }
+}

java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/conformance/SerializeJava.kt

@@ -0,0 +1,119 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+// Serializes the shared test Dags with this SDK, for
+// scripts/ci/lang_sdk_serialization/compare.py:
+//
+//   java -cp <sdk test runtime classpath> org.apache.airflow.sdk.conformance.SerializeJavaKt \
+//       scripts/ci/lang_sdk_serialization/test_dags.yaml serialized_java.json
+//
+// Each Dag is built with the interface API and written as the runtime answers
+// a parse request, keyed by Dag ID. A task's `upstream` becomes an ordering
+// edge, which a Java task serializes the same way as a data edge.
+package org.apache.airflow.sdk.conformance
+
+import com.fasterxml.jackson.databind.JsonNode
+import com.fasterxml.jackson.databind.ObjectMapper
+import com.fasterxml.jackson.dataformat.yaml.YAMLMapper
+import org.apache.airflow.sdk.Bundle
+import org.apache.airflow.sdk.Client
+import org.apache.airflow.sdk.Context
+import org.apache.airflow.sdk.DagDef
+import org.apache.airflow.sdk.Deps
+import org.apache.airflow.sdk.Task
+import org.apache.airflow.sdk.TaskGroupRef
+import org.apache.airflow.sdk.TaskRef
+import org.apache.airflow.sdk.execution.serializeDag
+import org.apache.airflow.sdk.internal.Field
+import org.apache.airflow.sdk.internal.FieldType
+import org.apache.airflow.sdk.internal.SchemaFields
+import java.io.File
+import java.time.Duration
+import java.time.OffsetDateTime
+
+class ConformanceTask : Task {
+  override fun execute(
+    context: Context,
+    client: Client,
+  ) = Unit
+}
+
+fun main(args: Array<String>) {
+  require(args.size == 2) { "usage: SerializeJava <test_dags.yaml> <output.json>" }
+  val cases = YAMLMapper().readTree(File(args[0])).path("dags")
+  val bundle = Bundle()
+  cases.forEach { bundle.register(buildDag(it)) }
+  val serialized =
+    bundle.dags.values.associate { dag ->
+      dag.id to mapOf("__version" to 3, "dag" to serializeDag(dag, "", "."))
+    }
+  ObjectMapper().writerWithDefaultPrettyPrinter().writeValue(File(args[1]), serialized)
+}
+
+private fun buildDag(case: JsonNode): DagDef {
+  val dag = DagDef(case.path("dag_id").asText())
+  case.path("spec").fields().forEach { (key, value) -> dag.config(key, toValue(SchemaFields.DAG, key, value)) }
+
+  // A group ID is fully qualified, so its parent is whatever comes before the last dot.
+  val groups = linkedMapOf<String, TaskGroupRef>()
+  case.path("groups").forEach { node ->
+    val groupId = node.asText()
+    val parentId = groupId.substringBeforeLast('.', "")
+    val localId = groupId.substringAfterLast('.')
+    groups[groupId] = if (parentId.isEmpty()) dag.taskGroup(localId) else groups.getValue(parentId).taskGroup(localId)
+  }
+
+  val tasks = linkedMapOf<String, TaskRef<Unit>>()
+  case.path("tasks").forEach { task ->
+    val groupId = task.path("group").asText("")
+    val localId = task.path("task_id").asText()
+    val ref =
+      if (groupId.isEmpty()) {
+        dag.task<Unit>(localId, ConformanceTask::class.java)
+      } else {
+        groups.getValue(groupId).task(localId, ConformanceTask::class.java)
+      }
+    task.path("spec").fields().forEach { (key, value) -> ref.config(key, toValue(SchemaFields.TASK, key, value)) }
+    task.path("upstream").forEach { ref.after(tasks.getValue(it.asText())) }
+    tasks[ref.def.id] = ref
+  }
+
+  case.path("order_edges").forEach { edge ->
+    val node = { id: String -> groups[id] as Deps.Flow? ?: tasks.getValue(id) }
+    node(edge[0].asText()).before(node(edge[1].asText()))
+  }
+  return dag
+}
+
+/** Reads a YAML value as the Java type the config key takes. */
+private fun toValue(
+  table: Map<String, Field>,
+  key: String,
+  node: JsonNode,
+): Any =
+  when (requireNotNull(table[key]) { "Unknown config key: '$key'" }.type) {
+    FieldType.STRING -> node.asText()
+    FieldType.BOOLEAN -> node.asBoolean()
+    FieldType.INTEGER -> node.asInt()
+    FieldType.NUMBER -> node.numberValue()
+    // `!datetime` is an ISO 8601 timestamp, and `!timedelta` a number of seconds.
+    FieldType.DATETIME -> OffsetDateTime.parse(node.asText())
+    FieldType.TIMEDELTA -> Duration.ofNanos((node.asText().toDouble() * 1e9).toLong())
+    FieldType.STRING_ARRAY -> node.map { it.asText() }
+  }

scripts/ci/lang_sdk_serialization/compare.py

@@ -30,6 +30,9 @@
     python3 scripts/ci/lang_sdk_serialization/compare.py --sdk typescript -- \
         pnpm --dir ts-sdk exec tsx tests/conformance/serialize_typescript.ts
 
+Each SDK's own prek hook runs it this way, such as
+java-sdk/scripts/ci/prek/check_serialization_conformance.py for the Java SDK.
+
 On a failure the files are kept, and their directory is printed.
 """
 
@@ -83,21 +86,25 @@ def get_task_defaults() -> dict[str, Any]:
     return {key: field["default"] for key, field in fields.items() if field.get("default") is not None}
 
 
-def normalize_as_javascript(value: Any) -> Any:
-    """Read a JSON value as JavaScript does: one number type, and a bool that is not a number."""
+def normalize_numbers(value: Any) -> Any:
+    """
+    Read a JSON value with one number type, and a bool that is not a number.
+
+    An SDK may write ``2`` where Python writes ``2.0``, as JavaScript does, which has one number type.
+    """
     if isinstance(value, bool):
         return ("bool", value)
     if isinstance(value, (int, float)):
         return float(value)
     if isinstance(value, list):
-        return [normalize_as_javascript(item) for item in value]
+        return [normalize_numbers(item) for item in value]
     if isinstance(value, dict):
-        return {key: normalize_as_javascript(item) for key, item in value.items()}
+        return {key: normalize_numbers(item) for key, item in value.items()}
     return value
 
 
 def is_same_json(python: Any, sdk: Any) -> bool:
-    return normalize_as_javascript(python) == normalize_as_javascript(sdk)
+    return normalize_numbers(python) == normalize_numbers(sdk)
 
 
 def find_differences(path: str, python: Any, sdk: Any) -> list[str]:

scripts/ci/lang_sdk_serialization/test_dags.yaml

@@ -62,7 +62,7 @@ dags:
           wait_for_past_depends_before_skipping: true
           wait_for_downstream: true
           retries: 2
-          queue: typescript
+          queue: workers
           pool: tiny
           pool_slots: 2
           execution_timeout: !timedelta 120
@@ -126,7 +126,7 @@ dags:
 
   # Fan-out and fan-in, so downstream_task_ids has to be sorted rather than kept in wiring
   # order. The schedule is not written "@daily" because Python expands a preset before
-  # serializing it, which serde.test.ts covers separately.
+  # serializing it, which each SDK's own serializer tests cover separately.
   - dag_id: conformance_diamond
     spec:
       schedule: "0 0 * * *"