diff --git a/airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst b/airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst index f22f54c0cdf..2e8aba27521 100644 --- a/airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst +++ b/airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst @@ -719,6 +719,35 @@ The ``Client`` A value from a secrets backend, such as an ``AIRFLOW_VAR_*`` environment variable, still wins over the stored value when the Variable is read back. ``setVariable`` without a description clears the existing description. +The Dag run conf +~~~~~~~~~~~~~~~~ + +``context.dagRun.conf`` is the conf the Dag run was triggered with. It is a ``Map``, and its values have the Java +types listed in :ref:`the XCom type mapping `. ``context.dagRun.getConf(type)`` reads the conf into +an instance of a class with typed fields instead: + +.. code-block:: java + + public static class EtlConf { + public String targetTable; + public int limit = 100; + public LocalDate since; + } + + public class Load implements Task { + @Override + public void execute(Context context, Client client) { + EtlConf etl = context.dagRun.getConf(EtlConf.class); + } + } + +For a run triggered with the conf ``{"target_table": "sales", "since": "2024-01-31"}``, ``etl.targetTable`` is +``"sales"``, ``etl.since`` is ``LocalDate.of(2024, 1, 31)``, and ``etl.limit`` stays ``100``. Keys match field +names when case and underscores are ignored, and keys that match no field are ignored. ``getConf`` does not add +the defaults of the Dag's params, and the conf of a scheduled run is empty. For the classes that ``getConf`` +accepts and the cases where it throws, see ``DagRun.getConf`` in the `Java SDK API reference +`__. + .. _java-sdk/task-state-store: Task state store diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Context.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Context.kt index 69badbd49f6..d6a6426abf6 100644 --- a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Context.kt +++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Context.kt @@ -21,6 +21,7 @@ package org.apache.airflow.sdk import org.apache.airflow.sdk.execution.Logger import org.apache.airflow.sdk.execution.comm.StartupDetails +import org.apache.airflow.sdk.internal.bindConf import java.time.OffsetDateTime import org.apache.airflow.sdk.execution.comm.DagRun.DagRunType as CommDagRunType @@ -81,7 +82,8 @@ enum class DagRunType { * @property dataIntervalEnd End of the data interval. * @property runAfter A date-time tells the scheduler when the Dag run can be scheduled. * @property runType How the run was created. - * @property conf The configuration for this run. + * @property conf The configuration for this run. [getConf] reads it into an + * instance of a class with typed fields. */ data class DagRun( @JvmField val dagId: String, @@ -92,7 +94,67 @@ data class DagRun( @JvmField val runAfter: OffsetDateTime?, @JvmField val runType: DagRunType?, @JvmField val conf: Map, -) +) { + /** + * Reads [conf] into a new instance of [type]. For example, [type] can be a + * class with one field per conf key. + * + * ```java + * public static class EtlConf { + * public String targetTable; // set from the "target_table" key + * public int limit = 100; // stays 100 when the conf has no "limit" key + * } + * + * EtlConf etl = context.dagRun.getConf(EtlConf.class); + * ``` + * + * Each key goes to the field, setter or record component whose name matches + * the key when case and underscores are ignored. If one key matches a name + * exactly, that key wins over keys that match the same name only when case + * and underscores are ignored. Objects nested in the conf follow the same + * rules, including objects in lists and objects that are map values. The keys + * of a map field stay as written. + * + * Keys that match no field are ignored, because other tasks in the run may + * read them. A field with no matching key keeps its default value. The conf + * holds only the values that the run was triggered with. This method does not + * add the defaults of the Dag's params, so in a scheduled run, where the conf + * is empty, no field has a matching key. + * + * The class needs a no-argument constructor and public fields or setters, so + * a nested class must be `static`. A record also works. A record component + * with no matching key is null. If the component has a primitive type, this + * method throws instead, so a task that reads the conf into a record with a + * primitive component fails in a scheduled run. In Kotlin, declare the fields + * with `@JvmField` so that a property such as `isEnabled` keeps its name. + * + * Annotating a field with [ArgName] has no effect on this method. To set a + * field from a key that no Java name can match, such as `target-table`, add + * `@JsonProperty("target-table")` to that field. `@JsonProperty` needs + * `com.fasterxml.jackson.core:jackson-annotations` in your build, because the + * SDK keeps Jackson off your compile classpath. You can also read the key + * from [conf]. + * + * This method throws [IllegalArgumentException] in these cases: + * + * * A value, such as an empty string for an `int` field, cannot be converted + * to the type of its field. + * * A key holds null for a primitive field, a list holds null for a primitive + * array such as `int[]`, or a primitive record component has no matching + * key. To let a field or a component receive null, declare it with a boxed + * type, such as `Integer`. + * * Two keys, such as `target_table` and `TARGET_TABLE`, match the same field + * after case and underscores are ignored, but neither key matches the field + * name exactly. + * * The class has two fields whose names differ only in case or underscores, + * such as `targetTable` and `target_table`. This method checks the class + * even when the conf is empty, so it throws in a scheduled run too. + * + * @param type Class to read the conf into. + * @throws IllegalArgumentException in the cases listed above. + */ + fun getConf(type: Class): T = bindConf(conf, type) +} /** * Identifies the task instance that is currently executing. diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/internal/ConfBinding.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/internal/ConfBinding.kt new file mode 100644 index 00000000000..f5bf61cf2ed --- /dev/null +++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/internal/ConfBinding.kt @@ -0,0 +1,184 @@ +/* + * 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. + */ + +package org.apache.airflow.sdk.internal + +import com.fasterxml.jackson.databind.DeserializationFeature +import com.fasterxml.jackson.databind.JavaType +import com.fasterxml.jackson.databind.ObjectMapper +import com.fasterxml.jackson.databind.introspect.BeanPropertyDefinition +import com.fasterxml.jackson.databind.json.JsonMapper +import java.util.concurrent.ConcurrentHashMap + +// Other tasks in the Dag run may read keys that the target class does not +// declare, so this mapper must not fail on an unknown key. By default, Jackson +// moves a date-time that carries an offset to UTC. For example, +// 2024-01-31T02:00+08:00 would become 2024-01-30T18:00Z and fall on another +// date. This mapper keeps the offset as written instead. +// +// FAIL_ON_NULL_FOR_PRIMITIVES stops Jackson from turning a null into 0, for +// example in an int[]. It also rejects a record whose primitive component has +// no key, because Jackson handles the missing value as null. matchFields still +// checks a key that holds null, because its message names the key and the +// matching field. +private val confMapper: ObjectMapper by lazy { + JsonMapper + .builder() + .disable(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES) + .disable(DeserializationFeature.ADJUST_DATES_TO_CONTEXT_TIME_ZONE) + .enable(DeserializationFeature.FAIL_ON_NULL_FOR_PRIMITIVES) + .build() + .findAndRegisterModules() +} + +/** + * Converts a Dag run conf into an instance of [type] for + * [org.apache.airflow.sdk.DagRun.getConf]. That function documents how keys + * match fields. + * + * Jackson matches a key to a field only when the key is exactly the field + * name. So this function first uses [foldArgName] to find the field that each + * key matches, then renames the key to the exact name of that field. + */ +internal fun bindConf( + conf: Map, + type: Class, +): T { + val javaType = confMapper.constructType(type) + val matched = matchKeys(conf, javaType) + return try { + confMapper.convertValue(matched, javaType) + } catch (e: IllegalArgumentException) { + throw IllegalArgumentException("Cannot convert the Dag run conf to ${type.name}: ${e.message}", e) + } +} + +/** + * Renames the keys of every object in [value] that becomes an instance of a + * class when [value] is converted to [type]. The renaming also reaches objects + * in lists and arrays, objects that are map values, and the value inside a + * reference type such as `Optional`. The keys of a map are data rather than + * field names, so they stay as written. + */ +private fun matchKeys( + value: Any?, + type: JavaType, +): Any? = + when { + type.isReferenceType -> matchKeys(value, type.contentType) + value is Map<*, *> && type.isMapLikeType -> value.mapValues { (_, v) -> matchKeys(v, type.contentType) } + value is List<*> && (type.isCollectionLikeType || type.isArrayType) -> value.map { matchKeys(it, type.contentType) } + value is Map<*, *> -> matchFields(value, type) + else -> value + } + +/** + * Renames each key of [obj] to the name of the [type] field that the key + * matches. If a field has exactly the key's name, the key matches that field. + * Otherwise the key matches the field whose name folds to the same token as the + * key, unless another key has exactly that field's name. A key that matches no + * field keeps its name, and Jackson ignores that key. + * + * @throws IllegalArgumentException if the names of two fields fold to the same + * token, if two keys match the same field only after folding, or if a key + * holds null for a primitive field. + */ +private fun matchFields( + obj: Map<*, *>, + type: JavaType, +): Map { + val fields = findFields(type) + val className = type.rawClass.simpleName + val keys = obj.keys.mapTo(mutableSetOf()) { it.toString() } + val keyOfField = mutableMapOf() + val matched = LinkedHashMap() + obj.forEach { (rawKey, value) -> + val key = rawKey.toString() + val field = fields.byName[key] ?: fields.byFold[foldArgName(key)]?.takeUnless { it.name in keys } + if (field == null) { + matched[key] = value + return@forEach + } + val other = keyOfField.put(field.name, key) + require(other == null) { + "Dag run conf keys '$other' and '$key' both match field $className.${field.name} when case and " + + "underscores are ignored, so the SDK cannot tell which one to use. Name the field exactly as the key " + + "it should take." + } + require(value != null || !field.primaryType.isPrimitive) { + "Dag run conf key '$key' is null, but field $className.${field.name} has a primitive type that cannot " + + "be null. Declare a boxed type (e.g. Integer instead of int) to receive null." + } + matched[field.name] = matchKeys(value, field.primaryType) + } + return matched +} + +private class ConfFields( + val byName: Map, + val byFold: Map, +) + +/** + * Finds the fields of [type] so that each field can be looked up by its exact + * name or by its folded token from [foldArgName]. + * + * The result is cached per type. Jackson takes far longer to find the fields of + * a class than to convert one object of that class. Without the cache, + * converting a list of many objects would spend most of the time finding the + * same fields again. + * + * @throws IllegalArgumentException if two fields fold to the same token. + */ +private fun findFields(type: JavaType): ConfFields = confFieldsByType.computeIfAbsent(type) { indexFields(it) } + +private val confFieldsByType = ConcurrentHashMap() + +private fun indexFields(type: JavaType): ConfFields { + val fields = + confMapper.deserializationConfig + .introspect(type) + .findProperties() + .filter { it.isWritable() } + val className = type.rawClass.simpleName + val byFold = mutableMapOf() + fields.forEach { field -> + val previous = byFold.put(foldArgName(field.name), field) + require(previous == null) { + "Fields $className.${previous?.name} and $className.${field.name} differ only in case or underscores, " + + "so the SDK cannot tell which one a Dag run conf key should go to. Rename one of them." + } + } + return ConfFields(fields.associateBy { it.name }, byFold) +} + +/** + * Whether Jackson can set this property from a key. Jackson also lists + * read-only properties. For example, a `getBaseURL()` getter next to a + * `baseUrl` field gives a read-only property named `baseURL`. A read-only + * property must not take a key, and it must not clash with a field, as + * `baseURL` would clash with `baseUrl`. A collection or a map that has only a + * getter is not read-only, because Jackson adds the elements to the object that + * the getter returns. + */ +private fun BeanPropertyDefinition.isWritable(): Boolean = + hasSetter() || hasField() || hasConstructorParameter() || getter?.rawType?.let(::isCollectionOrMap) == true + +private fun isCollectionOrMap(type: Class<*>): Boolean = + Collection::class.java.isAssignableFrom(type) || Map::class.java.isAssignableFrom(type) diff --git a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/ContextTest.kt b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/ContextTest.kt index 2c48be12443..4a76e4f05f2 100644 --- a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/ContextTest.kt +++ b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/ContextTest.kt @@ -19,13 +19,137 @@ package org.apache.airflow.sdk +import com.fasterxml.jackson.annotation.JsonCreator +import com.fasterxml.jackson.annotation.JsonProperty import org.apache.airflow.sdk.execution.CoordinatorComm import org.apache.airflow.sdk.execution.byteArrayFromHexString import org.apache.airflow.sdk.execution.comm.StartupDetails import org.apache.airflow.sdk.execution.comm.TIRunContext import org.junit.jupiter.api.Assertions import org.junit.jupiter.api.Test +import java.time.LocalDate import java.time.OffsetDateTime +import java.util.concurrent.atomic.AtomicReference + +private class EtlConf { + @JvmField + var targetTable: String? = null + + @JvmField + var limit: Int = 100 + + @JvmField + var since: LocalDate? = null + + @JvmField + var runAt: OffsetDateTime? = null + + @JvmField + var source: EtlSource? = null + + @JvmField + var current: AtomicReference? = null + + @JvmField + var tables: List? = null + + @JvmField + var tableArray: Array? = null + + @JvmField + var tablesByName: Map? = null + + @JvmField + var options: Map? = null +} + +private class EtlSource { + @JvmField + var bucketName: String? = null +} + +private class EtlTable { + @JvmField + var tableName: String? = null +} + +private class FoldClash { + @JvmField + var targetTable: String? = null + + @JvmField + var targettable: String? = null +} + +private class UrlConf { + @JvmField + var baseUrl: String? = null + + fun getBaseURL(): String? = baseUrl +} + +private class SetterConf { + private var stored: String? = null + + fun setTargetTable(value: String?) { + stored = value + } + + fun stored(): String? = stored +} + +// Jackson binds a Java record through its constructor, but these tests compile +// for Java 11, which has no records. Jackson binds a class with a @JsonCreator +// constructor the same way, so this class stands in for a record. +private class CreatorConf + @JsonCreator + constructor( + @JsonProperty("targetTable") targetTable: String?, + @JsonProperty("rowLimit") rowLimit: Int, + ) { + val stored = targetTable + val storedLimit = rowLimit + } + +private class CountsConf { + @JvmField + var counts: IntArray? = null + + @JvmField + var matrix: Array? = null + + @JvmField + var boxedCounts: List? = null +} + +private class PinnedConf { + @JvmField + @JsonProperty("target-table") + var table: String? = null +} + +private class GetterOnlyTables { + private val backingList = mutableListOf() + private val backingMap = mutableMapOf() + + fun getTables(): MutableList = backingList + + fun getTablesByName(): MutableMap = backingMap +} + +private fun dagRunWithConf(conf: Map) = DagRun("d", "r", null, null, null, null, null, conf) + +private fun assertConfFailsOnField( + conf: Map, + type: Class<*>, + field: String, +) { + val error = Assertions.assertThrows(IllegalArgumentException::class.java) { dagRunWithConf(conf).getConf(type) } + + val message = error.message.orEmpty() + Assertions.assertTrue(message.startsWith("Cannot convert the Dag run conf to ${type.name}: ")) { message } + Assertions.assertTrue(message.contains("[\"$field\"]")) { message } +} class ContextTest { @Test @@ -125,4 +249,191 @@ class ContextTest { Assertions.assertNull(dr.dataIntervalStart) Assertions.assertTrue(dr.conf.isEmpty()) } + + @Test + fun getConfMatchesKeysIgnoringCaseAndUnderscores() { + val conf = dagRunWithConf(mapOf("target_table" to "sales", "LIMIT" to 5L)).getConf(EtlConf::class.java) + + Assertions.assertEquals("sales", conf.targetTable) + Assertions.assertEquals(5, conf.limit) + } + + @Test + fun getConfIgnoresKeysWithoutAFieldAndKeepsDefaults() { + val conf = dagRunWithConf(mapOf("notify_email" to "ops@example.com")).getConf(EtlConf::class.java) + + Assertions.assertNull(conf.targetTable) + Assertions.assertEquals(100, conf.limit) + } + + @Test + fun getConfPrefersAnExactKeyOverKeysThatMatchOnlyIgnoringCase() { + val conf = + dagRunWithConf(mapOf("target_table" to "folded", "targetTable" to "exact", "TARGET_TABLE" to "folded")) + .getConf(EtlConf::class.java) + + Assertions.assertEquals("exact", conf.targetTable) + } + + @Test + fun getConfMatchesKeysOfNestedObjects() { + val conf = + dagRunWithConf( + mapOf( + "source" to mapOf("bucket_name" to "raw"), + "current" to mapOf("table_name" to "staging"), + "tables" to listOf(mapOf("table_name" to "orders")), + "table_array" to listOf(mapOf("table_name" to "items")), + "tables_by_name" to mapOf("Sales_Key" to mapOf("table_name" to "sales")), + ), + ).getConf(EtlConf::class.java) + + Assertions.assertEquals("raw", conf.source?.bucketName) + Assertions.assertEquals("staging", conf.current?.get()?.tableName) + Assertions.assertEquals("orders", conf.tables?.single()?.tableName) + Assertions.assertEquals("items", conf.tableArray?.single()?.tableName) + Assertions.assertEquals(setOf("Sales_Key"), conf.tablesByName?.keys) + Assertions.assertEquals("sales", conf.tablesByName?.get("Sales_Key")?.tableName) + } + + @Test + fun getConfKeepsTheKeysOfAMapWithUntypedValues() { + val options = mapOf("Rows" to listOf(mapOf("Key_A" to "a")), "Inner_Map" to mapOf("Key_B" to "b")) + + val conf = dagRunWithConf(mapOf("options" to options)).getConf(EtlConf::class.java) + + Assertions.assertEquals(options, conf.options) + } + + @Test + fun getConfReadsIsoDateStringsIntoJavaTimeFields() { + val conf = + dagRunWithConf(mapOf("since" to "2024-01-31", "run_at" to "2024-01-31T02:00:00+08:00")) + .getConf(EtlConf::class.java) + + Assertions.assertEquals(LocalDate.of(2024, 1, 31), conf.since) + Assertions.assertEquals(OffsetDateTime.parse("2024-01-31T02:00:00+08:00"), conf.runAt) + } + + @Test + fun getConfMatchesKeysToSettersAndConstructorParameters() { + val conf = dagRunWithConf(mapOf("target_table" to "sales", "row_limit" to 5L)) + + Assertions.assertEquals("sales", conf.getConf(SetterConf::class.java).stored()) + Assertions.assertEquals("sales", conf.getConf(CreatorConf::class.java).stored) + Assertions.assertEquals(5, conf.getConf(CreatorConf::class.java).storedLimit) + } + + @Test + fun getConfRejectsAPrimitiveConstructorParameterWithNoKey() { + Assertions.assertNull(dagRunWithConf(mapOf("row_limit" to 5L)).getConf(CreatorConf::class.java).stored) + + assertConfFailsOnField(emptyMap(), CreatorConf::class.java, "rowLimit") + } + + @Test + fun getConfRejectsAnEmptyStringForAPrimitiveField() { + listOf("", " ").forEach { assertConfFailsOnField(mapOf("limit" to it), EtlConf::class.java, "limit") } + } + + @Test + fun getConfMatchesKeysToNamesFromJsonProperty() { + listOf("target-table", "TARGET-TABLE").forEach { key -> + Assertions.assertEquals("sales", dagRunWithConf(mapOf(key to "sales")).getConf(PinnedConf::class.java).table) + } + Assertions.assertNull(dagRunWithConf(mapOf("table" to "sales")).getConf(PinnedConf::class.java).table) + } + + @Test + fun getConfRejectsNullElementsOnlyForPrimitiveArrays() { + Assertions.assertEquals( + listOf(1, null), + dagRunWithConf(mapOf("boxed_counts" to listOf(1L, null))).getConf(CountsConf::class.java).boxedCounts, + ) + + assertConfFailsOnField(mapOf("counts" to listOf(1L, null)), CountsConf::class.java, "counts") + assertConfFailsOnField(mapOf("matrix" to listOf(listOf(1L, null))), CountsConf::class.java, "matrix") + } + + @Test + fun getConfSkipsPropertiesThatOnlyHaveAGetter() { + val conf = dagRunWithConf(mapOf("base_url" to "https://example.com")).getConf(UrlConf::class.java) + + Assertions.assertEquals("https://example.com", conf.baseUrl) + } + + @Test + fun getConfFillsCollectionsAndMapsThatOnlyHaveAGetter() { + val conf = + dagRunWithConf( + mapOf( + "TABLES" to listOf(mapOf("table_name" to "orders")), + "tables_by_name" to mapOf("Sales_Key" to mapOf("table_name" to "sales")), + ), + ).getConf(GetterOnlyTables::class.java) + + Assertions.assertEquals("orders", conf.getTables().single().tableName) + Assertions.assertEquals("sales", conf.getTablesByName()["Sales_Key"]?.tableName) + } + + @Test + fun getConfRejectsTwoKeysThatMatchOneField() { + val error = + Assertions.assertThrows(IllegalArgumentException::class.java) { + dagRunWithConf(mapOf("target_table" to "a", "TargetTable" to "b")).getConf(EtlConf::class.java) + } + + Assertions.assertEquals( + "Dag run conf keys 'target_table' and 'TargetTable' both match field EtlConf.targetTable when case and " + + "underscores are ignored, so the SDK cannot tell which one to use. Name the field exactly as the key " + + "it should take.", + error.message, + ) + } + + @Test + fun getConfRejectsNullOnlyForPrimitiveFields() { + Assertions.assertNull(dagRunWithConf(mapOf("target_table" to null)).getConf(EtlConf::class.java).targetTable) + + listOf("limit", "LIMIT").forEach { key -> + val error = + Assertions.assertThrows(IllegalArgumentException::class.java) { + dagRunWithConf(mapOf(key to null)).getConf(EtlConf::class.java) + } + + Assertions.assertEquals( + "Dag run conf key '$key' is null, but field EtlConf.limit has a primitive type that cannot be null. " + + "Declare a boxed type (e.g. Integer instead of int) to receive null.", + error.message, + ) + } + } + + @Test + fun getConfRejectsFieldsThatDifferOnlyInCaseOrUnderscores() { + val error = + Assertions.assertThrows(IllegalArgumentException::class.java) { + dagRunWithConf(emptyMap()).getConf(FoldClash::class.java) + } + + Assertions.assertEquals( + "Fields FoldClash.targetTable and FoldClash.targettable differ only in case or underscores, so the SDK " + + "cannot tell which one a Dag run conf key should go to. Rename one of them.", + error.message, + ) + } + + @Test + fun getConfNamesTheTargetClassWhenAValueDoesNotConvert() { + val error = + Assertions.assertThrows(IllegalArgumentException::class.java) { + dagRunWithConf(mapOf("limit" to "many")).getConf(EtlConf::class.java) + } + + val message = error.message.orEmpty() + Assertions.assertTrue(message.startsWith("Cannot convert the Dag run conf to org.apache.airflow.sdk.EtlConf: ")) { + message + } + Assertions.assertInstanceOf(IllegalArgumentException::class.java, error.cause) + } }