jason810496 commented on code in PR #69079: URL: https://github.com/apache/airflow/pull/69079#discussion_r3503185522
########## ts-sdk/src/generated/supervisor.ts: ########## @@ -0,0 +1,1871 @@ +/*! + * 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. + */ + +// AUTO-GENERATED by scripts/generate-supervisor.mjs — do not edit by hand. +// Source: ../task-sdk/src/airflow/sdk/execution_time/schema/schema.json +// +// Re-run with: pnpm run generate:supervisor + +export type Name = string; +export type Id = number; +export type Timestamp = string; +export type Extra = { + [k: string]: JsonValue; +} | null; +/** + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "JsonValue". + */ +export type JsonValue = unknown; +export type Name1 = string; +export type Uri = string; +export type Group = string; +export type Extra1 = { + [k: string]: JsonValue; +} | null; +export type RunId = string; +export type DagId = string; +export type LogicalDate = string | null; +export type StartDate = string; +export type EndDate = string | null; +export type State = string; +export type DataIntervalStart = string | null; +export type DataIntervalEnd = string | null; +export type PartitionKey = string | null; +export type CreatedDagruns = DagRunAssetReference[]; +export type SourceTaskId = string | null; +export type SourceDagId = string | null; +export type SourceRunId = string | null; +export type SourceMapIndex = number | null; +export type PartitionKey1 = string | null; +export type AssetEvents = AssetEventResponse[]; +export type Type = "AssetEventsResult"; +export type Name2 = string | null; +export type Uri1 = string | null; +export type Type1 = string; +export type Name3 = string; +export type Uri2 = string; +export type Name4 = string; +export type Uri3 = string; +export type Group1 = string; +export type Extra3 = { + [k: string]: JsonValue; +} | null; +export type Type2 = "AssetResult"; +export type Type3 = "AssetStateStoreResult"; +export type Assets = AssetResult[]; +export type Type4 = "AssetsByAliasResult"; +export type State1 = "awaiting_input" | null; +export type Timeout = string | null; +export type NextMethod = string; +export type NextKwargs = { + [k: string]: JsonValue; +} | null; +export type RenderedMapIndex = string | null; +export type Type5 = "AwaitInputTask"; +export type Name5 = string; +export type Version = string | null; +export type VersionData = { + [k: string]: unknown; +} | null; +export type Name6 = string; +export type Type6 = "ClearAssetStateStoreByName"; +export type Uri4 = string; +export type Type7 = "ClearAssetStateStoreByUri"; +export type TiId = string; +export type Type8 = "ClearTaskStateStore"; +export type ConnId = string; +export type ConnType = string; +export type Host = string | null; +export type Schema = string | null; +export type Login = string | null; +export type Password = string | null; +export type Port = number | null; +export type Extra4 = string | null; +export type Type9 = "ConnectionResult"; +export type TiId1 = string; +/** + * @minItems 1 + */ +export type Options = [string, ...string[]]; +export type Subject = string; +export type Body = string | null; +export type Defaults = string[] | null; +export type Multiple = boolean | null; +export type Params = { + [k: string]: unknown; +} | null; +export type AssignedUsers = HITLUser[] | null; +export type Id1 = string; +export type Name7 = string; +export type Type10 = "CreateHITLDetailPayload"; +export type Count = number; +export type Type11 = "DRCount"; +export type Filepath = string; +export type BundleName = string; +export type BundleVersion = string | null; +export type Msg = string | null; +export type DagId1 = string; +export type RunId1 = string; +export type DagId2 = string; +export type RunId2 = string; +export type LogicalDate1 = string | null; +export type DataIntervalStart1 = string | null; +export type DataIntervalEnd1 = string | null; +export type RunAfter = string; +export type StartDate1 = string | null; +export type EndDate1 = string | null; +export type ClearNumber = number; +/** + * Class with DagRun types. + * + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "DagRunType". + */ +export type DagRunType = + | "backfill" + | "scheduled" + | "manual" + | "operator_triggered" + | "asset_triggered" + | "asset_materialization"; +/** + * All possible states that a DagRun can be in. + * + * These are "shared" with TaskInstanceState in some parts of the code, + * so please ensure that their values always match the ones with the + * same name in TaskInstanceState. + * + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "DagRunState". + */ +export type DagRunState = "queued" | "running" | "success" | "failed"; +export type Conf = { + [k: string]: unknown; +} | null; +export type TriggeringUserName = string | null; +export type SourceTaskId1 = string | null; +export type SourceDagId1 = string | null; +export type SourceRunId1 = string | null; +export type SourceMapIndex1 = number | null; +export type SourceAliases = AssetAliasReferenceAssetEventDagRun[]; +export type Timestamp1 = string; +export type PartitionKey2 = string | null; +export type ConsumedAssetEvents = AssetEventDagRunReference[]; +export type PartitionKey3 = string | null; +export type PartitionDate = string | null; +export type Note = string | null; +export type TeamName = string | null; +export type Id2 = string; +export type TaskId = string; +export type DagId3 = string; +export type RunId3 = string; +export type TryNumber = number; +export type DagVersionId = string; +export type MapIndex = number; +export type Hostname = string | null; +export type ContextCarrier = { + [k: string]: unknown; +} | null; +export type Queue = string; +export type IsFailureCallback = boolean | null; +export type Type12 = "DagCallbackRequest"; +export type File = string; +export type BundlePath = string; +export type BundleName1 = string; +export type Filepath1 = string; +export type BundleName2 = string; +export type BundleVersion1 = string | null; +export type Msg1 = string | null; +/** + * All possible states that a Task Instance can be in. + * + * Note that None is also allowed, so always use this in a type hint with Optional. + * + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "TaskInstanceState". + */ +export type TaskInstanceState = + | "removed" + | "scheduled" + | "queued" + | "running" + | "success" + | "restarting" + | "failed" + | "up_for_retry" + | "up_for_reschedule" + | "upstream_failed" + | "skipped" + | "deferred" + | "awaiting_input"; +export type TaskRescheduleCount = number; +export type MaxTries = number; +export type Key = string; +export type Value = string | null; +export type Variables = VariableResponse[]; +export type ConnId1 = string; +export type ConnType1 = string; +export type Host1 = string | null; +export type Schema1 = string | null; +export type Login1 = string | null; +export type Password1 = string | null; +export type Port1 = number | null; +export type Extra6 = string | null; +export type Connections = ConnectionResponse[]; +export type NextMethod1 = string | null; +export type NextKwargs1 = + | { + [k: string]: unknown; + } + | string + | null; +export type XcomKeysToClear = string[]; +export type ShouldRetry = boolean; +export type StartDate2 = string | null; +export type Type13 = "TaskCallbackRequest"; +export type Filepath2 = string; +export type BundleName3 = string; +export type BundleVersion2 = string | null; +export type Msg2 = string | null; +export type EmailType = "failure" | "retry"; +export type Type14 = "EmailRequest"; +export type CallbackRequests = (DagCallbackRequest | TaskCallbackRequest | EmailRequest)[]; +export type Type15 = "DagFileParseRequest"; +export type Fileloc = string; +export type LastLoaded = string | null; +export type SerializedDags = LazyDeserializedDAG[]; +export type Warnings = unknown[] | null; +export type ImportErrors = { + [k: string]: string; +} | null; +export type Type16 = "DagFileParsingResult"; +export type DagId4 = string; +export type IsPaused = boolean; +export type BundleName4 = string | null; +export type BundleVersion3 = string | null; +export type RelativeFileloc = string | null; +export type Owners = string | null; +export type Tags = string[]; +export type NextDagrun = string | null; +export type Type17 = "DagResult"; +export type DagId5 = string; +export type RunId4 = string; +export type LogicalDate2 = string | null; +export type DataIntervalStart2 = string | null; +export type DataIntervalEnd2 = string | null; +export type RunAfter1 = string; +export type StartDate3 = string | null; +export type EndDate2 = string | null; +export type ClearNumber1 = number | null; +export type Conf1 = { + [k: string]: unknown; +} | null; +export type TriggeringUserName1 = string | null; +export type ConsumedAssetEvents1 = AssetEventDagRunReference[]; +export type PartitionKey4 = string | null; +export type PartitionDate1 = string | null; +export type Note1 = string | null; +export type TeamName1 = string | null; +export type Type18 = "DagRunResult"; +export type Type19 = "DagRunStateResult"; +export type State2 = "deferred" | null; +export type Classpath = string; +export type TriggerKwargs = + | { + [k: string]: JsonValue; + } + | string + | null; +export type TriggerTimeout = string | null; +export type Queue1 = string | null; +export type NextMethod2 = string; +export type NextKwargs2 = { + [k: string]: JsonValue; +} | null; +export type RenderedMapIndex1 = string | null; +export type Type20 = "DeferTask"; +export type Name8 = string; +export type Key1 = string; +export type Type21 = "DeleteAssetStateStoreByName"; +export type Uri5 = string; +export type Key2 = string; +export type Type22 = "DeleteAssetStateStoreByUri"; +export type TiId2 = string; +export type Key3 = string; +export type Type23 = "DeleteTaskStateStore"; +export type Key4 = string; +export type Type24 = "DeleteVariable"; +export type Key5 = string; +export type DagId6 = string; +export type RunId5 = string; +export type TaskId1 = string; +export type MapIndex1 = number | null; +export type Type25 = "DeleteXCom"; +/** + * Error types used in the API client. + */ +export type ErrorType = + | "CONNECTION_NOT_FOUND" + | "VARIABLE_NOT_FOUND" + | "XCOM_NOT_FOUND" + | "ASSET_NOT_FOUND" + | "TASK_STORE_NOT_FOUND" + | "ASSET_STORE_NOT_FOUND" + | "DAGRUN_ALREADY_EXISTS" + | "PERMISSION_DENIED" + | "GENERIC_ERROR" + | "API_SERVER_ERROR"; +export type Detail = { + [k: string]: unknown; +} | null; +export type Type26 = "ErrorResponse"; +/** + * Error types used in the API client. + * + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "ErrorType". + */ +export type ErrorType1 = + | "CONNECTION_NOT_FOUND" + | "VARIABLE_NOT_FOUND" + | "XCOM_NOT_FOUND" + | "ASSET_NOT_FOUND" + | "TASK_STORE_NOT_FOUND" + | "ASSET_STORE_NOT_FOUND" + | "DAGRUN_ALREADY_EXISTS" + | "PERMISSION_DENIED" + | "GENERIC_ERROR" + | "API_SERVER_ERROR"; +export type Name9 = string; +export type Type27 = "GetAssetByName"; +export type Uri6 = string; +export type Type28 = "GetAssetByUri"; +export type Name10 = string | null; +export type Uri7 = string | null; +export type After = string | null; +export type Before = string | null; +export type Limit = number | null; +export type Ascending = boolean; +export type Type29 = "GetAssetEventByAsset"; +export type AliasName = string; +export type After1 = string | null; +export type Before1 = string | null; +export type Limit1 = number | null; +export type Ascending1 = boolean; +export type Type30 = "GetAssetEventByAssetAlias"; +export type Name11 = string; +export type Key6 = string; +export type Type31 = "GetAssetStateStoreByName"; +export type Uri8 = string; +export type Key7 = string; +export type Type32 = "GetAssetStateStoreByUri"; +export type AliasName1 = string; +export type Type33 = "GetAssetsByAlias"; +export type ConnId2 = string; +export type Type34 = "GetConnection"; +export type DagId7 = string; +export type LogicalDates = string[] | null; +export type RunIds = string[] | null; +export type States = string[] | null; +export type Type35 = "GetDRCount"; +export type DagId8 = string; +export type Type36 = "GetDag"; +export type DagId9 = string; +export type RunId6 = string; +export type Type37 = "GetDagRun"; +export type DagId10 = string; +export type RunId7 = string; +export type Type38 = "GetDagRunState"; +export type TiId3 = string; +export type Type39 = "GetHITLDetailResponse"; +export type TiId4 = string; +export type Type40 = "GetPrevSuccessfulDagRun"; +export type DagId11 = string; +export type LogicalDate3 = string; +export type State3 = string | null; +export type Type41 = "GetPreviousDagRun"; +export type DagId12 = string; +export type TaskId2 = string; +export type LogicalDate4 = string | null; +export type MapIndex2 = number; +export type Type42 = "GetPreviousTI"; +export type DagId13 = string; +export type MapIndex3 = number | null; +export type TaskIds = string[] | null; +export type TaskGroupId = string | null; +export type LogicalDates1 = string[] | null; +export type RunIds1 = string[] | null; +export type States1 = string[] | null; +export type Type43 = "GetTICount"; +export type DagId14 = string; +export type RunId8 = string; +export type Type44 = "GetTaskBreadcrumbs"; +export type TiId5 = string; +export type TryNumber1 = number; +export type Type45 = "GetTaskRescheduleStartDate"; +export type TiId6 = string; +export type Key8 = string; +export type Type46 = "GetTaskStateStore"; +export type DagId15 = string; +export type MapIndex4 = number | null; +export type TaskIds1 = string[] | null; +export type TaskGroupId1 = string | null; +export type LogicalDates2 = string[] | null; +export type RunIds2 = string[] | null; +export type Type47 = "GetTaskStates"; +export type Key9 = string; +export type Type48 = "GetVariable"; +export type Prefix = string | null; +export type Limit2 = number; +export type Offset = number; +export type Type49 = "GetVariableKeys"; +export type Key10 = string; +export type DagId16 = string; +export type RunId9 = string; +export type TaskId3 = string; +export type MapIndex5 = number | null; +export type IncludePriorDates = boolean; +export type Type50 = "GetXCom"; +export type Key11 = string; +export type DagId17 = string; +export type RunId10 = string; +export type TaskId4 = string; +export type Type51 = "GetXComCount"; +export type Key12 = string; +export type DagId18 = string; +export type RunId11 = string; +export type TaskId5 = string; +export type Offset1 = number; +export type Type52 = "GetXComSequenceItem"; +export type Key13 = string; +export type DagId19 = string; +export type RunId12 = string; +export type TaskId6 = string; +export type Start = number | null; +export type Stop = number | null; +export type Step = number | null; +export type IncludePriorDates1 = boolean; +export type Type53 = "GetXComSequenceSlice"; +export type TiId7 = string; +/** + * @minItems 1 + */ +export type Options1 = [string, ...string[]]; +export type Subject1 = string; +export type Body1 = string | null; +export type Defaults1 = string[] | null; +export type Multiple1 = boolean | null; +export type Params1 = { + [k: string]: unknown; +} | null; +export type AssignedUsers1 = HITLUser[] | null; +export type Type54 = "HITLDetailRequestResult"; +export type InactiveAssets = AssetProfile[] | null; +export type Type55 = "InactiveAssetsResult"; +export type Name12 = string | null; +export type Type56 = "MaskSecret"; +export type Ok = boolean; +export type Type57 = "OKResponse"; +export type DataIntervalStart3 = string | null; +export type DataIntervalEnd3 = string | null; +export type StartDate4 = string | null; +export type EndDate3 = string | null; +export type Type58 = "PrevSuccessfulDagRunResult"; +export type Type59 = "PreviousDagRunResult"; +export type TaskId7 = string; +export type DagId20 = string; +export type RunId13 = string; +export type LogicalDate5 = string | null; +export type StartDate5 = string | null; +export type EndDate4 = string | null; +export type State4 = string | null; +export type TryNumber2 = number; +export type MapIndex6 = number | null; +export type Duration = number | null; +export type Type60 = "PreviousTIResult"; +export type Key14 = string; +export type Value1 = string | null; +export type Description = string | null; +export type Type61 = "PutVariable"; +export type State5 = "up_for_reschedule" | null; +export type RescheduleDate = string; +export type EndDate5 = string; +export type Type62 = "RescheduleTask"; +export type Type63 = "ResendLoggingFD"; +export type State6 = "up_for_retry" | null; +export type EndDate6 = string; +export type RenderedMapIndex2 = string | null; +export type RetryDelaySeconds = number | null; +export type RetryReason = string | null; +export type Type64 = "RetryTask"; +export type Type65 = "SentFDs"; +export type Fds = number[]; +export type Name13 = string; +export type Key15 = string; +export type Type66 = "SetAssetStateStoreByName"; +export type Uri9 = string; +export type Key16 = string; +export type Type67 = "SetAssetStateStoreByUri"; +export type Type68 = "SetRenderedFields"; +export type RenderedMapIndex3 = string; +export type Type69 = "SetRenderedMapIndex"; +export type TiId8 = string; +export type Key17 = string; +export type ExpiresAt = string | null; +export type Type70 = "SetTaskStateStore"; +export type Key18 = string; +export type DagId21 = string; +export type RunId14 = string; +export type TaskId8 = string; +export type MapIndex7 = number | null; +export type DagResult1 = boolean; +export type MappedLength = number | null; +export type Type71 = "SetXCom"; +export type Tasks = (string | [unknown, unknown])[]; +export type Type72 = "SkipDownstreamTasks"; +export type DagRelPath = string; +export type StartDate6 = string; +export type SentryIntegration = string; +export type Type73 = "StartupDetails"; +export type State7 = "success" | null; +export type EndDate7 = string; +export type TaskOutlets = AssetProfile[] | null; +export type OutletEvents = + | { + [k: string]: unknown; + }[] + | null; +export type RenderedMapIndex4 = string | null; +export type Type74 = "SucceedTask"; +export type Count1 = number; +export type Type75 = "TICount"; +export type Breadcrumbs = { + [k: string]: unknown; +}[]; +export type Type76 = "TaskBreadcrumbsResult"; +export type StartDate7 = string | null; +export type Type77 = "TaskRescheduleStartDate"; +export type State8 = "failed" | "skipped" | "removed"; +export type EndDate8 = string | null; +export type Type78 = "TaskState"; +export type RenderedMapIndex5 = string | null; +export type Type79 = "TaskStateStoreResult"; +export type Type80 = "TaskStatesResult"; +export type LogicalDate6 = string | null; +export type RunAfter2 = string | null; +export type Conf2 = { + [k: string]: unknown; +} | null; +export type ResetDagRun = boolean | null; +export type PartitionKey5 = string | null; +export type Note2 = string | null; +export type DagId22 = string; +export type DagRunId = string; +export type Type81 = "TriggerDagRun"; +export type TiId9 = string; +/** + * @minItems 1 + */ +export type ChosenOptions = [string, ...string[]]; +export type ParamsInput = { + [k: string]: unknown; +} | null; +export type Type82 = "UpdateHITLDetail"; +export type TiId10 = string; +export type Type83 = "ValidateInletsAndOutlets"; +export type Keys = string[]; +export type TotalEntries = number; +export type Type84 = "VariableKeysResult"; +export type Key19 = string; +export type Value2 = string | null; +export type Type85 = "VariableResult"; +export type Len = number; +export type Type86 = "XComCountResponse"; +export type Key20 = string; +export type Type87 = "XComResult"; +export type Type88 = "XComSequenceIndexResult"; +export type Root = JsonValue[]; +export type Type89 = "XComSequenceSliceResult"; + +export interface SupervisorWireSchema {} +/** + * Schema for AssetAliasModel used in AssetEventDagRunReference. + * + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "AssetAliasReferenceAssetEventDagRun". + */ +export interface AssetAliasReferenceAssetEventDagRun { + name: Name; +} +/** + * Asset event schema with fields that are needed for Runtime. + * + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "AssetEventResponse". + */ +export interface AssetEventResponse { + id: Id; + timestamp: Timestamp; + extra?: Extra; + asset: AssetResponse; + created_dagruns: CreatedDagruns; + source_task_id?: SourceTaskId; + source_dag_id?: SourceDagId; + source_run_id?: SourceRunId; + source_map_index?: SourceMapIndex; + partition_key?: PartitionKey1; +} +/** + * Asset schema for responses with fields that are needed for Runtime. + * + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "AssetResponse". + */ +export interface AssetResponse { + name: Name1; + uri: Uri; + group: Group; + extra?: Extra1; +} +/** + * DagRun serializer for asset responses. + * + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "DagRunAssetReference". + */ +export interface DagRunAssetReference { + run_id: RunId; + dag_id: DagId; + logical_date?: LogicalDate; + start_date: StartDate; + end_date?: EndDate; + state: State; + data_interval_start?: DataIntervalStart; + data_interval_end?: DataIntervalEnd; + partition_key?: PartitionKey; +} +/** + * Response to GetAssetEvent request. + * + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "AssetEventsResult". + */ +export interface AssetEventsResult { + asset_events: AssetEvents; + type?: Type; +} +/** + * Profile of an asset-like object. + * + * Asset will have name, uri defined, with type set to 'Asset'. + * AssetNameRef will have name defined, type set to 'AssetNameRef'. + * AssetUriRef will have uri defined, type set to 'AssetUriRef'. + * AssetAlias will have name defined, type set to 'AssetAlias'. + * + * Note that 'type' here is distinct from 'asset_type' the user declares on an + * Asset (or subclass). This field is for distinguishing between different + * asset-related types (Asset, AssetRef, or AssetAlias). + * + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "AssetProfile". + */ +export interface AssetProfile { + name?: Name2; + uri?: Uri1; + type: Type1; +} +/** + * Schema for AssetModel used in AssetEventDagRunReference. + * + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "AssetReferenceAssetEventDagRun". + */ +export interface AssetReferenceAssetEventDagRun { + name: Name3; + uri: Uri2; + extra: Extra2; +} +export interface Extra2 { + [k: string]: JsonValue; +} +/** + * Response to ReadXCom request. + * + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "AssetResult". + */ +export interface AssetResult { + name: Name4; + uri: Uri3; + group: Group1; + extra?: Extra3; + type?: Type2; +} +/** + * Response to GetAssetStateStore; wraps the generated API response for supervisor to worker comms. + * + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "AssetStateStoreResult". + */ +export interface AssetStateStoreResult { + value: JsonValue; + type?: Type3; +} +/** + * Response to GetAssetsByAlias; list of concrete assets resolved from an alias. + * + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "AssetsByAliasResult". + */ +export interface AssetsByAliasResult { + assets: Assets; + type?: Type4; +} +/** + * Park a task instance awaiting human input (Human-in-the-loop), without a trigger. + * + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "AwaitInputTask". + */ +export interface AwaitInputTask { + state?: State1; + timeout?: Timeout; + next_method: NextMethod; + next_kwargs?: NextKwargs; + rendered_map_index?: RenderedMapIndex; + type?: Type5; +} +/** + * Schema for telling task which bundle to run with. + * + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "BundleInfo". + */ +export interface BundleInfo { + name: Name5; + version?: Version; + version_data?: VersionData; +} +/** + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "ClearAssetStateStoreByName". + */ +export interface ClearAssetStateStoreByName { + name: Name6; + type?: Type6; +} Review Comment: Would it be possible to introduce a post step to normalize the type naming? Since there're quite a low of type annotated with number suffix, just wonder could we remove them (Though it should not impact the comms itself). Not sure will this be too complicated or it's not possible to make it? ########## ts-sdk/src/coordinator/client.ts: ########## @@ -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. + */ + +import type { CommChannel } from "./comm-channel.js"; +import type { LogChannel } from "./log-channel.js"; +import type { TaskContext } from "../sdk/task.js"; +import type { TaskClient } from "../sdk/client.js"; +import type { ConnectionResult, GetXComOpts, SetXComOpts } from "../sdk/client-types.js"; +import { VariableNotFoundError } from "../sdk/client.js"; +import type { + GetVariable, + GetXCom, + SetXCom, + GetConnection, + ConnectionResult as WireConnectionResult, +} from "./protocol.js"; + +function toWireMapIndex(m: number | null | undefined): number | null { + return m == null || m < 0 ? null : m; +} + +function fromWireConnection(body: WireConnectionResult): ConnectionResult { + return { + id: body.conn_id, + type: body.conn_type, + host: body.host ?? null, + schema: body.schema ?? null, + login: body.login ?? null, + password: body.password ?? null, + port: body.port ?? null, + extra: body.extra ?? null, + }; +} + +export function createCoordinatorClient( + comm: CommChannel, + ctx: TaskContext, + logs: LogChannel | null = null, +): TaskClient { + async function rpc<T>( + op: string, + expectedType: string | null, + request: unknown, + extract: (body: Record<string, unknown> | null) => T, + ): Promise<T | null> { + logs?.debug(`${op} request`); + const frame = await comm.request(request); + const err = parseFrameError(frame); + if (err) { + if (isNotFound(err)) { + logs?.debug(`${op} not found`, { error: err.code }); + return null; + } + logs?.warning(`${op} failed`, { error: err.code }); + throw new Error(`${op} failed: ${err.code}`); + } + const body = frame.body as Record<string, unknown> | null; + if (expectedType !== null && body?.type !== expectedType) { + logs?.error(`${op} unexpected response type`, { + expected: expectedType, + got: body?.type ?? null, + }); + throw new Error(`${op}: unexpected response type ${JSON.stringify(body?.type)}`); + } + logs?.debug(`${op} ok`); + return extract(body); + } + + const client: TaskClient = { + // ---- Variables ---- + + async getVariable(key: string): Promise<string | null> { + const msg: GetVariable = { type: "GetVariable", key }; + return rpc("GetVariable", "VariableResult", msg, (body) => (body!.value as string) ?? null); + }, + + async getVariableOrThrow(key: string): Promise<string> { + const value = await client.getVariable(key); + if (value == null) throw new VariableNotFoundError(key); + return value; + }, + + // ---- XCom ---- + + async getXCom<T = unknown>(opts: GetXComOpts): Promise<T | null> { + const msg: GetXCom = { + type: "GetXCom", + key: opts.key, + dag_id: opts.dagId ?? ctx.dagId, + task_id: opts.taskId ?? ctx.taskId, + run_id: opts.runId ?? ctx.runId, + map_index: toWireMapIndex(opts.mapIndex ?? ctx.mapIndex), + include_prior_dates: opts.includePriorDates ?? false, + }; + return rpc("GetXCom", "XComResult", msg, (body) => (body!.value as T) ?? null); + }, + + async setXCom(opts: SetXComOpts): Promise<void> { + const msg: SetXCom = { + type: "SetXCom", + key: opts.key, + value: opts.value, + dag_id: opts.dagId ?? ctx.dagId, + task_id: opts.taskId ?? ctx.taskId, + run_id: opts.runId ?? ctx.runId, + map_index: toWireMapIndex(opts.mapIndex ?? ctx.mapIndex), + }; Review Comment: It seems we will drop the explicit null or -1 from user. ########## ts-sdk/src/coordinator/client.ts: ########## @@ -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. + */ + +import type { CommChannel } from "./comm-channel.js"; +import type { LogChannel } from "./log-channel.js"; +import type { TaskContext } from "../sdk/task.js"; +import type { TaskClient } from "../sdk/client.js"; +import type { ConnectionResult, GetXComOpts, SetXComOpts } from "../sdk/client-types.js"; +import { VariableNotFoundError } from "../sdk/client.js"; +import type { + GetVariable, + GetXCom, + SetXCom, + GetConnection, + ConnectionResult as WireConnectionResult, +} from "./protocol.js"; + +function toWireMapIndex(m: number | null | undefined): number | null { + return m == null || m < 0 ? null : m; +} + +function fromWireConnection(body: WireConnectionResult): ConnectionResult { + return { + id: body.conn_id, + type: body.conn_type, + host: body.host ?? null, + schema: body.schema ?? null, + login: body.login ?? null, + password: body.password ?? null, + port: body.port ?? null, + extra: body.extra ?? null, + }; +} + +export function createCoordinatorClient( + comm: CommChannel, + ctx: TaskContext, + logs: LogChannel | null = null, +): TaskClient { + async function rpc<T>( + op: string, + expectedType: string | null, + request: unknown, + extract: (body: Record<string, unknown> | null) => T, + ): Promise<T | null> { + logs?.debug(`${op} request`); + const frame = await comm.request(request); + const err = parseFrameError(frame); + if (err) { + if (isNotFound(err)) { + logs?.debug(`${op} not found`, { error: err.code }); + return null; + } + logs?.warning(`${op} failed`, { error: err.code }); + throw new Error(`${op} failed: ${err.code}`); + } + const body = frame.body as Record<string, unknown> | null; + if (expectedType !== null && body?.type !== expectedType) { + logs?.error(`${op} unexpected response type`, { + expected: expectedType, + got: body?.type ?? null, + }); + throw new Error(`${op}: unexpected response type ${JSON.stringify(body?.type)}`); + } + logs?.debug(`${op} ok`); + return extract(body); + } + + const client: TaskClient = { + // ---- Variables ---- + + async getVariable(key: string): Promise<string | null> { + const msg: GetVariable = { type: "GetVariable", key }; + return rpc("GetVariable", "VariableResult", msg, (body) => (body!.value as string) ?? null); + }, + + async getVariableOrThrow(key: string): Promise<string> { + const value = await client.getVariable(key); + if (value == null) throw new VariableNotFoundError(key); + return value; + }, + + // ---- XCom ---- + + async getXCom<T = unknown>(opts: GetXComOpts): Promise<T | null> { + const msg: GetXCom = { + type: "GetXCom", + key: opts.key, + dag_id: opts.dagId ?? ctx.dagId, + task_id: opts.taskId ?? ctx.taskId, + run_id: opts.runId ?? ctx.runId, + map_index: toWireMapIndex(opts.mapIndex ?? ctx.mapIndex), + include_prior_dates: opts.includePriorDates ?? false, + }; + return rpc("GetXCom", "XComResult", msg, (body) => (body!.value as T) ?? null); + }, + + async setXCom(opts: SetXComOpts): Promise<void> { + const msg: SetXCom = { + type: "SetXCom", + key: opts.key, + value: opts.value, + dag_id: opts.dagId ?? ctx.dagId, + task_id: opts.taskId ?? ctx.taskId, + run_id: opts.runId ?? ctx.runId, + map_index: toWireMapIndex(opts.mapIndex ?? ctx.mapIndex), + }; + await rpc("SetXCom", null, msg, () => undefined); + }, + + // ---- Connections ---- + + async getConnection(connId: string): Promise<ConnectionResult | null> { + const msg: GetConnection = { type: "GetConnection", conn_id: connId }; + return rpc("GetConnection", "ConnectionResult", msg, (body) => + fromWireConnection(body as unknown as WireConnectionResult), + ); + }, + }; + return client; +} + +// -------- Error handling (two functions) -------- +// +// parseFrameError: extract a structured error from the frame (once). +// isNotFound: decide if the error means "absent" (return null to caller) +// or "failed" (throw). + +interface FrameError { + code: string; + statusCode?: number; +} + +/** Extract the error code and optional HTTP status from a response frame. + * Returns `null` for non-error frames. */ +function parseFrameError(frame: { body: unknown; error?: unknown }): FrameError | null { + // Case 1: error field on the frame itself + if (frame.error != null) { + if (typeof frame.error === "string") return { code: frame.error }; + if (typeof frame.error === "object") { + const e = frame.error as Record<string, unknown>; + if (typeof e.error === "string") { + const detail = e.detail as { status_code?: number } | undefined; + return { code: e.error, statusCode: detail?.status_code }; + } + } + } + // Case 2: ErrorResponse in body + const body = frame.body as Record<string, unknown> | null; + if (body?.type === "ErrorResponse" && typeof body.error === "string") { + const detail = body.detail as { status_code?: number } | undefined; + return { code: body.error, statusCode: detail?.status_code }; + } + return null; +} + +// Exact supervisor ErrorType codes that mean "absent", not "failed". +const NOT_FOUND_CODES = new Set(["VARIABLE_NOT_FOUND", "XCOM_NOT_FOUND", "CONNECTION_NOT_FOUND"]); + +/** Is this error a "not found" (caller should get null, not a throw)? + * Covers both dedicated NOT_FOUND codes and API_SERVER_ERROR with 404. */ +function isNotFound(err: FrameError): boolean { + if (NOT_FOUND_CODES.has(err.code)) return true; + // The supervisor wraps API server 404s as API_SERVER_ERROR with + // detail.status_code=404 (supervisor.py: WatchedSubprocess.handle_requests). + // Dag / Dag run lookups hit this path. + return err.code === "API_SERVER_ERROR" && err.statusCode === 404; Review Comment: Worthwhile to add TODO for this part as folding all the not found errors as null is not accurate (it's fine for now if only support XCom, Variable, Connection, but it doesn't seem fine in the future). ########## ts-sdk/tests/coordinator/log-channel.test.ts: ########## @@ -0,0 +1,133 @@ +/*! + * 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. + */ + +import * as net from "node:net"; +import { afterEach, beforeEach, describe, expect, it } from "vitest"; +import { LogChannel } from "../../src/coordinator/log-channel.js"; + +interface Fixture { + server: net.Server; + port: number; + received: Buffer[]; + sockClosed: Promise<void>; +} + +async function makeServer(): Promise<Fixture> { + const received: Buffer[] = []; + let resolveSockClosed: () => void; + const sockClosed = new Promise<void>((r) => { + resolveSockClosed = r; + }); + const server = net.createServer((sock) => { + sock.on("data", (chunk) => received.push(chunk)); + sock.on("close", () => resolveSockClosed()); + }); Review Comment: We did protect the socket for comms with `sock.on("error"` but we miss for the log socket here. -- 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]
