shivaam commented on code in PR #74484: URL: https://github.com/apache/airflow/pull/74484#discussion_r4235807706
########## ts-sdk/src/hitl/operator.ts: ########## @@ -0,0 +1,279 @@ +/*! + * 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. + */ + +// Mirrors `HITLOperator.execute` / `execute_complete` and `ApprovalOperator.execute_complete` +// in the standard provider, on the Airflow 3.3+ path that parks the task in `awaiting_input`. +// +// The task runs twice, in two processes. The first run writes the request and parks; the +// second starts with `next_method` set and the response in `next_kwargs.event`. +// +// As `HITLOperator` does, the second run reads the options, the params and the reject policy from +// the task it is started for, and works out what a rejection skips from that task's Dag; nothing +// is carried over from the first run. A versioned Dag bundle starts both runs from the same +// version. With an unversioned bundle, a bundle deployed while the task waits is the one the +// response is read against, as it is for any task that runs after a deployment. + +import type { JsonValue } from "../sdk/client-types.js"; +import { isPlainRecord, type Dag } from "../sdk/dag.js"; +import { getDagDownstreamTaskIds } from "../sdk/dag-graph.js"; +import { brandOperator, type OperatorContext, type OperatorOutcome } from "../sdk/operator.js"; +import { runInTaskScope } from "../sdk/task.js"; +import type { HITLResult, HITLTask, HITLTaskFields, HITLText, HITLUser } from "./spec.js"; + +export const APPROVE = "Approve"; +export const REJECT = "Reject"; + +/** Build the operator for a task whose options `spec.ts` has checked. */ +export function createHITLTask<TArgs extends object | void>( + fields: HITLTaskFields, +): HITLTask<TArgs> { + const isApproval = fields.kind === "approval"; + const task: HITLTask<never> = { + ...fields, + operatorName: isApproval ? "ApprovalOperator" : "HITLOperator", + // An approval skips what follows it on "Reject". + canSkipDownstream: isApproval, + requiresTaskId: true, + label: "human-in-the-loop", + taskIdExample: isApproval + ? 'dag.task("sign_off", approval({ subject: "..." }))' + : 'dag.task("choose", hitl({ subject: "...", options: ["..."] }))', + executionTimeoutAlternative: "responseTimeout", + execute: (op) => executeHITL(task, op), + executeComplete: (op, event) => resumeHITL(task, op, event), + }; + return brandOperator(task) as HITLTask<TArgs>; +} + +/** The first run: write the request, then park the task until a response arrives. */ +async function executeHITL(task: HITLTask<never>, op: OperatorContext): Promise<OperatorOutcome> { + const { details, ctx, client, logs } = op; + // The inputs are read here only, for the text: the resumed run must not depend on upstream + // XComs still being there. + const args = await op.resolveArgs(task.argNames); + const render = (text: HITLText<never, string | null>) => + runInTaskScope({ ctx, client }, () => renderText(text, args)); + const subject = await render(task.subject); + const body = task.body === undefined ? null : await render(task.body); + if (typeof subject !== "string" || subject.length === 0) { + return op.fail(`The subject of task "${ctx.taskId}" has to be a non-empty string`); + } + if (body !== null && body !== undefined && typeof body !== "string") { + return op.fail(`The body of task "${ctx.taskId}" has to be a string or null`); + } + + const options = task.options as [string, ...string[]]; + const defaults = task.defaults === undefined ? null : [...task.defaults]; + const assignedUsers = task.assignedUsers.map((user) => ({ id: user.id, name: user.name })); + await client.createHITLDetail({ + ti_id: details.ti.id, + options, + subject, + body: body ?? null, + defaults, + multiple: task.multiple, + params: serializeParams(task.params), + assigned_users: assignedUsers, + }); + + logs.info("Waiting for response", { task_id: ctx.taskId }); + return op.awaitInput({ timeoutSeconds: task.responseTimeout }); +} + +/** + * The params as `Param.serialize()` writes them. The source is always "task": + * `HITLOperator` drops the params that come from the Dag before it sends them. + */ +function serializeParams(params: HITLTask<never>["params"]): Record<string, JsonValue> { + return Object.fromEntries( + Object.entries(params).map(([name, { value, description, schema }]) => [ + name, + { value, description: description ?? null, schema: schema ?? {}, source: "task" }, + ]), + ); +} + +/** The second run: read the response against the task's options, then finish. */ +async function resumeHITL( + task: HITLTask<never>, + op: OperatorContext, + event: unknown, +): Promise<OperatorOutcome> { + const { dag, ctx, logs } = op; + // Throws, failing the task, when the answer cannot be read or is not allowed. + const result = readAnswer(event, task); + const responder = result.responded_by_user?.name ?? "the response timeout default"; + logs.info("Received response", { chosen_options: result.chosen_options, responder }); + + // `ApprovalOperator.execute_complete`: on "Reject", fail, or skip what follows. + if (task.kind === "approval" && result.chosen_options[0] === REJECT) { + if (task.failOnReject) return op.fail(`Rejected by ${responder}`); + const skipped = getSkipTargets(dag, ctx.taskId, task.ignoreDownstreamTriggerRules); + if (skipped.length > 0) { + // `SkipMixin.skip` ends the task as soon as it skips, so Python pushes no response then. + logs.info("Skipping downstream tasks", { task_ids: skipped }); + await op.skip(skipped); + return op.succeed(); + } + logs.info("No downstream tasks; nothing to do."); + } else if (task.kind === "approval") { + logs.info("Approved. Proceeding with downstream tasks..."); + } + return op.succeed(result as unknown as JsonValue); Review Comment: HITLResult and HITLUser are type aliases and the cast is gone. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
