Merge branch 'main' into airflow-db2-adapter

Author: dabla

airflow-core/docs/authoring-and-scheduling/language-sdks/typescript.rst

@@ -86,13 +86,14 @@ value routes the task to the Node.js coordinator.
 TypeScript implementation
 ~~~~~~~~~~~~~~~~~~~~~~~~~
 
-A task is an ordinary (usually ``async``) function receiving ``TaskHandlerArgs``. Register it with the
-``dag_id`` and ``task_id`` it implements, then start the coordinator runtime; the registrations and the
-top-level ``await startCoordinator()`` make the module a runnable bundle entry point.
+A task is an ordinary (usually ``async``) function receiving ``TaskHandlerArgs``. Create a ``Dag`` with
+the ``dag_id`` it implements, attach each handler with ``dag.task``, collect the Dags in a ``DagRegistry``,
+then serve them to Airflow with ``serveDags``; that top-level ``await`` makes the module a runnable bundle
+entry point.
 
 .. code-block:: typescript
 
-    import { registerTask, startCoordinator, type TaskHandlerArgs } from "@apache-airflow/ts-sdk";
+    import { Dag, DagRegistry, serveDags, type TaskHandlerArgs } from "@apache-airflow/ts-sdk";
 
     export async function buildMessage({ ctx, client }: TaskHandlerArgs) {
       const upstream = await client.getXCom<string>({
@@ -103,12 +104,22 @@ top-level ``await startCoordinator()`` make the module a runnable bundle entry p
       return `${greeting ?? "hello from TypeScript"}; upstream=${upstream ?? "missing"}`;
     }
 
-    registerTask({ dagId: "typescript_example", taskId: "build_message" }, buildMessage);
+    const dag = new Dag("typescript_example");
+    dag.task("build_message", buildMessage);
 
-    await startCoordinator();
+    await serveDags(new DagRegistry(dag));
 
-The ``dagId`` passed to ``registerTask`` must match the ``dag_id`` of the Python Dag, and each ``taskId``
-must match a ``@task.stub`` function in that Dag.
+The ``dagId`` passed to ``new Dag(...)`` must match the ``dag_id`` of the Python Dag, and each ``taskId``
+passed to ``dag.task`` must match a ``@task.stub`` function in that Dag. The registry passed to
+``serveDags`` is the bundle's complete set of Dags; a second ``serveDags`` call is rejected. A Dag left out
+of the registry is not part of the packed bundle, and its tasks are marked removed at runtime.
+
+``DagRegistry`` holds no sockets and starts nothing, so a unit test can build one and dispatch a handler
+through ``registry.getTaskHandler(dagId, taskId)`` without a coordinator runtime. A bundle that collects
+its Dags across several modules can add them incrementally with ``registry.register(...)``.
+
+``new Dag`` and ``dag.task`` take a trailing options object — ``spec`` on both, plus ``inputs`` on a task.
+These are not used yet; do not set them. Any other key is rejected.
 
 .. note::
 

airflow-core/newsfragments/71704.bugfix.rst

@@ -0,0 +1 @@
+The scheduler's Dag cache is now a bounded LRU of 512 versions, so scheduler memory no longer grows with every Dag version the process has ever seen.

airflow-core/src/airflow/jobs/scheduler_job_runner.py

@@ -155,6 +155,18 @@
 TASK_STUCK_IN_QUEUED_RESCHEDULE_EVENT = "stuck in queued reschedule"
 """:meta private:"""
 
+SCHEDULER_DAG_CACHE_SIZE = 512
+"""
+Max deserialized Dag versions the scheduler keeps in memory.
+
+The scheduler reaches its DagBag through the Dag version of each active Dag run, so an
+unbounded cache retains every version the process has ever seen and grows for the life of
+the process. Sized to sit above the versions-with-runs-in-flight working set of a typical
+deployment, so eviction costs a re-fetch only where that working set is genuinely larger.
+
+:meta private:
+"""
+
 # Per-tick cap on pending AssetPartitionDagRun rows the scheduler evaluates.
 # Bounds the per-tick transaction so executor heartbeats and regular scheduling
 # aren't starved; remaining APDRs drain across subsequent ticks.
@@ -370,7 +382,7 @@ def __init__(
         if log:
             self._log = log
 
-        self.scheduler_dag_bag = DBDagBag(load_op_links=False)
+        self.scheduler_dag_bag = DBDagBag(load_op_links=False, cache_size=SCHEDULER_DAG_CACHE_SIZE)
 
         # Set of (dag_id, asset_name, asset_uri) tuples for trigger policies that
         # are permanently unreachable for the rollup window's cardinality — the

airflow-core/tests/unit/jobs/test_scheduler_job.py

@@ -64,7 +64,7 @@
 from airflow.executors.executor_utils import ExecutorName
 from airflow.executors.local_executor import LocalExecutor
 from airflow.jobs.job import Job, run_job
-from airflow.jobs.scheduler_job_runner import SchedulerJobRunner
+from airflow.jobs.scheduler_job_runner import SCHEDULER_DAG_CACHE_SIZE, SchedulerJobRunner
 from airflow.models.asset import (
     AssetActive,
     AssetAliasModel,
@@ -414,6 +414,15 @@ def test_executor_loaded_in_scheduler_job(self, mock_init_executors, mock_defaul
         assert scheduler_job.executor == mock_local_executor
         assert scheduler_job.executors == [mock_local_executor]
 
+    def test_scheduler_dag_bag_is_bounded(self):
+        """The scheduler's Dag cache must evict, or it retains every version it has ever seen."""
+        from cachetools import LRUCache
+
+        job_runner = SchedulerJobRunner(Job())
+
+        assert isinstance(job_runner.scheduler_dag_bag._dags, LRUCache)
+        assert job_runner.scheduler_dag_bag._dags.maxsize == SCHEDULER_DAG_CACHE_SIZE
+
     @pytest.mark.parametrize(
         "heartrate",
         [10, 5],

ts-sdk/README.md

@@ -33,14 +33,17 @@ runtime used to execute registered TypeScript handlers from Airflow.
 ## Task Handlers
 
 ```ts
-import { registerTask, type TaskHandlerArgs } from "@apache-airflow/ts-sdk";
+import { Dag, DagRegistry, serveDags, type TaskHandlerArgs } from "@apache-airflow/ts-sdk";
 
 export async function sayHello({ ctx, client }: TaskHandlerArgs) {
   const greeting = await client.getVariable("greeting");
   return { message: `Hello from ${ctx.taskId}: ${greeting}` };
 }
 
-registerTask({ dagId: "example_dag", taskId: "say_hello" }, sayHello);
+const dag = new Dag("example_dag");
+dag.task("say_hello", sayHello);
+
+await serveDags(new DagRegistry(dag));
 ```
 
 Non-`undefined` return values are pushed to XCom under the `"return_value"`
@@ -95,7 +98,7 @@ Airflow metadata in the bundle itself.
 TypeScript entrypoint:
 
 ```ts
-import { registerTask, startCoordinator, type TaskHandlerArgs } from "@apache-airflow/ts-sdk";
+import { Dag, DagRegistry, serveDags, type TaskHandlerArgs } from "@apache-airflow/ts-sdk";
 
 export async function extract({ client }: TaskHandlerArgs) {
   const connection = await client.getConnection("sales_db");
@@ -118,33 +121,49 @@ export async function transform({ client }: TaskHandlerArgs) {
   };
 }
 
-registerTask({ dagId: "sales_pipeline", taskId: "extract" }, extract);
-registerTask({ dagId: "sales_pipeline", taskId: "transform" }, transform);
+const salesPipeline = new Dag("sales_pipeline");
+salesPipeline.task("extract", extract);
+salesPipeline.task("transform", transform);
 
-await startCoordinator();
+await serveDags(new DagRegistry(salesPipeline));
 ```
 
 The Python stub defines the Dag dependency graph. The TypeScript handler does
-the work and uses `TaskClient` for task-time Airflow data access. Register each
-handler with the Python Dag's `dag_id` and the stub task's `task_id`. The
-handler function is the reusable task implementation; `registerTask` binds that
-handler to a Python stub Dag/task identity for coordinator mode.
+the work and uses `TaskClient` for task-time Airflow data access. Create a
+`Dag` with the Python Dag's `dag_id` and attach each handler with the stub
+task's `task_id`. The handler function is the reusable task implementation;
+`dag.task` binds that handler to a Python stub task identity, a `DagRegistry`
+collects the Dags this bundle can execute, and `serveDags` serves them to
+Airflow.
+
+`serveDags` is the entrypoint, and the registry it is given is the whole bundle:
+a Dag left out of the registry is not part of the bundle, and its tasks are
+marked removed at runtime. The registry itself holds no sockets and starts
+nothing, so a unit test can build one and dispatch through
+`registry.getTaskHandler(dagId, taskId)` without any runtime involved.
 
-For larger projects, keep one Airflow entrypoint that imports every module that
-registers tasks, then starts the coordinator:
+`new Dag` and `dag.task` take a trailing options object — `spec` on both, plus
+`inputs` on a task. These are not used yet; do not set them.
+
+For larger projects, declare each Dag in its own module and keep one Airflow
+entrypoint that serves them all:
 
 ```ts
-import "./sales/tasks";
-import "./billing/tasks";
-import { startCoordinator } from "@apache-airflow/ts-sdk";
+import { salesDag } from "./sales/dag";
+import { billingDag } from "./billing/dag";
+import { DagRegistry, serveDags } from "@apache-airflow/ts-sdk";
 
-await startCoordinator();
+await serveDags(new DagRegistry(salesDag, billingDag));
 ```
 
+A bundle that collects its Dags across several modules can add them
+incrementally with `registry.register(...)` instead of passing them all to the
+constructor.
+
 Airflow launches the bundled entrypoint with `--comm=host:port` and
-`--logs=host:port`. `startCoordinator()` connects to those sockets, receives
-the task startup message, finds the registered handler for the Dag/task pair,
-and reports the terminal task state back to Airflow.
+`--logs=host:port`. `serveDags()` connects to those sockets, receives the task
+startup message, finds the registered handler for the Dag/task pair, and
+reports the terminal task state back to Airflow.
 
 See [`example/`](example/) for a coordinator-runtime example that packs a
 bundle with `airflow-ts-pack` and uses a Python stub Dag.

ts-sdk/docs/index.md

@@ -34,19 +34,23 @@ The SDK is currently distributed as source in the `ts-sdk/` directory of the
 Apache Airflow repository. Build it there and add it as a local dependency of
 your task bundle; it is not yet published to a public npm registry.
 
-Register a task handler. Handlers receive a `TaskContext` and a `TaskClient`;
-any non-`undefined` return value is pushed to XCom under the `"return_value"`
-key by the active runtime, matching Python `@task` behavior:
+Define a Dag and register its task handlers. Handlers receive a `TaskContext`
+and a `TaskClient`; any non-`undefined` return value is pushed to XCom under
+the `"return_value"` key by the active runtime, matching Python `@task`
+behavior:
 
 ```ts
-import { registerTask, type TaskHandlerArgs } from "@apache-airflow/ts-sdk";
+import { Dag, DagRegistry, serveDags, type TaskHandlerArgs } from "@apache-airflow/ts-sdk";
 
 export async function sayHello({ ctx, client }: TaskHandlerArgs) {
   const greeting = await client.getVariable("greeting");
   return { message: `Hello from ${ctx.taskId}: ${greeting}` };
 }
 
-registerTask({ dagId: "example_dag", taskId: "say_hello" }, sayHello);
+const dag = new Dag("example_dag");
+dag.task("say_hello", sayHello);
+
+await serveDags(new DagRegistry(dag));
 ```
 
 ## Coordinators

ts-sdk/example/src/main.ts

@@ -17,9 +17,9 @@
  * under the License.
  */
 
-import { registerTask, startCoordinator, type TaskHandlerArgs } from "@apache-airflow/ts-sdk";
+import { Dag, DagRegistry, serveDags, type TaskHandlerArgs } from "@apache-airflow/ts-sdk";
 
-const DAG_ID = "typescript_example";
+const dag = new Dag("typescript_example");
 
 export async function buildMessage({ client }: TaskHandlerArgs) {
   const upstream = await client.getXCom<string>({
@@ -49,7 +49,7 @@ export async function readConnection({ client }: TaskHandlerArgs) {
   };
 }
 
-registerTask({ dagId: DAG_ID, taskId: "build_message" }, buildMessage);
-registerTask({ dagId: DAG_ID, taskId: "read_connection" }, readConnection);
+dag.task("build_message", buildMessage);
+dag.task("read_connection", readConnection);
 
-await startCoordinator();
+await serveDags(new DagRegistry(dag));

ts-sdk/src/cli/pack.ts

@@ -129,9 +129,16 @@ function readBundleManifest(bundlePath: string): BundleManifest {
       maxBuffer: MANIFEST_MAX_BUFFER_BYTES,
     });
   } catch (error) {
-    throw new Error(`Running the bundle with ${AIRFLOW_METADATA_FLAG} failed: ${String(error)}`, {
-      cause: error,
-    });
+    const stderr = (error as { stderr?: string }).stderr ?? "";
+    const reported = stderr
+      .split("\n")
+      .reverse()
+      .find((line) => /^\w*Error: /.test(line.trim()))
+      ?.trim();
+    throw new Error(
+      reported ?? `Running the bundle with ${AIRFLOW_METADATA_FLAG} failed: ${String(error)}`,
+      { cause: error },
+    );
   }
 
   // Import-time logging from user code lands on stdout too; pick the sentinel line.
@@ -143,20 +150,52 @@ function readBundleManifest(bundlePath: string): BundleManifest {
     throw new Error(`Bundle produced no ${AIRFLOW_METADATA_FLAG} output`);
   }
 
-  let manifest: BundleManifest;
+  let parsed: unknown;
   try {
-    manifest = JSON.parse(line.slice(AIRFLOW_METADATA_SENTINEL.length)) as BundleManifest;
+    parsed = JSON.parse(line.slice(AIRFLOW_METADATA_SENTINEL.length));
   } catch (error) {
     throw new Error(`Bundle produced invalid ${AIRFLOW_METADATA_FLAG} output: ${String(error)}`, {
       cause: error,
     });
   }
-  if (!manifest.supervisor_schema_version || !manifest.dags || typeof manifest.dags !== "object") {
+  if (!isBundleManifest(parsed)) {
     throw new Error(`Bundle produced incomplete ${AIRFLOW_METADATA_FLAG} output`);
   }
+  const manifest = parsed;
+  // The line is whatever the bundle printed and nothing downstream re-validates
+  // it, so check each Dag entry down to the task-id element.
+  for (const [dagId, dag] of Object.entries(manifest.dags)) {
+    if (dag == null || !isTaskIdList(dag.tasks)) {
+      throw new Error(
+        `Bundle produced ${AIRFLOW_METADATA_FLAG} output with a malformed entry for Dag "${dagId}"`,
+      );
+    }
+  }
   return manifest;
 }
 
+// The document is checked before anything is read off it: JSON.parse also yields
+// null and primitives, and `null.supervisor_schema_version` would surface as a
+// raw TypeError rather than a report about the bundle.
+function isBundleManifest(value: unknown): value is BundleManifest {
+  if (typeof value !== "object" || value === null || Array.isArray(value)) return false;
+  const { supervisor_schema_version: version, dags } = value as Partial<BundleManifest>;
+  return (
+    // Rendered into the manifest verbatim, where the schema requires a non-empty
+    // string, so a truthy number or boolean would travel to Airflow as-is.
+    typeof version === "string" &&
+    version.length > 0 &&
+    typeof dags === "object" &&
+    dags !== null &&
+    // An array would pass the typeof check and yield Dags named "0", "1", ...
+    !Array.isArray(dags)
+  );
+}
+
+function isTaskIdList(value: unknown): value is string[] {
+  return Array.isArray(value) && value.every((item) => typeof item === "string" && item.length > 0);
+}
+
 // esbuild keeps an entry hashbang as line 1, where the metadata comment must go;
 // NodeCoordinator always runs the bundle through `node`, so drop it.
 function stripShebang(bundle: string): string {
@@ -193,10 +232,16 @@ export async function runPack(argv: readonly string[]): Promise<void> {
     });
 
     const manifest = readBundleManifest(stagingPath);
-    if (Object.keys(manifest.dags).length === 0) {
-      throw new Error(
-        `${args.entry} registered no tasks; call registerTask(...) before startCoordinator()`,
-      );
+    const dagEntries = Object.entries(manifest.dags);
+    if (dagEntries.length === 0) {
+      throw new Error(`${args.entry} served no Dags; pass them to serveDags(new DagRegistry(...))`);
+    }
+    // Warn rather than fail, as airflow-go-pack does: the shared schema allows a
+    // Dag with no tasks.
+    for (const [dagId, dag] of dagEntries) {
+      if (dag.tasks.length === 0) {
+        process.stderr.write(`warning: dag ${JSON.stringify(dagId)} has no tasks\n`);
+      }
     }
 
     const metadataYaml = renderMetadataYaml({

ts-sdk/src/coordinator/index.ts

@@ -21,8 +21,11 @@
 //
 // TaskClient and related types are exported from the package root. This
 // barrel only exports coordinator-specific entry points.
+//
+// `startCoordinator` is deliberately not exported: Dag authors reach the
+// runtime through `serveDags()`, and never name the coordinator itself.
 
-export { startCoordinator, type StartCoordinatorOptions } from "./runtime.js";
+export { serveDags } from "./runtime.js";
 /** Cadwyn schema version this SDK was generated against. Not sent on
  *  the wire — exposed so callers can read it for bundle metadata,
  *  health checks, or to confirm which schema their build is pinned to. */