apache-airflow-task-sdk 1.3.0b2__tar.gz → 1.3.0rc1__tar.gz

{apache_airflow_task_sdk-1.3.0b2 → apache_airflow_task_sdk-1.3.0rc1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: apache-airflow-task-sdk
-Version: 1.3.0b2
+Version: 1.3.0rc1
 Summary: Python Task SDK for Apache Airflow DAG Authors
 Project-URL: Bug Tracker, https://github.com/apache/airflow/issues
 Project-URL: Documentation, https://airflow.staged.apache.org/docs/
@@ -29,7 +29,7 @@ Classifier: Programming Language :: Python :: 3.13
 Classifier: Programming Language :: Python :: 3.14
 Classifier: Topic :: System :: Monitoring
 Requires-Python: !=3.15,>=3.10
-Requires-Dist: apache-airflow-core<3.4.0,>=3.3.0b1
+Requires-Dist: apache-airflow-core<3.4.0,>=3.3.0rc1
 Requires-Dist: asgiref>=2.3.0; python_version < '3.14'
 Requires-Dist: asgiref>=3.11.1; python_version >= '3.14'
 Requires-Dist: attrs!=25.2.0,>=24.2.0

{apache_airflow_task_sdk-1.3.0b2 → apache_airflow_task_sdk-1.3.0rc1}/dev/generate_task_sdk_models.py

@@ -33,8 +33,6 @@ from datamodel_code_generator import (
 from openapi_spec_validator import validate_spec
 
 os.environ["_AIRFLOW__AS_LIBRARY"] = "1"
-# Set a placeholder secret to allow the in-process FastAPI app to run its lifecycle hooks.
-os.environ.setdefault("AIRFLOW__API_AUTH__JWT_SECRET", "task-sdk-model-generation")
 
 AIRFLOW_ROOT_PATH = Path(__file__).parents[2].resolve()
 AIRFLOW_TASK_SDK_ROOT_PATH = AIRFLOW_ROOT_PATH / "task-sdk"

{apache_airflow_task_sdk-1.3.0b2 → apache_airflow_task_sdk-1.3.0rc1}/docs/api.rst

@@ -214,7 +214,7 @@ Timetables
 
 .. autoapiclass:: airflow.sdk.MultipleCronTriggerTimetable
 
-.. autoapiclass:: airflow.sdk.PartitionAtRuntime
+.. autoapiclass:: airflow.sdk.PartitionedAtRuntime
 
 .. autoapiclass:: airflow.sdk.PartitionedAssetTimetable
 

{apache_airflow_task_sdk-1.3.0b2 → apache_airflow_task_sdk-1.3.0rc1}/pyproject.toml

@@ -48,7 +48,7 @@ classifiers = [
     "Topic :: System :: Monitoring",
 ]
 dependencies = [
-    "apache-airflow-core<3.4.0,>=3.3.0b1",
+    "apache-airflow-core<3.4.0,>=3.3.0rc1",
     "asgiref>=2.3.0; python_version < '3.14'",
     "asgiref>=3.11.1; python_version >= '3.14'",
     "attrs>=24.2.0, !=25.2.0",

{apache_airflow_task_sdk-1.3.0b2 → apache_airflow_task_sdk-1.3.0rc1}/src/airflow/sdk/__init__.py

@@ -67,7 +67,7 @@ __all__ = [
     "ObjectStoragePath",
     "Param",
     "ParamsDict",
-    "PartitionAtRuntime",
+    "PartitionedAtRuntime",
     "PartitionedAssetTimetable",
     "PartitionMapper",
     "PokeReturnValue",
@@ -117,7 +117,7 @@ __all__ = [
     "teardown",
 ]
 
-__version__ = "1.3.0b2"
+__version__ = "1.3.0rc1"
 
 if TYPE_CHECKING:
     from airflow.sdk.api.datamodels._generated import DagRunState, TaskInstanceState, TriggerRule, WeightRule
@@ -199,8 +199,8 @@ if TYPE_CHECKING:
     from airflow.sdk.definitions.template import literal
     from airflow.sdk.definitions.timetables.assets import (
         AssetOrTimeSchedule,
-        PartitionAtRuntime,
         PartitionedAssetTimetable,
+        PartitionedAtRuntime,
     )
     from airflow.sdk.definitions.timetables.events import EventsTimetable
     from airflow.sdk.definitions.timetables.interval import (
@@ -269,7 +269,7 @@ __lazy_imports: dict[str, str] = {
     "ObjectStoragePath": ".io.path",
     "Param": ".definitions.param",
     "ParamsDict": ".definitions.param",
-    "PartitionAtRuntime": ".definitions.timetables.assets",
+    "PartitionedAtRuntime": ".definitions.timetables.assets",
     "PartitionedAssetTimetable": ".definitions.timetables.assets",
     "PartitionMapper": ".definitions.partition_mappers.base",
     "PokeReturnValue": ".bases.sensor",

{apache_airflow_task_sdk-1.3.0b2 → apache_airflow_task_sdk-1.3.0rc1}/src/airflow/sdk/__init__.pyi

@@ -110,8 +110,8 @@ from airflow.sdk.definitions.taskgroup import TaskGroup as TaskGroup
 from airflow.sdk.definitions.template import literal as literal
 from airflow.sdk.definitions.timetables.assets import (
     AssetOrTimeSchedule,
-    PartitionAtRuntime,
     PartitionedAssetTimetable,
+    PartitionedAtRuntime,
 )
 from airflow.sdk.definitions.timetables.events import EventsTimetable
 from airflow.sdk.definitions.timetables.interval import (
@@ -177,7 +177,7 @@ __all__ = [
     "ObjectStoragePath",
     "Param",
     "PokeReturnValue",
-    "PartitionAtRuntime",
+    "PartitionedAtRuntime",
     "PartitionedAssetTimetable",
     "PartitionMapper",
     "ProductMapper",

{apache_airflow_task_sdk-1.3.0b2 → apache_airflow_task_sdk-1.3.0rc1}/src/airflow/sdk/api/datamodels/_generated.py

@@ -725,6 +725,7 @@ class AssetEventDagRunReference(BaseModel):
     source_map_index: Annotated[int | None, Field(title="Source Map Index")] = None
     source_aliases: Annotated[list[AssetAliasReferenceAssetEventDagRun], Field(title="Source Aliases")]
     timestamp: Annotated[AwareDatetime, Field(title="Timestamp")]
+    partition_key: Annotated[str | None, Field(title="Partition Key")] = None
 
 
 class AssetEventResponse(BaseModel):

{apache_airflow_task_sdk-1.3.0b2 → apache_airflow_task_sdk-1.3.0rc1}/src/airflow/sdk/bases/timetable.py

@@ -51,7 +51,7 @@ class BaseTimetable:
     """
     Whether this timetable defers partition selection to task runtime.
 
-    *True* for :class:`~airflow.sdk.PartitionAtRuntime`; downstream code can
+    *True* for :class:`~airflow.sdk.PartitionedAtRuntime`; downstream code can
     branch on this flag instead of using ``isinstance``.
     """
 

{apache_airflow_task_sdk-1.3.0b2 → apache_airflow_task_sdk-1.3.0rc1}/src/airflow/sdk/coordinators/_subprocess.py

@@ -72,16 +72,26 @@ def _socket_address(value: tuple | str) -> tuple[str, int] | None:
         return None
     host, port = value[:2]
     host = str(host)
-    # Canonicalize IPv4-mapped IPv6 ("::ffff:127.0.0.1" -> "127.0.0.1") so a dual-stack
-    # client (e.g. the JVM, shown v4-mapped in /proc/net/tcp6) matches the AF_INET
-    # supervisor socket's plain-IPv4 address in the ownership check below.
+    # Canonicalize an IPv4 address that a dual-stack client embeds in IPv6 so it matches
+    # the AF_INET supervisor socket's plain-IPv4 address in the ownership check below. A
+    # dual-stack JVM's loopback connection is rendered in two different forms depending on
+    # the platform, and both must collapse to plain "127.0.0.1":
+    #   * IPv4-mapped     "::ffff:127.0.0.1" -> "127.0.0.1"  (Linux, via /proc/net/tcp6)
+    #   * IPv4-compatible "::127.0.0.1"      -> "127.0.0.1"  (macOS, via psutil)
+    # Otherwise the JVM's connection fails the check and every Java task is rejected with
+    # "process exited with 1 before connecting".
     try:
         parsed = ipaddress.ip_address(host)
     except ValueError:
         pass
     else:
-        if isinstance(parsed, ipaddress.IPv6Address) and parsed.ipv4_mapped is not None:
-            host = str(parsed.ipv4_mapped)
+        if isinstance(parsed, ipaddress.IPv6Address):
+            if parsed.ipv4_mapped is not None:
+                host = str(parsed.ipv4_mapped)
+            elif 1 < int(parsed) <= 0xFFFFFFFF:
+                # IPv4-compatible IPv6: ::/96 with the IPv4 in the low 32 bits. Exclude
+                # "::" (unspecified) and "::1" (IPv6 loopback), which are not IPv4.
+                host = str(ipaddress.IPv4Address(int(parsed)))
     return host, int(port)
 
 

{apache_airflow_task_sdk-1.3.0b2 → apache_airflow_task_sdk-1.3.0rc1}/src/airflow/sdk/coordinators/executable/coordinator.py

@@ -370,12 +370,11 @@ class ExecutableCoordinator(SubprocessCoordinator):
     Configuration is taken from the ``[sdk] coordinators`` entry that constructs
     this instance::
 
-        {
-            "name": "go",
+        "go": {
             "classpath": "airflow.sdk.coordinators.executable.ExecutableCoordinator",
             "kwargs": {
-                "executables_root": ["~/airflow/executable-bundles"],
-            },
+                "executables_root": ["~/airflow/executable-bundles"]
+            }
         }
 
     :param executables_root: A list of directories scanned for executable

{apache_airflow_task_sdk-1.3.0b2 → apache_airflow_task_sdk-1.3.0rc1}/src/airflow/sdk/coordinators/java/coordinator.py

@@ -174,14 +174,13 @@ class JavaCoordinator(SubprocessCoordinator):
     Configuration is taken from the ``[sdk] coordinators`` entry that constructs
     this instance::
 
-        {
-            "name": "jdk-17",
+        "jdk-17": {
             "classpath": "airflow.sdk.coordinators.java.JavaCoordinator",
             "kwargs": {
-                "java_executable": "/usr/lib/jvm/java-17-openjdk/bin/java",
-                "jvm_args": ["-Xmx1024m"],
                 "jars_root": ["~/airflow/jars"],
-            },
+                "java_executable": "/usr/lib/jvm/java-17-openjdk/bin/java",
+                "jvm_args": ["-Xmx1024m"]
+            }
         }
 
     :param java_executable: Path to the ``java`` command (defaults to

{apache_airflow_task_sdk-1.3.0b2 → apache_airflow_task_sdk-1.3.0rc1}/src/airflow/sdk/definitions/dag.py

@@ -1404,8 +1404,23 @@ class DAG:
                 # triggerer may mark tasks scheduled so we read from DB
                 all_tis = set(dr.get_task_instances(session=session))
                 scheduled_tis = {x for x in all_tis if x.state == TaskInstanceState.SCHEDULED}
-                ids_unrunnable = {x for x in all_tis if x.state not in FINISHED_STATES} - scheduled_tis
-                if not scheduled_tis and ids_unrunnable:
+                awaiting_input_tis = {x for x in all_tis if x.state == TaskInstanceState.AWAITING_INPUT}
+                ids_unrunnable = (
+                    {x for x in all_tis if x.state not in FINISHED_STATES}
+                    - scheduled_tis
+                    - awaiting_input_tis
+                )
+                if not scheduled_tis and awaiting_input_tis:
+                    # Human-in-the-loop tasks stay parked in AWAITING_INPUT: dag.test() never
+                    # resolves them itself. Keep the run alive until a response recorded from
+                    # outside -- the Required Actions UI or the HITL REST API of an api-server
+                    # sharing this metadata DB -- flips them back to SCHEDULED.
+                    log.info(
+                        "Waiting for Human-in-the-loop input for tasks: %s",
+                        sorted(x.task_id for x in awaiting_input_tis),
+                    )
+                    time.sleep(1)
+                elif not scheduled_tis and ids_unrunnable:
                     log.warning("No tasks to run. unrunnable tasks: %s", ids_unrunnable)
                     time.sleep(1)
 

{apache_airflow_task_sdk-1.3.0b2 → apache_airflow_task_sdk-1.3.0rc1}/src/airflow/sdk/definitions/timetables/assets.py

@@ -54,7 +54,7 @@ class PartitionedAssetTimetable(AssetTriggeredTimetable):
     default_partition_mapper: PartitionMapper = IdentityMapper()
 
 
-class PartitionAtRuntime(BaseTimetable):
+class PartitionedAtRuntime(BaseTimetable):
     """Marker timetable indicating that partition key(s) are determined at runtime."""
 
     can_be_scheduled = False

{apache_airflow_task_sdk-1.3.0b2 → apache_airflow_task_sdk-1.3.0rc1}/src/airflow/sdk/execution_time/request_handlers.py

@@ -31,6 +31,7 @@ from typing import TYPE_CHECKING
 from uuid import UUID
 
 from airflow.sdk.api.datamodels._generated import (
+    AssetStateStoreResponse,
     ConnectionResponse,
     DagRunStateResponse,
     TaskStatesResponse,
@@ -40,10 +41,17 @@ from airflow.sdk.api.datamodels._generated import (
     XComSequenceSliceResponse,
 )
 from airflow.sdk.execution_time.comms import (
+    AssetStateStoreResult,
+    ClearAssetStateStoreByName,
+    ClearAssetStateStoreByUri,
     ConnectionResult,
     DagRunStateResult,
+    DeleteAssetStateStoreByName,
+    DeleteAssetStateStoreByUri,
     DeleteVariable,
     DeleteXCom,
+    GetAssetStateStoreByName,
+    GetAssetStateStoreByUri,
     GetConnection,
     GetDagRunState,
     GetDRCount,
@@ -60,6 +68,8 @@ from airflow.sdk.execution_time.comms import (
     MaskSecret,
     PrevSuccessfulDagRunResult,
     PutVariable,
+    SetAssetStateStoreByName,
+    SetAssetStateStoreByUri,
     SetXCom,
     TaskStatesResult,
     VariableKeysResult,
@@ -272,3 +282,87 @@ def handle_get_xcom(client: Client, msg: GetXCom) -> tuple[BaseModel | None, dic
         xcom_result = XComResult.from_xcom_response(xcom)
         return xcom_result, {"exclude_unset": True}
     return xcom, {}
+
+
+def handle_get_asset_state_store_by_name(
+    client: Client, msg: GetAssetStateStoreByName
+) -> tuple[BaseModel | None, dict[str, bool]]:
+    asset_state = client.asset_state_store.get(
+        key=msg.key,
+        name=msg.name,
+    )
+
+    if isinstance(asset_state, AssetStateStoreResponse):
+        return AssetStateStoreResult.from_asset_state_store_response(asset_state), {}
+
+    return asset_state, {}
+
+
+def handle_get_asset_state_store_by_uri(
+    client: Client, msg: GetAssetStateStoreByUri
+) -> tuple[BaseModel | None, dict[str, bool]]:
+    asset_state = client.asset_state_store.get(
+        key=msg.key,
+        uri=msg.uri,
+    )
+
+    if isinstance(asset_state, AssetStateStoreResponse):
+        return AssetStateStoreResult.from_asset_state_store_response(asset_state), {}
+
+    return asset_state, {}
+
+
+def handle_set_asset_state_store_by_name(
+    client: Client, msg: SetAssetStateStoreByName
+) -> tuple[BaseModel | None, dict[str, bool]]:
+    client.asset_state_store.set(
+        key=msg.key,
+        value=msg.value,
+        name=msg.name,
+    )
+    return None, {}
+
+
+def handle_set_asset_state_store_by_uri(
+    client: Client, msg: SetAssetStateStoreByUri
+) -> tuple[BaseModel | None, dict[str, bool]]:
+    client.asset_state_store.set(
+        key=msg.key,
+        value=msg.value,
+        uri=msg.uri,
+    )
+    return None, {}
+
+
+def handle_delete_asset_state_store_by_name(
+    client: Client, msg: DeleteAssetStateStoreByName
+) -> tuple[BaseModel | None, dict[str, bool]]:
+    client.asset_state_store.delete(
+        key=msg.key,
+        name=msg.name,
+    )
+    return None, {}
+
+
+def handle_delete_asset_state_store_by_uri(
+    client: Client, msg: DeleteAssetStateStoreByUri
+) -> tuple[BaseModel | None, dict[str, bool]]:
+    client.asset_state_store.delete(
+        key=msg.key,
+        uri=msg.uri,
+    )
+    return None, {}
+
+
+def handle_clear_asset_state_store_by_name(
+    client: Client, msg: ClearAssetStateStoreByName
+) -> tuple[BaseModel | None, dict[str, bool]]:
+    client.asset_state_store.clear(name=msg.name)
+    return None, {}
+
+
+def handle_clear_asset_state_store_by_uri(
+    client: Client, msg: ClearAssetStateStoreByUri
+) -> tuple[BaseModel | None, dict[str, bool]]:
+    client.asset_state_store.clear(uri=msg.uri)
+    return None, {}

{apache_airflow_task_sdk-1.3.0b2 → apache_airflow_task_sdk-1.3.0rc1}/src/airflow/sdk/execution_time/schema/__init__.py

@@ -52,7 +52,6 @@ from airflow.sdk.execution_time.schema.migrator import (
     SchemaVersionMigrator,
     get_schema_version_migrator,
 )
-from airflow.sdk.execution_time.schema.versions import bundle
 
 if TYPE_CHECKING:
     from collections.abc import Iterable
@@ -119,6 +118,18 @@ def resolve_body_class(body: Any) -> type[BaseModel] | None:
     return registered_models_by_name().get(name)
 
 
+def __getattr__(name: str) -> Any:
+    # Re-export ``bundle`` lazily so importing this package does not import ``cadwyn`` (-> FastAPI/
+    # Starlette) until something actually accesses ``schema.bundle``. The Task SDK supervisor imports
+    # this package on every worker, but only the foreign-language-SDK migration path touches the
+    # bundle, so a pure-Python worker never pays the cadwyn import.
+    if name == "bundle":
+        from airflow.sdk.execution_time.schema.versions import get_bundle
+
+        return get_bundle()
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+
 __all__ = [
     "SchemaVersionMigrator",
     "bundle",

{apache_airflow_task_sdk-1.3.0b2 → apache_airflow_task_sdk-1.3.0rc1}/src/airflow/sdk/execution_time/schema/migrator.py

@@ -33,7 +33,6 @@ import functools
 from typing import TYPE_CHECKING, Any, cast
 
 import attrs
-from cadwyn import generate_versioned_models
 
 if TYPE_CHECKING:
     from cadwyn import VersionBundle
@@ -103,6 +102,11 @@ class SchemaVersionMigrator:
 
     def _versioned_class(self, version: str, model: type[BaseModel]) -> type[BaseModel]:
         """Get the Cadwyn-generated class for *model* at *version*."""
+        # Imported here, not at module scope, so importing this module (and the parent ``schema``
+        # package the Task SDK supervisor pulls in) does not import ``cadwyn`` -> FastAPI/Starlette.
+        # Only the foreign-language-SDK migration path reaches this; a pure-Python worker never does.
+        from cadwyn import generate_versioned_models
+
         if self._versioned_models is None:
             self._versioned_models = generate_versioned_models(self._bundle)
         return self._versioned_models[version][model]
@@ -207,8 +211,9 @@ def get_schema_version_migrator() -> SchemaVersionMigrator:
     no per-call state, so concurrent callers can share a single
     instance safely.
     """
-    from airflow.sdk.execution_time.schema.versions import bundle
+    from airflow.sdk.execution_time.schema.versions import get_bundle
 
+    bundle = get_bundle()
     return SchemaVersionMigrator(bundle=bundle, supervisor_version=bundle.versions[0].value)
 
 

{apache_airflow_task_sdk-1.3.0b2 → apache_airflow_task_sdk-1.3.0rc1}/src/airflow/sdk/execution_time/schema/schema.json

@@ -4502,6 +4502,18 @@
           "format": "date-time",
           "title": "Timestamp",
           "type": "string"
+        },
+        "partition_key": {
+          "anyOf": [
+            {
+              "type": "string"
+            },
+            {
+              "type": "null"
+            }
+          ],
+          "default": null,
+          "title": "Partition Key"
         }
       },
       "required": [

apache_airflow_task_sdk-1.3.0rc1/src/airflow/sdk/execution_time/schema/versions/__init__.py

@@ -0,0 +1,51 @@
+# 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.
+
+from __future__ import annotations
+
+import functools
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from cadwyn import VersionBundle
+
+
+@functools.cache
+def get_bundle() -> VersionBundle:
+    """
+    Build the supervisor schema ``VersionBundle`` lazily.
+
+    The ``cadwyn`` import is deferred to here so importing this module -- and the parent ``schema``
+    package, which the Task SDK supervisor pulls in on every worker -- does not import ``cadwyn``,
+    which drags in FastAPI/Starlette/Jinja2. The bundle (and its cadwyn machinery) is only needed on
+    the foreign-language-SDK migration path; a pure-Python worker never builds it. Cached so the
+    bundle is constructed once per process.
+    """
+    from cadwyn import HeadVersion, Version, VersionBundle
+
+    return VersionBundle(
+        HeadVersion(),
+        Version("2026-06-16"),
+    )
+
+
+def __getattr__(name: str) -> Any:
+    # Keep ``from ...versions import bundle`` working for existing callers without forcing the
+    # cadwyn import at module load (only resolves the bundle when ``bundle`` is actually accessed).
+    if name == "bundle":
+        return get_bundle()
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

{apache_airflow_task_sdk-1.3.0b2 → apache_airflow_task_sdk-1.3.0rc1}/tests/task_sdk/coordinators/test_subprocess.py

@@ -372,6 +372,41 @@ class TestConnectionFromProcess:
             client.close()
             server.close()
 
+    def test_matches_dual_stack_ipv4_compatible_connection(self):
+        """A dual-stack child whose loopback is rendered in IPv4-compatible form is accepted.
+
+        Companion to :meth:`test_matches_dual_stack_ipv4_mapped_connection` for macOS
+        (#68938): there the JVM's loopback connection is reported by ``psutil`` as the
+        deprecated IPv4-compatible ``::127.0.0.1`` rather than the IPv4-mapped
+        ``::ffff:127.0.0.1`` seen on Linux. Both forms must canonicalize to plain
+        ``127.0.0.1`` or the ownership check rejects the Java task. The OS will not
+        reliably establish a routable ``::`` connection on demand, so ``psutil``'s view of
+        the child's connections is mocked to the form macOS actually reports.
+        """
+        server = _start_server()
+        _, server_port = server.getsockname()
+        client = socket.socket()
+        client.connect(("127.0.0.1", server_port))
+        conn, _ = server.accept()
+        child_port = conn.getpeername()[1]
+        mock_proc = MagicMock(spec=subprocess.Popen)
+        mock_proc.pid = os.getpid()
+
+        # On macOS psutil reports the child's dual-stack loopback in IPv4-compatible form.
+        compat_conn = MagicMock(
+            laddr=("::127.0.0.1", child_port),
+            raddr=("::127.0.0.1", server_port),
+        )
+        try:
+            with patch("airflow.sdk.coordinators._subprocess.psutil.Process") as mock_process:
+                mock_process.return_value.children.return_value = []
+                mock_process.return_value.net_connections.return_value = [compat_conn]
+                assert _is_connection_from_process(conn, mock_proc) is True
+        finally:
+            conn.close()
+            client.close()
+            server.close()
+
     def test_rejects_tcp_connection_not_owned_by_child_process(self):
         server = _start_server()
         _, port = server.getsockname()

{apache_airflow_task_sdk-1.3.0b2 → apache_airflow_task_sdk-1.3.0rc1}/tests/task_sdk/definitions/test_dag.py

@@ -25,7 +25,16 @@ from unittest import mock
 
 import pytest
 
-from airflow.sdk import DAG, Context, Label, Param, PartitionAtRuntime, TaskGroup, dag as dag_decorator, task
+from airflow.sdk import (
+    DAG,
+    Context,
+    Label,
+    Param,
+    PartitionedAtRuntime,
+    TaskGroup,
+    dag as dag_decorator,
+    task,
+)
 from airflow.sdk.bases.operator import BaseOperator
 from airflow.sdk.bases.timetable import BaseTimetable
 from airflow.sdk.definitions.param import DagParam, ParamsDict
@@ -438,8 +447,8 @@ class TestDag:
         with pytest.raises(ValueError, match="ContinuousTimetable requires max_active_runs <= 1"):
             dag = DAG("continuous", start_date=DEFAULT_DATE, schedule="@continuous", max_active_runs=25)
 
-    def test_only_partition_at_runtime_has_partitioned_at_runtime_flag(self):
-        """Regression guard: across every BaseTimetable subclass, only PartitionAtRuntime sets partitioned_at_runtime=True."""
+    def test_only_partitioned_at_runtime_has_partitioned_at_runtime_flag(self):
+        """Regression guard: across every BaseTimetable subclass, only PartitionedAtRuntime sets partitioned_at_runtime=True."""
 
         def all_subclasses(cls):
             for sub in cls.__subclasses__():
@@ -447,7 +456,7 @@ class TestDag:
                 yield from all_subclasses(sub)
 
         flagged = {c for c in all_subclasses(BaseTimetable) if c.partitioned_at_runtime}
-        assert flagged == {PartitionAtRuntime}
+        assert flagged == {PartitionedAtRuntime}
 
     def test_dag_add_task_checks_trigger_rule(self):
         # A non fail stop dag should allow any trigger rule

{apache_airflow_task_sdk-1.3.0b2 → apache_airflow_task_sdk-1.3.0rc1}/tests/task_sdk/execution_time/schema/test_migrator.py

@@ -28,6 +28,8 @@ schema diverges from head.
 
 from __future__ import annotations
 
+import subprocess
+import sys
 from typing import Literal
 
 import pytest
@@ -334,3 +336,36 @@ class TestResolveBodyClass:
     def test_extra_fields_in_body_do_not_affect_resolution(self, mock_registry):
         body = {"type": "_LangSdkRequest", "ti_id": "t1", "field_a": 7}
         assert resolve_body_class(body) is _LangSdkRequest
+
+
+class TestLazyCadwynImport:
+    """``cadwyn`` (which imports FastAPI/Starlette/Jinja2) must stay off the worker import path.
+
+    The Task SDK supervisor imports the schema package on every Celery pool worker, but ``cadwyn`` is
+    only needed on the foreign-language-SDK migration path, so its import is deferred. Each check runs
+    in a fresh interpreter because ``sys.modules`` is process-global and other tests import cadwyn.
+    """
+
+    @pytest.mark.parametrize(
+        "module",
+        [
+            "airflow.sdk.execution_time.schema",
+            "airflow.sdk.execution_time.supervisor",
+        ],
+    )
+    def test_importing_worker_path_does_not_load_cadwyn(self, module):
+        code = (
+            f"import sys; import {module}; "
+            "assert 'cadwyn' not in sys.modules, sorted(m for m in sys.modules if 'cadwyn' in m); "
+            "assert 'fastapi' not in sys.modules, 'fastapi was imported'"
+        )
+        subprocess.run([sys.executable, "-c", code], check=True, capture_output=True, text=True)
+
+    def test_accessing_bundle_loads_cadwyn(self):
+        # The foreign-SDK migration path does need cadwyn; accessing the bundle is where it loads.
+        code = (
+            "import sys; from airflow.sdk.execution_time.schema import bundle; "
+            "assert bundle.versions, 'bundle should be built'; "
+            "assert 'cadwyn' in sys.modules, 'cadwyn should load when the bundle is accessed'"
+        )
+        subprocess.run([sys.executable, "-c", code], check=True, capture_output=True, text=True)

{apache_airflow_task_sdk-1.3.0b2 → apache_airflow_task_sdk-1.3.0rc1}/tests/task_sdk/execution_time/test_context.py

@@ -660,6 +660,24 @@ class TestTriggeringAssetEventsAccessor:
         assert mock_supervisor_comms.send.mock_calls == [mock.call(GetAssetByUri(uri=uri))]
         assert _AssetRefResolutionMixin._asset_ref_cache
 
+    def test_partition_key_exposed(self):
+        """A consumed asset event's partition key is reachable via triggering_asset_events."""
+        event = {
+            "asset": {"name": "1", "uri": "1", "extra": {}},
+            "extra": {},
+            "source_task_id": "t1",
+            "source_dag_id": "d1",
+            "source_run_id": "r1",
+            "source_map_index": -1,
+            "source_aliases": [],
+            "timestamp": "2025-01-01T00:00:12Z",
+            "partition_key": "2024-01-15",
+        }
+        accessor = TriggeringAssetEventsAccessor.build(
+            [AssetEventDagRunReferenceResult.model_validate(event)]
+        )
+        assert [e.partition_key for e in accessor[Asset("1")]] == ["2024-01-15"]
+
     def test_source_task_instance_xcom_pull(self, mock_supervisor_comms, accessor):
         events = accessor[Asset("2")]
         assert len(events) == 1