diff --git a/task-sdk/src/airflow/sdk/execution_time/schema/schema.json b/task-sdk/src/airflow/sdk/execution_time/schema/schema.json index 0ec8fe4e49a1c..13b107d3b43db 100644 --- a/task-sdk/src/airflow/sdk/execution_time/schema/schema.json +++ b/task-sdk/src/airflow/sdk/execution_time/schema/schema.json @@ -34,7 +34,14 @@ "anyOf": [ { "additionalProperties": { - "$ref": "#/$defs/JsonValue" + "anyOf": [ + { + "$ref": "#/$defs/JsonValue" + }, + { + "type": "null" + } + ] }, "type": "object" }, @@ -187,34 +194,6 @@ "title": "AssetProfile", "type": "object" }, - "AssetReferenceAssetEventDagRun": { - "additionalProperties": false, - "description": "Schema for AssetModel used in AssetEventDagRunReference.", - "properties": { - "name": { - "title": "Name", - "type": "string" - }, - "uri": { - "title": "Uri", - "type": "string" - }, - "extra": { - "additionalProperties": { - "$ref": "#/$defs/JsonValue" - }, - "title": "Extra", - "type": "object" - } - }, - "required": [ - "name", - "uri", - "extra" - ], - "title": "AssetReferenceAssetEventDagRun", - "type": "object" - }, "AssetResponse": { "description": "Asset schema for responses with fields that are needed for Runtime.", "properties": { @@ -234,7 +213,14 @@ "anyOf": [ { "additionalProperties": { - "$ref": "#/$defs/JsonValue" + "anyOf": [ + { + "$ref": "#/$defs/JsonValue" + }, + { + "type": "null" + } + ] }, "type": "object" }, @@ -273,7 +259,14 @@ "anyOf": [ { "additionalProperties": { - "$ref": "#/$defs/JsonValue" + "anyOf": [ + { + "$ref": "#/$defs/JsonValue" + }, + { + "type": "null" + } + ] }, "type": "object" }, @@ -304,7 +297,14 @@ "description": "Response to GetAssetStateStore; wraps the generated API response for supervisor to worker comms.", "properties": { "value": { - "$ref": "#/$defs/JsonValue" + "anyOf": [ + { + "$ref": "#/$defs/JsonValue" + }, + { + "type": "null" + } + ] }, "type": { "const": "AssetStateStoreResult", @@ -347,17 +347,10 @@ "description": "Park a task instance awaiting human input (Human-in-the-loop), without a trigger.", "properties": { "state": { - "anyOf": [ - { - "const": "awaiting_input", - "type": "string" - }, - { - "type": "null" - } - ], + "const": "awaiting_input", "default": "awaiting_input", - "title": "State" + "title": "State", + "type": "string" }, "timeout": { "anyOf": [ @@ -380,7 +373,14 @@ "anyOf": [ { "additionalProperties": { - "$ref": "#/$defs/JsonValue" + "anyOf": [ + { + "$ref": "#/$defs/JsonValue" + }, + { + "type": "null" + } + ] }, "type": "object" }, @@ -513,6 +513,97 @@ "title": "ClearTaskStateStore", "type": "object" }, + "ConnectionResponse": { + "description": "Connection schema for responses with fields that are needed for Runtime.", + "properties": { + "conn_id": { + "title": "Conn Id", + "type": "string" + }, + "conn_type": { + "title": "Conn Type", + "type": "string" + }, + "host": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Host" + }, + "schema": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Schema" + }, + "login": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Login" + }, + "password": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Password" + }, + "port": { + "anyOf": [ + { + "type": "integer" + }, + { + "type": "null" + } + ], + "title": "Port" + }, + "extra": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Extra" + } + }, + "required": [ + "conn_id", + "conn_type", + "host", + "schema", + "login", + "password", + "port", + "extra" + ], + "title": "ConnectionResponse", + "type": "object" + }, "ConnectionResult": { "properties": { "conn_id": { @@ -532,7 +623,6 @@ "type": "null" } ], - "default": null, "title": "Host" }, "schema": { @@ -544,7 +634,6 @@ "type": "null" } ], - "default": null, "title": "Schema" }, "login": { @@ -556,7 +645,6 @@ "type": "null" } ], - "default": null, "title": "Login" }, "password": { @@ -568,7 +656,6 @@ "type": "null" } ], - "default": null, "title": "Password" }, "port": { @@ -580,7 +667,6 @@ "type": "null" } ], - "default": null, "title": "Port" }, "extra": { @@ -592,7 +678,6 @@ "type": "null" } ], - "default": null, "title": "Extra" }, "type": { @@ -604,7 +689,13 @@ }, "required": [ "conn_id", - "conn_type" + "conn_type", + "host", + "schema", + "login", + "password", + "port", + "extra" ], "title": "ConnectionResult", "type": "object" @@ -957,7 +1048,6 @@ "type": "null" } ], - "default": null, "title": "Bundle Name" }, "bundle_version": { @@ -969,7 +1059,6 @@ "type": "null" } ], - "default": null, "title": "Bundle Version" }, "relative_fileloc": { @@ -981,7 +1070,6 @@ "type": "null" } ], - "default": null, "title": "Relative Fileloc" }, "owners": { @@ -993,7 +1081,6 @@ "type": "null" } ], - "default": null, "title": "Owners" }, "tags": { @@ -1013,7 +1100,6 @@ "type": "null" } ], - "default": null, "title": "Next Dagrun" }, "type": { @@ -1026,7 +1112,12 @@ "required": [ "dag_id", "is_paused", - "tags" + "bundle_name", + "bundle_version", + "relative_fileloc", + "owners", + "tags", + "next_dagrun" ], "title": "DagResult", "type": "object" @@ -1053,7 +1144,6 @@ "type": "null" } ], - "default": null, "title": "Logical Date" }, "start_date": { @@ -1071,7 +1161,6 @@ "type": "null" } ], - "default": null, "title": "End Date" }, "state": { @@ -1088,7 +1177,6 @@ "type": "null" } ], - "default": null, "title": "Data Interval Start" }, "data_interval_end": { @@ -1101,7 +1189,6 @@ "type": "null" } ], - "default": null, "title": "Data Interval End" }, "partition_key": { @@ -1113,15 +1200,19 @@ "type": "null" } ], - "default": null, "title": "Partition Key" } }, "required": [ "run_id", "dag_id", + "logical_date", "start_date", - "state" + "end_date", + "state", + "data_interval_start", + "data_interval_end", + "partition_key" ], "title": "DagRunAssetReference", "type": "object" @@ -1176,7 +1267,6 @@ "type": "null" } ], - "default": null, "title": "Logical Date" }, "data_interval_start": { @@ -1189,7 +1279,6 @@ "type": "null" } ], - "default": null, "title": "Data Interval Start" }, "data_interval_end": { @@ -1202,7 +1291,6 @@ "type": "null" } ], - "default": null, "title": "Data Interval End" }, "run_after": { @@ -1220,7 +1308,6 @@ "type": "null" } ], - "default": null, "title": "Start Date" }, "end_date": { @@ -1233,7 +1320,6 @@ "type": "null" } ], - "default": null, "title": "End Date" }, "clear_number": { @@ -1295,7 +1381,6 @@ "type": "null" } ], - "default": null, "title": "Partition Key" }, "partition_date": { @@ -1345,10 +1430,16 @@ "required": [ "dag_id", "run_id", + "logical_date", + "data_interval_start", + "data_interval_end", "run_after", + "start_date", + "end_date", "run_type", "state", - "consumed_asset_events" + "consumed_asset_events", + "partition_key" ], "title": "DagRunResult", "type": "object" @@ -1400,17 +1491,10 @@ "description": "Update a task instance state to deferred.", "properties": { "state": { - "anyOf": [ - { - "const": "deferred", - "type": "string" - }, - { - "type": "null" - } - ], + "const": "deferred", "default": "deferred", - "title": "State" + "title": "State", + "type": "string" }, "classpath": { "title": "Classpath", @@ -1420,7 +1504,14 @@ "anyOf": [ { "additionalProperties": { - "$ref": "#/$defs/JsonValue" + "anyOf": [ + { + "$ref": "#/$defs/JsonValue" + }, + { + "type": "null" + } + ] }, "type": "object" }, @@ -1467,7 +1558,14 @@ "anyOf": [ { "additionalProperties": { - "$ref": "#/$defs/JsonValue" + "anyOf": [ + { + "$ref": "#/$defs/JsonValue" + }, + { + "type": "null" + } + ] }, "type": "object" }, @@ -3388,17 +3486,10 @@ "description": "Update a task instance state to reschedule/up_for_reschedule.", "properties": { "state": { - "anyOf": [ - { - "const": "up_for_reschedule", - "type": "string" - }, - { - "type": "null" - } - ], + "const": "up_for_reschedule", "default": "up_for_reschedule", - "title": "State" + "title": "State", + "type": "string" }, "reschedule_date": { "format": "date-time", @@ -3441,17 +3532,10 @@ "description": "Update a task instance state to up_for_retry.", "properties": { "state": { - "anyOf": [ - { - "const": "up_for_retry", - "type": "string" - }, - { - "type": "null" - } - ], + "const": "up_for_retry", "default": "up_for_retry", - "title": "State" + "title": "State", + "type": "string" }, "end_date": { "format": "date-time", @@ -3825,17 +3909,10 @@ "description": "Update a task's state to success. Includes task_outlets and outlet_events for registering asset events.", "properties": { "state": { - "anyOf": [ - { - "const": "success", - "type": "string" - }, - { - "type": "null" - } - ], + "const": "success", "default": "success", - "title": "State" + "title": "State", + "type": "string" }, "end_date": { "format": "date-time", @@ -4132,7 +4209,14 @@ "description": "Response to GetTaskStateStore; wraps the generated API response for supervisor to worker comms.", "properties": { "value": { - "$ref": "#/$defs/JsonValue" + "anyOf": [ + { + "$ref": "#/$defs/JsonValue" + }, + { + "type": "null" + } + ] }, "type": { "const": "TaskStateStoreResult", @@ -4357,6 +4441,33 @@ "title": "VariableKeysResult", "type": "object" }, + "VariableResponse": { + "additionalProperties": false, + "description": "Variable schema for responses with fields that are needed for Runtime.", + "properties": { + "key": { + "title": "Key", + "type": "string" + }, + "value": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Value" + } + }, + "required": [ + "key", + "value" + ], + "title": "VariableResponse", + "type": "object" + }, "VariableResult": { "additionalProperties": false, "properties": { @@ -4373,7 +4484,6 @@ "type": "null" } ], - "default": null, "title": "Value" }, "type": { @@ -4384,7 +4494,8 @@ } }, "required": [ - "key" + "key", + "value" ], "title": "VariableResult", "type": "object" @@ -4416,7 +4527,14 @@ "type": "string" }, "value": { - "$ref": "#/$defs/JsonValue" + "anyOf": [ + { + "$ref": "#/$defs/JsonValue" + }, + { + "type": "null" + } + ] }, "type": { "const": "XComResult", @@ -4472,97 +4590,6 @@ "title": "XComSequenceSliceResult", "type": "object" }, - "ConnectionResponse": { - "description": "Connection schema for responses with fields that are needed for Runtime.", - "properties": { - "conn_id": { - "title": "Conn Id", - "type": "string" - }, - "conn_type": { - "title": "Conn Type", - "type": "string" - }, - "host": { - "anyOf": [ - { - "type": "string" - }, - { - "type": "null" - } - ], - "title": "Host" - }, - "schema": { - "anyOf": [ - { - "type": "string" - }, - { - "type": "null" - } - ], - "title": "Schema" - }, - "login": { - "anyOf": [ - { - "type": "string" - }, - { - "type": "null" - } - ], - "title": "Login" - }, - "password": { - "anyOf": [ - { - "type": "string" - }, - { - "type": "null" - } - ], - "title": "Password" - }, - "port": { - "anyOf": [ - { - "type": "integer" - }, - { - "type": "null" - } - ], - "title": "Port" - }, - "extra": { - "anyOf": [ - { - "type": "string" - }, - { - "type": "null" - } - ], - "title": "Extra" - } - }, - "required": [ - "conn_id", - "conn_type", - "host", - "schema", - "login", - "password", - "port", - "extra" - ], - "title": "ConnectionResponse", - "type": "object" - }, "AssetEventDagRunReference": { "additionalProperties": false, "description": "Schema for AssetEvent model used in DagRun.", @@ -4659,6 +4686,34 @@ "title": "AssetEventDagRunReference", "type": "object" }, + "AssetReferenceAssetEventDagRun": { + "additionalProperties": false, + "description": "Schema for AssetModel used in AssetEventDagRunReference.", + "properties": { + "name": { + "title": "Name", + "type": "string" + }, + "uri": { + "title": "Uri", + "type": "string" + }, + "extra": { + "additionalProperties": { + "$ref": "#/$defs/JsonValue" + }, + "title": "Extra", + "type": "object" + } + }, + "required": [ + "name", + "uri", + "extra" + ], + "title": "AssetReferenceAssetEventDagRun", + "type": "object" + }, "DagRun": { "additionalProperties": false, "description": "Schema for DagRun model with minimal required fields needed for Runtime.", @@ -4914,6 +4969,11 @@ "title": "Should Retry", "type": "boolean" }, + "has_mapped_dependants": { + "default": false, + "title": "Has Mapped Dependants", + "type": "boolean" + }, "start_date": { "anyOf": [ { @@ -5010,33 +5070,6 @@ ], "title": "TaskInstance", "type": "object" - }, - "VariableResponse": { - "additionalProperties": false, - "description": "Variable schema for responses with fields that are needed for Runtime.", - "properties": { - "key": { - "title": "Key", - "type": "string" - }, - "value": { - "anyOf": [ - { - "type": "string" - }, - { - "type": "null" - } - ], - "title": "Value" - } - }, - "required": [ - "key", - "value" - ], - "title": "VariableResponse", - "type": "object" } } -} +} \ No newline at end of file diff --git a/ts-sdk/src/coordinator/runtime.ts b/ts-sdk/src/coordinator/runtime.ts index a3a057b4c0176..64b9e831dad00 100644 --- a/ts-sdk/src/coordinator/runtime.ts +++ b/ts-sdk/src/coordinator/runtime.ts @@ -376,6 +376,7 @@ function buildContext(details: StartupDetails, signal: AbortSignal): TaskContext runId: details.ti.run_id, tryNumber: details.ti.try_number, mapIndex: details.ti.map_index ?? -1, + hasMappedDependants: details.ti_context.has_mapped_dependants ?? false, signal, }; } diff --git a/ts-sdk/src/generated/supervisor.ts b/ts-sdk/src/generated/supervisor.ts index ab2632831ab95..1eaef4f64e2a2 100644 --- a/ts-sdk/src/generated/supervisor.ts +++ b/ts-sdk/src/generated/supervisor.ts @@ -26,18 +26,13 @@ export type Name = string; export type Id = number; export type Timestamp = string; export type Extra = { - [k: string]: JsonValue; + [k: string]: unknown; } | 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; + [k: string]: unknown; } | null; export type RunId = string; export type DagId = string; @@ -61,32 +56,30 @@ 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; +export type Extra2 = { + [k: string]: unknown; } | 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 State1 = "awaiting_input"; export type Timeout = string | null; export type NextMethod = string; export type NextKwargs = { - [k: string]: JsonValue; + [k: string]: unknown; } | null; export type RenderedMapIndex = string | null; export type Type5 = "AwaitInputTask"; -export type Name5 = string; +export type Name4 = string; export type Version = string | null; export type VersionData = { [k: string]: unknown; } | null; -export type Name6 = string; +export type Name5 = string; export type Type6 = "ClearAssetStateStoreByName"; -export type Uri4 = string; +export type Uri3 = string; export type Type7 = "ClearAssetStateStoreByUri"; export type TiId = string; export type Type8 = "ClearTaskStateStore"; @@ -97,6 +90,14 @@ export type Schema = string | null; export type Login = string | null; export type Password = string | null; export type Port = number | null; +export type Extra3 = string | null; +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 Extra4 = string | null; export type Type9 = "ConnectionResult"; export type TiId1 = string; @@ -113,7 +114,7 @@ export type Params = { } | null; export type AssignedUsers = HITLUser[] | null; export type Id1 = string; -export type Name7 = string; +export type Name6 = string; export type Type10 = "CreateHITLDetailPayload"; export type Count = number; export type Type11 = "DRCount"; @@ -142,12 +143,7 @@ export type ClearNumber = number; * via the `definition` "DagRunType". */ export type DagRunType = - | "backfill" - | "scheduled" - | "manual" - | "operator_triggered" - | "asset_triggered" - | "asset_materialization"; + "backfill" | "scheduled" | "manual" | "operator_triggered" | "asset_triggered" | "asset_materialization"; /** * All possible states that a DagRun can be in. * @@ -163,6 +159,13 @@ export type Conf = { [k: string]: unknown; } | null; export type TriggeringUserName = string | null; +export type Name7 = string; +export type Uri4 = string; +/** + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "JsonValue". + */ +export type JsonValue = unknown; export type SourceTaskId1 = string | null; export type SourceDagId1 = string | null; export type SourceRunId1 = string | null; @@ -226,14 +229,6 @@ 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 = @@ -244,6 +239,7 @@ export type NextKwargs1 = | null; export type XcomKeysToClear = string[]; export type ShouldRetry = boolean; +export type HasMappedDependants = boolean; export type StartDate2 = string | null; export type Type13 = "TaskCallbackRequest"; export type Filepath2 = string; @@ -294,11 +290,11 @@ 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 State2 = "deferred"; export type Classpath = string; export type TriggerKwargs = | { - [k: string]: JsonValue; + [k: string]: unknown; } | string | null; @@ -306,7 +302,7 @@ export type TriggerTimeout = string | null; export type Queue1 = string | null; export type NextMethod2 = string; export type NextKwargs2 = { - [k: string]: JsonValue; + [k: string]: unknown; } | null; export type RenderedMapIndex1 = string | null; export type Type20 = "DeferTask"; @@ -523,12 +519,12 @@ 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 State5 = "up_for_reschedule"; 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 State6 = "up_for_retry"; export type EndDate6 = string; export type RenderedMapIndex2 = string | null; export type RetryDelaySeconds = number | null; @@ -563,7 +559,7 @@ export type DagRelPath = string; export type StartDate6 = string; export type SentryIntegration = string; export type Type73 = "StartupDetails"; -export type State7 = "success" | null; +export type State7 = "success"; export type EndDate7 = string; export type TaskOutlets = AssetProfile[] | null; export type OutletEvents = @@ -672,13 +668,13 @@ export interface AssetResponse { export interface DagRunAssetReference { run_id: RunId; dag_id: DagId; - logical_date?: LogicalDate; + logical_date: LogicalDate; start_date: StartDate; - end_date?: EndDate; + end_date: EndDate; state: State; - data_interval_start?: DataIntervalStart; - data_interval_end?: DataIntervalEnd; - partition_key?: PartitionKey; + data_interval_start: DataIntervalStart; + data_interval_end: DataIntervalEnd; + partition_key: PartitionKey; } /** * Response to GetAssetEvent request. @@ -710,20 +706,6 @@ export interface AssetProfile { 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. * @@ -731,10 +713,10 @@ export interface Extra2 { * via the `definition` "AssetResult". */ export interface AssetResult { - name: Name4; - uri: Uri3; + name: Name3; + uri: Uri2; group: Group1; - extra?: Extra3; + extra?: Extra2; type?: Type2; } /** @@ -744,7 +726,7 @@ export interface AssetResult { * via the `definition` "AssetStateStoreResult". */ export interface AssetStateStoreResult { - value: JsonValue; + value: unknown; type?: Type3; } /** @@ -778,7 +760,7 @@ export interface AwaitInputTask { * via the `definition` "BundleInfo". */ export interface BundleInfo { - name: Name5; + name: Name4; version?: Version; version_data?: VersionData; } @@ -787,7 +769,7 @@ export interface BundleInfo { * via the `definition` "ClearAssetStateStoreByName". */ export interface ClearAssetStateStoreByName { - name: Name6; + name: Name5; type?: Type6; } /** @@ -795,7 +777,7 @@ export interface ClearAssetStateStoreByName { * via the `definition` "ClearAssetStateStoreByUri". */ export interface ClearAssetStateStoreByUri { - uri: Uri4; + uri: Uri3; type?: Type7; } /** @@ -807,18 +789,34 @@ export interface ClearTaskStateStore { type?: Type8; } /** + * Connection schema for responses with fields that are needed for Runtime. + * * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema - * via the `definition` "ConnectionResult". + * via the `definition` "ConnectionResponse". */ -export interface ConnectionResult { +export interface ConnectionResponse { conn_id: ConnId; conn_type: ConnType; - host?: Host; - schema?: Schema; - login?: Login; - password?: Password; - port?: Port; - extra?: Extra4; + host: Host; + schema: Schema; + login: Login; + password: Password; + port: Port; + extra: Extra3; +} +/** + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "ConnectionResult". + */ +export interface ConnectionResult { + conn_id: ConnId1; + conn_type: ConnType1; + host: Host1; + schema: Schema1; + login: Login1; + password: Password1; + port: Port1; + extra: Extra4; type?: Type9; } /** @@ -846,7 +844,7 @@ export interface CreateHITLDetailPayload { */ export interface HITLUser { id: Id1; - name: Name7; + name: Name6; } /** * Response containing count of Dag Runs matching certain filters. @@ -920,7 +918,7 @@ export interface DagRun { */ export interface AssetEventDagRunReference { asset: AssetReferenceAssetEventDagRun; - extra: Extra5; + extra: Extra6; source_task_id: SourceTaskId1; source_dag_id: SourceDagId1; source_run_id: SourceRunId1; @@ -929,9 +927,23 @@ export interface AssetEventDagRunReference { timestamp: Timestamp1; partition_key?: PartitionKey2; } +/** + * Schema for AssetModel used in AssetEventDagRunReference. + * + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "AssetReferenceAssetEventDagRun". + */ +export interface AssetReferenceAssetEventDagRun { + name: Name7; + uri: Uri4; + extra: Extra5; +} export interface Extra5 { [k: string]: JsonValue; } +export interface Extra6 { + [k: string]: JsonValue; +} /** * Schema for TaskInstance model with minimal required fields needed for Runtime. * @@ -1002,6 +1014,7 @@ export interface TIRunContext { next_kwargs?: NextKwargs1; xcom_keys_to_clear?: XcomKeysToClear; should_retry?: ShouldRetry; + has_mapped_dependants?: HasMappedDependants; start_date?: StartDate2; } /** @@ -1014,22 +1027,6 @@ export interface VariableResponse { key: Key; value: Value; } -/** - * Connection schema for responses with fields that are needed for Runtime. - * - * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema - * via the `definition` "ConnectionResponse". - */ -export interface ConnectionResponse { - conn_id: ConnId1; - conn_type: ConnType1; - host: Host1; - schema: Schema1; - login: Login1; - password: Password1; - port: Port1; - extra: Extra6; -} /** * Email notification request for task failures/retries. * @@ -1086,12 +1083,12 @@ export interface Data { export interface DagResult { dag_id: DagId4; is_paused: IsPaused; - bundle_name?: BundleName4; - bundle_version?: BundleVersion3; - relative_fileloc?: RelativeFileloc; - owners?: Owners; + bundle_name: BundleName4; + bundle_version: BundleVersion3; + relative_fileloc: RelativeFileloc; + owners: Owners; tags: Tags; - next_dagrun?: NextDagrun; + next_dagrun: NextDagrun; type?: Type17; } /** @@ -1101,19 +1098,19 @@ export interface DagResult { export interface DagRunResult { dag_id: DagId5; run_id: RunId4; - logical_date?: LogicalDate2; - data_interval_start?: DataIntervalStart2; - data_interval_end?: DataIntervalEnd2; + logical_date: LogicalDate2; + data_interval_start: DataIntervalStart2; + data_interval_end: DataIntervalEnd2; run_after: RunAfter1; - start_date?: StartDate3; - end_date?: EndDate2; + start_date: StartDate3; + end_date: EndDate2; clear_number?: ClearNumber1; run_type: DagRunType; state: DagRunState; conf?: Conf1; triggering_user_name?: TriggeringUserName1; consumed_asset_events: ConsumedAssetEvents1; - partition_key?: PartitionKey4; + partition_key: PartitionKey4; partition_date?: PartitionDate1; note?: Note1; team_name?: TeamName1; @@ -1788,7 +1785,7 @@ export interface TaskState { * via the `definition` "TaskStateStoreResult". */ export interface TaskStateStoreResult { - value: JsonValue; + value: unknown; type?: Type79; } /** @@ -1852,7 +1849,7 @@ export interface VariableKeysResult { */ export interface VariableResult { key: Key19; - value?: Value2; + value: Value2; type?: Type85; } /** @@ -1871,7 +1868,7 @@ export interface XComCountResponse { */ export interface XComResult { key: Key20; - value: JsonValue; + value: unknown; type?: Type87; } /** diff --git a/ts-sdk/src/sdk/task.ts b/ts-sdk/src/sdk/task.ts index a61638745d4f0..ddf1d71b75e02 100644 --- a/ts-sdk/src/sdk/task.ts +++ b/ts-sdk/src/sdk/task.ts @@ -33,6 +33,16 @@ export interface TaskContext { readonly tryNumber: number; /** -1 for non-mapped tasks, 0..N-1 for mapped instances. */ readonly mapIndex: number; + /** + * Whether this task's return value feeds a downstream `.expand()`. + * + * When `true`, returning an array from this handler causes the + * supervisor to record `mapped_length = value.length` on the + * return-value XCom automatically, so the scheduler can expand the + * mapped dependants. When `false`, no `mapped_length` is recorded + * (no behaviour change from today). + */ + readonly hasMappedDependants: boolean; /** * AbortSignal that fires when Airflow terminates the task subprocess * with SIGTERM or SIGINT. diff --git a/ts-sdk/tests/coordinator/client.test.ts b/ts-sdk/tests/coordinator/client.test.ts index b2583e80e0a03..ccf0896abba3d 100644 --- a/ts-sdk/tests/coordinator/client.test.ts +++ b/ts-sdk/tests/coordinator/client.test.ts @@ -35,6 +35,7 @@ const FAKE_CTX: TaskContext = { runId: "r", tryNumber: 1, mapIndex: -1, + hasMappedDependants: false, signal: new AbortController().signal, }; diff --git a/ts-sdk/tests/coordinator/integration.test.ts b/ts-sdk/tests/coordinator/integration.test.ts index a580159073ec2..d1dacd128d2f5 100644 --- a/ts-sdk/tests/coordinator/integration.test.ts +++ b/ts-sdk/tests/coordinator/integration.test.ts @@ -579,4 +579,39 @@ describe("coordinator runtime integration", () => { const setXComReqs = result.runtimeRequests.filter((r) => r.type === "SetXCom"); expect(setXComReqs).toHaveLength(0); }); + + it("surfaces hasMappedDependants=true on TaskContext when the server sets it", async () => { + let observedCtx: unknown = null; + registerTask({ dagId: "test_dag", taskId: "mapped_producer" }, async ({ ctx }) => { + observedCtx = ctx; + return ["a", "b", "c"]; + }); + + const responder: Responder = (msgType) => (msgType === "SetXCom" ? { body: null } : null); + + const result = await driveSupervisor( + makeStartupDetails("mapped_producer", "test_dag", "r1", { + has_mapped_dependants: true, + max_tries: 1, + }), + responder, + ); + + expect(result.firstResponse!.body).toMatchObject({ type: "SucceedTask" }); + expect(observedCtx).toMatchObject({ hasMappedDependants: true }); + }); + + it("defaults hasMappedDependants to false when the server omits it", async () => { + let observedCtx: unknown = null; + registerTask({ dagId: "test_dag", taskId: "plain_task" }, async ({ ctx }) => { + observedCtx = ctx; + }); + + const result = await driveSupervisor( + makeStartupDetails("plain_task", "test_dag", "r1", { max_tries: 1 }), + ); + + expect(result.firstResponse!.body).toMatchObject({ type: "SucceedTask" }); + expect(observedCtx).toMatchObject({ hasMappedDependants: false }); + }); }); diff --git a/ts-sdk/tests/public-api.test.ts b/ts-sdk/tests/public-api.test.ts index 751b749539722..62cd1c3dd0ed6 100644 --- a/ts-sdk/tests/public-api.test.ts +++ b/ts-sdk/tests/public-api.test.ts @@ -72,6 +72,7 @@ describe("public API", () => { readonly runId: string; readonly tryNumber: number; readonly mapIndex: number; + readonly hasMappedDependants: boolean; readonly signal: AbortSignal; }>(); expectTypeOf().toEqualTypeOf<{