Skip to content

Warn on suspicious Dag and task IDs at Java SDK build time#69937

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Andrushika wants to merge 1 commit into
base:main
Choose a base branch
from
Open

Warn on suspicious Dag and task IDs at Java SDK build time #69937

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 9 additions & 0 deletions java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Bundle.kt
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -25,13 +25,22 @@ package org.apache.airflow.sdk
* Build a [Bundle] by implementing [BundleBuilder], then pass it to
* [Server.serve] to start accepting task-execution requests.
*
* Dag and task IDs the Airflow server would reject (longer than 250
* characters, or containing anything other than letters, digits, dashes,
* dots, and underscores) produce a best-effort warning in the task logs;
* the server remains the authoritative validator.
*
* @property dags All registered Dags keyed by [DagDef.id].
* @throws IllegalArgumentException if any two Dags share the same ID.
*/
class Bundle(
dags: Iterable<DagDef>,
) {
internal val dags: Map<String, DagDef> = dags.associateByDagId()

init {
IdValidation.warnOnSuspiciousIds(this.dags.values)
}
}

private fun Iterable<DagDef>.associateByDagId(): Map<String, DagDef> {
Expand Down
9 changes: 6 additions & 3 deletions java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/DagDef.kt
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -36,13 +36,16 @@ import kotlin.Throws
* .addTask("load", Load.class);
* ```
*
* @param id Dag identifier. Must contain only ASCII alphanumeric characters,
* dashes, dots, or underscores; must be unique within a [Bundle].
* Dag and task IDs are validated authoritatively by the Airflow server. This
* class accepts any ID; IDs the server would reject produce a best-effort
* warning when the Dags are assembled into a [Bundle].
*
* @param id Dag identifier; must be unique within a [Bundle].
*
* @see Builder.Dag
*/
class DagDef(
val id: String, // TODO: charset check?
val id: String,
) {
internal val tasks = linkedMapOf<String, TaskDef>()

Expand Down
71 changes: 71 additions & 0 deletions java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/IdValidation.kt
View file
Open in desktop
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,71 @@
/*
* 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

import org.apache.airflow.sdk.execution.Logger

/**
* Best-effort checks of Dag and task IDs against the rules the Airflow server
* enforces (`airflow.utils.helpers.validate_key`). Warn-only: the server
* validates authoritatively, and checks like the `..` one depend on server
* configuration this process cannot see.
*/
internal object IdValidation {
private const val MAX_ID_LENGTH = 250
private val ID_REGEX = Regex("""^[\p{L}\p{N}_.-]+$""")

private val logger = Logger(IdValidation::class)

fun warnOnSuspiciousIds(dags: Iterable<DagDef>) {
for (dag in dags) {
warnOnSuspiciousId("Dag id", dag.id, mapOf("dag_id" to dag.id))
for (taskId in dag.tasks.keys) {
warnOnSuspiciousId("Task id", taskId, mapOf("dag_id" to dag.id, "task_id" to taskId))
}
}
}

private fun warnOnSuspiciousId(
label: String,
id: String,
arguments: Map<String, Any>,
) {
val length = id.codePointCount(0, id.length)
if (length > MAX_ID_LENGTH) {
logger.warning(
"$label is longer than $MAX_ID_LENGTH characters; the Airflow server will reject it",
arguments + ("length" to length),
)
}
if (!ID_REGEX.matches(id)) {
logger.warning(
"$label must be made of alphanumeric characters, dashes, dots, and underscores; " +
"the Airflow server will reject it",
arguments,
)
} else if (id.contains("..")) {
logger.warning(
"$label contains '..'; the Airflow server will reject it " +
"unless [core] allow_double_dot_in_ids is enabled",
arguments,
)
}
}
}
31 changes: 31 additions & 0 deletions java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/BundleTest.kt
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -19,11 +19,26 @@

package org.apache.airflow.sdk

import org.apache.airflow.sdk.execution.Level
import org.apache.airflow.sdk.execution.LogCapture
import org.junit.jupiter.api.Assertions
import org.junit.jupiter.api.BeforeEach
import org.junit.jupiter.api.DisplayName
import org.junit.jupiter.api.Test

internal class BundleTest {
private class NoOp : Task {
override fun execute(
context: Context,
client: Client,
) = Unit
}

@BeforeEach
fun setUp() {
LogCapture.drain()
}

@Test
@DisplayName("Should index dags by dagId")
fun shouldIndexDagsByDagId() {
Expand All@@ -32,6 +47,22 @@ internal class BundleTest {
val bundle = Bundle(listOf(dag))

Assertions.assertEquals(mapOf("dag" to dag), bundle.dags)
Assertions.assertEquals(emptyList<Any>(), LogCapture.drain())
}

@Test
@DisplayName("Should warn on dag and task ids the server would reject")
fun shouldWarnOnSuspiciousIds() {
Bundle(listOf(DagDef("bad dag").addTask("bad task", NoOp::class.java)))

val warnings = LogCapture.drain()
Assertions.assertEquals(2, warnings.size)
Assertions.assertEquals(setOf(Level.WARNING), warnings.map { it.level }.toSet())
Assertions.assertEquals(mapOf<String, Any?>("dag_id" to "bad dag"), warnings[0].arguments)
Assertions.assertEquals(
mapOf<String, Any?>("dag_id" to "bad dag", "task_id" to "bad task"),
warnings[1].arguments,
)
}

@Test
Expand Down
127 changes: 127 additions & 0 deletions java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/IdValidationTest.kt
View file
Open in desktop
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,127 @@
/*
* 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

import org.apache.airflow.sdk.execution.CapturedLogMessage
import org.apache.airflow.sdk.execution.Level
import org.apache.airflow.sdk.execution.LogCapture
import org.junit.jupiter.api.Assertions.assertEquals
import org.junit.jupiter.api.BeforeEach
import org.junit.jupiter.api.DisplayName
import org.junit.jupiter.api.Test

private const val LOGGER_NAME = "org.apache.airflow.sdk.IdValidation"

private fun dagTooLongWarning(
id: String,
length: Int,
) = CapturedLogMessage(
Level.WARNING,
LOGGER_NAME,
"Dag id is longer than 250 characters; the Airflow server will reject it",
mapOf("dag_id" to id, "length" to length),
)

private fun dagCharsetWarning(id: String) =
CapturedLogMessage(
Level.WARNING,
LOGGER_NAME,
"Dag id must be made of alphanumeric characters, dashes, dots, and underscores; " +
"the Airflow server will reject it",
mapOf("dag_id" to id),
)

private fun dagDoubleDotWarning(id: String) =
CapturedLogMessage(
Level.WARNING,
LOGGER_NAME,
"Dag id contains '..'; the Airflow server will reject it " +
"unless [core] allow_double_dot_in_ids is enabled",
mapOf("dag_id" to id),
)

internal class IdValidationTest {
private class NoOp : Task {
override fun execute(
context: Context,
client: Client,
) = Unit
}

@BeforeEach
fun setUp() {
LogCapture.drain()
}

@Test
@DisplayName("dag id warnings — exact events across every branch")
fun dagIdWarnings() {
val astral = "𠀀"
val tooLongAndInvalid = "a".repeat(250) + " b"
val cases: List<Pair<String, List<CapturedLogMessage>>> =
listOf(
"simple" to emptyList(),
"with-dash" to emptyList(),
"with.dot" to emptyList(),
"with_underscore" to emptyList(),
"0numeric" to emptyList(),
"café_dag" to emptyList(),
"任務" to emptyList(),
"a".repeat(250) to emptyList(),
astral.repeat(250) to emptyList(),
"a".repeat(251) to listOf(dagTooLongWarning("a".repeat(251), 251)),
"任".repeat(251) to listOf(dagTooLongWarning("任".repeat(251), 251)),
astral.repeat(251) to listOf(dagTooLongWarning(astral.repeat(251), 251)),
"with space" to listOf(dagCharsetWarning("with space")),
"with/slash" to listOf(dagCharsetWarning("with/slash")),
"with:colon" to listOf(dagCharsetWarning("with:colon")),
"with\ttab" to listOf(dagCharsetWarning("with\ttab")),
"a..b c" to listOf(dagCharsetWarning("a..b c")),
"a..b" to listOf(dagDoubleDotWarning("a..b")),
tooLongAndInvalid to
listOf(dagTooLongWarning(tooLongAndInvalid, 252), dagCharsetWarning(tooLongAndInvalid)),
)
cases.forEach { (id, expected) ->
IdValidation.warnOnSuspiciousIds(listOf(DagDef(id)))
assertEquals(expected, LogCapture.drain(), "id=$id")
}
}

@Test
@DisplayName("a task warning carries its dag id")
fun taskWarningCarriesDagId() {
val dag = DagDef("my_dag").addTask("bad task", NoOp::class.java)

IdValidation.warnOnSuspiciousIds(listOf(dag))

assertEquals(
listOf(
CapturedLogMessage(
Level.WARNING,
LOGGER_NAME,
"Task id must be made of alphanumeric characters, dashes, dots, and underscores; " +
"the Airflow server will reject it",
mapOf("dag_id" to "my_dag", "task_id" to "bad task"),
),
),
LogCapture.drain(),
)
}
}