diff --git a/.dockerignore b/.dockerignore
index dce962375a92f..40954c20ec5f2 100644
--- a/.dockerignore
+++ b/.dockerignore
@@ -35,7 +35,7 @@
!docs
!licenses
!providers/
-!task_sdk/
+!task-sdk/
# Add those folders to the context so that they are available in the CI container
!scripts
diff --git a/.github/boring-cyborg.yml b/.github/boring-cyborg.yml
index 3c7e223209c8f..964e45a00c090 100644
--- a/.github/boring-cyborg.yml
+++ b/.github/boring-cyborg.yml
@@ -414,7 +414,7 @@ labelPRBasedOnFilePath:
- tests/system/**/*
area:task-sdk:
- - task_sdk/**/*
+ - task-sdk/**/*
area:db-migrations:
- airflow/migrations/versions/*
diff --git a/.github/dependabot.yml b/.github/dependabot.yml
index b243de7eda35f..3806c65061895 100644
--- a/.github/dependabot.yml
+++ b/.github/dependabot.yml
@@ -22,7 +22,7 @@ updates:
- /clients/python
- /dev/breeze
- /docker_tests
- - /task_sdk
+ - /task-sdk
- /
schedule:
interval: daily
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index baac8947ae294..6fd607e835ca5 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -676,7 +676,7 @@ repos:
name: Sync template context variable refs
language: python
entry: ./scripts/ci/pre_commit/template_context_key_sync.py
- files: ^airflow/models/taskinstance\.py$|^task_sdk/src/airflow/sdk/definitions/context\.py$|^docs/apache-airflow/templates-ref\.rst$
+ files: ^airflow/models/taskinstance\.py$|^task-sdk/src/airflow/sdk/definitions/context\.py$|^docs/apache-airflow/templates-ref\.rst$
- id: check-base-operator-usage
language: pygrep
name: Check BaseOperator core imports
@@ -1174,9 +1174,9 @@ repos:
^tests/ |
^providers/tests/ |
^providers/.*/tests/ |
- ^task_sdk/src/airflow/sdk/definitions/dag.py$ |
- ^task_sdk/src/airflow/sdk/execution_time/secrets_masker.py$ |
- ^task_sdk/src/airflow/sdk/definitions/_internal/node.py$ |
+ ^task-sdk/src/airflow/sdk/definitions/dag.py$ |
+ ^task-sdk/src/airflow/sdk/execution_time/secrets_masker.py$ |
+ ^task-sdk/src/airflow/sdk/definitions/_internal/node.py$ |
^dev/.*\.py$ |
^scripts/.*\.py$ |
^docker_tests/.*$ |
@@ -1293,7 +1293,7 @@ repos:
provider_packages|
providers/|
scripts|
- task_sdk/|
+ task-sdk/|
tests/dags/test_imports\.py
)
require_serial: true
@@ -1344,14 +1344,14 @@ repos:
name: Run mypy for Task SDK
language: python
entry: ./scripts/ci/pre_commit/mypy.py --namespace-packages
- files: ^task_sdk/src/airflow/sdk/.*\.py$|^task_sdk/tests//.*\.py$
+ files: ^task-sdk/src/airflow/sdk/.*\.py$|^task-sdk/tests//.*\.py$
require_serial: true
additional_dependencies: ['rich>=12.4.4']
- id: mypy-task-sdk
stages: ['manual']
name: Run mypy for Task SDK (manual)
language: python
- entry: ./scripts/ci/pre_commit/mypy_folder.py task_sdk/src/airflow/sdk
+ entry: ./scripts/ci/pre_commit/mypy_folder.py task-sdk/src/airflow/sdk
pass_filenames: false
files: ^.*\.py$
require_serial: true
@@ -1398,7 +1398,7 @@ repos:
- id: generate-tasksdk-datamodels
name: Generate Datamodels for TaskSDK client
language: python
- entry: uv run --active --group codegen --project apache-airflow-task-sdk --directory task_sdk -s dev/generate_models.py
+ entry: uv run --active --group codegen --project apache-airflow-task-sdk --directory task-sdk -s dev/generate_models.py
pass_filenames: false
files: ^airflow/api_fastapi/execution_api/.*\.py$
require_serial: true
diff --git a/Dockerfile b/Dockerfile
index f9f3b16753e4f..7cdfaec8aa211 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -815,7 +815,7 @@ function install_airflow() {
local installation_command_flags
if [[ ${AIRFLOW_INSTALLATION_METHOD} == "." ]]; then
# When installing from sources - we always use `--editable` mode
- installation_command_flags="--editable .[${AIRFLOW_EXTRAS}]${AIRFLOW_VERSION_SPECIFICATION} --editable ./task_sdk --editable ./devel-common"
+ installation_command_flags="--editable .[${AIRFLOW_EXTRAS}]${AIRFLOW_VERSION_SPECIFICATION} --editable ./task-sdk --editable ./devel-common"
while IFS= read -r -d '' pyproject_toml_file; do
project_folder=$(dirname ${pyproject_toml_file})
installation_command_flags="${installation_command_flags} --editable ${project_folder}"
diff --git a/Dockerfile.ci b/Dockerfile.ci
index f632d40dbd390..dc63936aa757d 100644
--- a/Dockerfile.ci
+++ b/Dockerfile.ci
@@ -585,7 +585,7 @@ function install_airflow() {
local installation_command_flags
if [[ ${AIRFLOW_INSTALLATION_METHOD} == "." ]]; then
# When installing from sources - we always use `--editable` mode
- installation_command_flags="--editable .[${AIRFLOW_EXTRAS}]${AIRFLOW_VERSION_SPECIFICATION} --editable ./task_sdk --editable ./devel-common"
+ installation_command_flags="--editable .[${AIRFLOW_EXTRAS}]${AIRFLOW_VERSION_SPECIFICATION} --editable ./task-sdk --editable ./devel-common"
while IFS= read -r -d '' pyproject_toml_file; do
project_folder=$(dirname ${pyproject_toml_file})
installation_command_flags="${installation_command_flags} --editable ${project_folder}"
@@ -1001,7 +1001,7 @@ function check_force_lowest_dependencies() {
echo
fi
set -x
- uv pip install --python "$(which python)" --resolution lowest-direct --upgrade --editable ".${EXTRA}" --editable "./task_sdk" --editable "./devel-common"
+ uv pip install --python "$(which python)" --resolution lowest-direct --upgrade --editable ".${EXTRA}" --editable "./task-sdk" --editable "./devel-common"
set +x
}
diff --git a/airflow/models/taskinstance.py b/airflow/models/taskinstance.py
index d2f0d46e0c45d..fbf3626f4916c 100644
--- a/airflow/models/taskinstance.py
+++ b/airflow/models/taskinstance.py
@@ -968,7 +968,7 @@ def get_triggering_events() -> dict[str, list[AssetEvent]]:
return triggering_events
# NOTE: If you add to this dict, make sure to also update the following:
- # * Context in task_sdk/src/airflow/sdk/definitions/context.py
+ # * Context in task-sdk/src/airflow/sdk/definitions/context.py
# * KNOWN_CONTEXT_KEYS in airflow/utils/context.py
# * Table in docs/apache-airflow/templates-ref.rst
diff --git a/airflow/utils/context.py b/airflow/utils/context.py
index 8ed44899bcb7f..507211492ebf1 100644
--- a/airflow/utils/context.py
+++ b/airflow/utils/context.py
@@ -47,7 +47,7 @@
from airflow.sdk.types import OutletEventAccessorsProtocol
# NOTE: Please keep this in sync with the following:
-# * Context in task_sdk/src/airflow/sdk/definitions/context.py
+# * Context in task-sdk/src/airflow/sdk/definitions/context.py
# * Table in docs/apache-airflow/templates-ref.rst
KNOWN_CONTEXT_KEYS: set[str] = {
"conn",
diff --git a/contributing-docs/07_local_virtualenv.rst b/contributing-docs/07_local_virtualenv.rst
index 177f2d433216a..18485dd57ab57 100644
--- a/contributing-docs/07_local_virtualenv.rst
+++ b/contributing-docs/07_local_virtualenv.rst
@@ -227,7 +227,7 @@ dependencies. For example, to install Amazon provider you need to install ``amaz
.. code:: bash
- pip install -e "./task_sdk"
+ pip install -e "./task-sdk"
pip install -e "./devel-common"
pip install -e "./providers/amazon"
pip install -e ".[amazon]"
@@ -284,7 +284,7 @@ You can run the following command in the venv that you have installed airflow in
.. code:: bash
pip install -e ".[google]"
- pip install -e "./task_sdk"
+ pip install -e "./task-sdk"
pip install -e "./devel-common"
pip install -e "./providers/google"
@@ -294,7 +294,7 @@ dependency groups will be implemented in ``pip`` - April 2025) - it will not be
when installing airflow - currently with ``pip`` it is the only way to install development dependencies
of the provider and is a bit convoluted.
-The second installs ``task_sdk`` project - where APIs for providers are kept.
+The second installs ``task-sdk`` project - where APIs for providers are kept.
The third one installs google provider source code in development mode, so that modifications
to the code are automatically reflected in your installed virtualenv.
diff --git a/contributing-docs/quick-start-ide/contributors_quick_start_pycharm.rst b/contributing-docs/quick-start-ide/contributors_quick_start_pycharm.rst
index 4343497329d98..cd40d8310a5b5 100644
--- a/contributing-docs/quick-start-ide/contributors_quick_start_pycharm.rst
+++ b/contributing-docs/quick-start-ide/contributors_quick_start_pycharm.rst
@@ -39,10 +39,10 @@ Setup your project
alt="Cloning github fork to Pycharm">
-3. Configure the source root directories well as for ``task_sdk`` and ``devel-common``.
+3. Configure the source root directories well as for ``task-sdk`` and ``devel-common``.
You also have to set "source" and "tests" root directories for each provider you want to develop (!).
- This is important in Airflow 3.0 we split ``task_sdk``, ``devel-common`` and each provider to be separate
+ This is important in Airflow 3.0 we split ``task-sdk``, ``devel-common`` and each provider to be separate
distribution - each with separate ``pyproject.toml`` file, so you need to separately
add "src" and "tests" directories for each provider you develop to be respectively
"source roots" and "test roots".
@@ -57,12 +57,12 @@ Setup your project
alt="Adding Source Root directories to Pycharm">
- You also need to add ``task_sdk`` sources (and ``devel-common`` in similar way).
+ You also need to add ``task-sdk`` sources (and ``devel-common`` in similar way).
.. raw:: html
-
diff --git a/dev/breeze/src/airflow_breeze/commands/release_management_commands.py b/dev/breeze/src/airflow_breeze/commands/release_management_commands.py
index 461c5592d012e..37b33a24e7168 100644
--- a/dev/breeze/src/airflow_breeze/commands/release_management_commands.py
+++ b/dev/breeze/src/airflow_breeze/commands/release_management_commands.py
@@ -627,7 +627,7 @@ def prepare_airflow_packages(
get_console().print("[success]Successfully prepared Airflow packages")
-TASK_SDK_DIR_PATH = AIRFLOW_SOURCES_ROOT / "task_sdk"
+TASK_SDK_DIR_PATH = AIRFLOW_SOURCES_ROOT / "task-sdk"
TASK_SDK_DIST_DIR_PATH = TASK_SDK_DIR_PATH / "dist"
@@ -692,7 +692,7 @@ def _build_package_with_docker(package_format: str):
"-e",
"GITHUB_ACTIONS",
"-w",
- "/opt/airflow/task_sdk",
+ "/opt/airflow/task-sdk",
AIRFLOW_BUILD_IMAGE_TAG,
"bash",
"-c",
@@ -707,7 +707,7 @@ def _build_package_with_docker(package_format: str):
DIST_DIR.mkdir(parents=True, exist_ok=True)
get_console().print()
# Copy all files in the dist directory in container to the host dist directory (note '/.' in SRC)
- run_command(["docker", "cp", f"{container_id}:/opt/airflow/task_sdk/dist/.", "./dist"], check=True)
+ run_command(["docker", "cp", f"{container_id}:/opt/airflow/task-sdk/dist/.", "./dist"], check=True)
run_command(["docker", "rm", "--force", container_id], check=False, stdout=DEVNULL, stderr=DEVNULL)
if use_local_hatch:
diff --git a/dev/breeze/src/airflow_breeze/commands/testing_commands.py b/dev/breeze/src/airflow_breeze/commands/testing_commands.py
index 83ebea2c96942..5939de725ed33 100644
--- a/dev/breeze/src/airflow_breeze/commands/testing_commands.py
+++ b/dev/breeze/src/airflow_breeze/commands/testing_commands.py
@@ -154,7 +154,7 @@ def docker_compose_tests(
sys.exit(return_code)
-TEST_PROGRESS_REGEXP = r"tests/.*|providers/.*/tests/.*|task_sdk/tests/.*|.*=====.*"
+TEST_PROGRESS_REGEXP = r"tests/.*|providers/.*/tests/.*|task-sdk/tests/.*|.*=====.*"
PERCENT_TEST_PROGRESS_REGEXP = r"^tests/.*\[[ \d%]*\].*|^\..*\[[ \d%]*\].*"
diff --git a/dev/breeze/src/airflow_breeze/global_constants.py b/dev/breeze/src/airflow_breeze/global_constants.py
index 9c64c567c413a..fbbd1b43bb354 100644
--- a/dev/breeze/src/airflow_breeze/global_constants.py
+++ b/dev/breeze/src/airflow_breeze/global_constants.py
@@ -287,7 +287,7 @@ def all_task_sdk_test_packages() -> list[str]:
return sorted(
[
candidate.name
- for candidate in (AIRFLOW_SOURCES_ROOT / "task_sdk" / "tests").iterdir()
+ for candidate in (AIRFLOW_SOURCES_ROOT / "task-sdk" / "tests").iterdir()
if candidate.is_dir() and candidate.name != "__pycache__"
]
)
diff --git a/dev/breeze/src/airflow_breeze/utils/docker_command_utils.py b/dev/breeze/src/airflow_breeze/utils/docker_command_utils.py
index 289466c7aaa8b..0c280698569b2 100644
--- a/dev/breeze/src/airflow_breeze/utils/docker_command_utils.py
+++ b/dev/breeze/src/airflow_breeze/utils/docker_command_utils.py
@@ -92,7 +92,7 @@
("hooks", "/opt/airflow/hooks"),
("logs", "/root/airflow/logs"),
("providers", "/opt/airflow/providers"),
- ("task_sdk", "/opt/airflow/task_sdk"),
+ ("task-sdk", "/opt/airflow/task-sdk"),
("pyproject.toml", "/opt/airflow/pyproject.toml"),
("scripts", "/opt/airflow/scripts"),
("scripts/docker/entrypoint_ci.sh", "/entrypoint"),
diff --git a/dev/breeze/src/airflow_breeze/utils/run_tests.py b/dev/breeze/src/airflow_breeze/utils/run_tests.py
index 174ad19975a07..814e1c91eacde 100644
--- a/dev/breeze/src/airflow_breeze/utils/run_tests.py
+++ b/dev/breeze/src/airflow_breeze/utils/run_tests.py
@@ -168,7 +168,7 @@ def get_excluded_provider_args(python_version: str) -> list[str]:
"Serialization": [
"tests/serialization",
],
- "TaskSDK": ["task_sdk/tests"],
+ "TaskSDK": ["task-sdk/tests"],
"WWW": [
"tests/www",
],
@@ -200,7 +200,7 @@ def get_excluded_provider_args(python_version: str) -> list[str]:
TEST_GROUP_TO_TEST_FOLDERS: dict[GroupOfTests, list[str]] = {
GroupOfTests.CORE: ["tests"],
GroupOfTests.PROVIDERS: ALL_PROVIDER_TEST_FOLDERS,
- GroupOfTests.TASK_SDK: ["task_sdk/tests"],
+ GroupOfTests.TASK_SDK: ["task-sdk/tests"],
GroupOfTests.HELM: ["helm_tests"],
GroupOfTests.INTEGRATION_CORE: ["tests/integration"],
GroupOfTests.INTEGRATION_PROVIDERS: ALL_PROVIDER_INTEGRATION_TEST_FOLDERS,
diff --git a/dev/breeze/src/airflow_breeze/utils/selective_checks.py b/dev/breeze/src/airflow_breeze/utils/selective_checks.py
index b4a0d2c4310b8..7d7be448088d3 100644
--- a/dev/breeze/src/airflow_breeze/utils/selective_checks.py
+++ b/dev/breeze/src/airflow_breeze/utils/selective_checks.py
@@ -190,7 +190,7 @@ def __hash__(self):
r"^chart",
r"^providers/.*/src/",
r"^providers/.*/docs/",
- r"^task_sdk/src/",
+ r"^task-sdk/src/",
r"^tests/system",
r"^CHANGELOG\.txt",
r"^airflow/config_templates/config\.yml",
@@ -226,8 +226,8 @@ def __hash__(self):
r"^providers/tests/",
r"^providers/.*/src/",
r"^providers/.*/tests/",
- r"^task_sdk/src/",
- r"^task_sdk/tests/",
+ r"^task-sdk/src/",
+ r"^task-sdk/tests/",
r"^tests",
r"^devel-common",
r"^kubernetes_tests",
@@ -246,18 +246,18 @@ def __hash__(self):
r"^devel-common/.*\.py$",
],
FileGroupForCi.TASK_SDK_FILES: [
- r"^task_sdk/src/airflow/sdk/.*\.py$",
- r"^task_sdk/tests/.*\.py$",
+ r"^task-sdk/src/airflow/sdk/.*\.py$",
+ r"^task-sdk/tests/.*\.py$",
],
FileGroupForCi.ASSET_FILES: [
r"^airflow/assets/",
r"^airflow/models/assets/",
- r"^task_sdk/src/airflow/sdk/definitions/asset/",
+ r"^task-sdk/src/airflow/sdk/definitions/asset/",
r"^airflow/datasets/",
],
FileGroupForCi.UNIT_TEST_FILES: [
r"^tests/",
- r"^task_sdk/tests/",
+ r"^task-sdk/tests/",
r"^providers/.*/tests/",
r"^dev/breeze/tests/",
],
@@ -274,8 +274,8 @@ def __hash__(self):
r"^docs/.*",
r"^providers/.*/tests/.*",
r"^tests/dags/test_imports.py",
- r"^task_sdk/src/airflow/sdk/.*\.py$",
- r"^task_sdk/tests/.*\.py$",
+ r"^task-sdk/src/airflow/sdk/.*\.py$",
+ r"^task-sdk/tests/.*\.py$",
]
}
)
@@ -307,8 +307,8 @@ def __hash__(self):
r"^providers/.*/tests/",
],
SelectiveTaskSdkTestType.TASK_SDK: [
- r"^task_sdk/src/",
- r"^task_sdk/tests/",
+ r"^task-sdk/src/",
+ r"^task-sdk/tests/",
],
SelectiveCoreTestType.SERIALIZATION: [
r"^airflow/serialization/",
diff --git a/dev/breeze/tests/test_pytest_args_for_test_types.py b/dev/breeze/tests/test_pytest_args_for_test_types.py
index 33db4c4974883..4500768a568ea 100644
--- a/dev/breeze/tests/test_pytest_args_for_test_types.py
+++ b/dev/breeze/tests/test_pytest_args_for_test_types.py
@@ -304,7 +304,7 @@ def test_pytest_args_for_missing_provider():
GroupOfTests.TASK_SDK,
"All",
[
- "task_sdk/tests",
+ "task-sdk/tests",
],
),
(
diff --git a/dev/breeze/tests/test_selective_checks.py b/dev/breeze/tests/test_selective_checks.py
index fd2afabf6ae31..250434548a9af 100644
--- a/dev/breeze/tests/test_selective_checks.py
+++ b/dev/breeze/tests/test_selective_checks.py
@@ -492,7 +492,7 @@ def assert_outputs_are_printed(expected_outputs: dict[str, str], stderr: str):
),
(
pytest.param(
- ("task_sdk/src/airflow/sdk/random.py",),
+ ("task-sdk/src/airflow/sdk/random.py",),
{
"all-python-versions": "['3.9']",
"all-python-versions-list-as-string": "3.9",
@@ -1688,7 +1688,7 @@ def test_expected_output_push(
(
"airflow/assets/",
"airflow/models/assets/",
- "task_sdk/src/airflow/sdk/definitions/asset/",
+ "task-sdk/src/airflow/sdk/definitions/asset/",
"airflow/datasets/",
),
{
@@ -2455,7 +2455,7 @@ def test_provider_compatibility_checks(labels: tuple[str, ...], expected_outputs
id="Airflow mypy checks on airflow files with model changes.",
),
pytest.param(
- ("task_sdk/src/airflow/sdk/a_file.py",),
+ ("task-sdk/src/airflow/sdk/a_file.py",),
{
"needs-mypy": "true",
"mypy-checks": "['mypy-providers', 'mypy-task-sdk']",
@@ -2594,7 +2594,7 @@ def test_pr_labels(
id="Caplog is in the git diff Providers",
),
pytest.param(
- ("task_sdk/tests/definitions/test_dag.py",),
+ ("task-sdk/tests/definitions/test_dag.py",),
(),
GithubEvents.PULL_REQUEST,
id="Caplog is in the git diff TaskSDK",
@@ -2651,7 +2651,7 @@ def test_is_log_mocked_in_the_tests_fail(
id="Caplog is in the git diff Providers",
),
pytest.param(
- ("task_sdk/tests/definitions/test_dag.py",),
+ ("task-sdk/tests/definitions/test_dag.py",),
(),
GithubEvents.PULL_REQUEST,
id="Caplog is in the git diff TaskSDK",
@@ -2711,7 +2711,7 @@ def test_is_log_mocked_in_the_tests_fail_formatted(
id="Caplog is in the git diff Providers",
),
pytest.param(
- ("task_sdk/tests/definitions/test_dag.py",),
+ ("task-sdk/tests/definitions/test_dag.py",),
(),
GithubEvents.PULL_REQUEST,
id="Caplog is in the git diff TaskSDK",
@@ -2759,7 +2759,7 @@ def test_is_log_mocked_in_the_tests_not_fail(
id="Caplog is in the git diff Providers",
),
pytest.param(
- ("task_sdk/tests/definitions/test_dag.py",),
+ ("task-sdk/tests/definitions/test_dag.py",),
(LOG_WITHOUT_MOCK_IN_TESTS_EXCEPTION_LABEL,),
GithubEvents.PULL_REQUEST,
id="Caplog is in the git diff TaskSDK",
diff --git a/docs/apache-airflow/administration-and-deployment/logging-monitoring/callbacks.rst b/docs/apache-airflow/administration-and-deployment/logging-monitoring/callbacks.rst
index 7d727bc01c421..417443192d0b0 100644
--- a/docs/apache-airflow/administration-and-deployment/logging-monitoring/callbacks.rst
+++ b/docs/apache-airflow/administration-and-deployment/logging-monitoring/callbacks.rst
@@ -97,4 +97,4 @@ In the following example, failures in any task call the ``task_failure_alert`` f
to be executed in the desired event. Simply pass a list of callback functions to the callback args when defining your DAG/task
callbacks: e.g ``on_failure_callback=[callback_func_1, callback_func_2]``
-Full list of variables available in ``context`` in :doc:`docs <../../templates-ref>` and `code `_.
+Full list of variables available in ``context`` in :doc:`docs <../../templates-ref>` and `code `_.
diff --git a/providers/openlineage/tests/unit/openlineage/extractors/test_manager.py b/providers/openlineage/tests/unit/openlineage/extractors/test_manager.py
index 5ad2e893c1aaa..256bc133f1a41 100644
--- a/providers/openlineage/tests/unit/openlineage/extractors/test_manager.py
+++ b/providers/openlineage/tests/unit/openlineage/extractors/test_manager.py
@@ -20,7 +20,7 @@
import logging
import tempfile
from datetime import datetime
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, Protocol
from unittest import mock
from unittest.mock import MagicMock
@@ -46,12 +46,16 @@
if TYPE_CHECKING:
try:
+ from airflow.sdk.api.datamodels._generated import TIRunContext
from airflow.sdk.definitions.context import Context
+
except ImportError:
# TODO: Remove once provider drops support for Airflow 2
+ # TIRunContext is only used in Airflow 3 tests
from airflow.utils.context import Context
- from task_sdk.tests.conftest import MakeTIContextCallable
+ TIRunContext = Any # type: ignore[misc, assignment]
+
if AIRFLOW_V_2_10_PLUS:
@@ -414,6 +418,24 @@ def set_dag(what: StartupDetails, dag_id: str, task: BaseOperator) -> RuntimeTas
return set_dag
+class MakeTIContextCallable(Protocol):
+ def __call__(
+ self,
+ dag_id: str = ...,
+ run_id: str = ...,
+ logical_date: str | datetime = ...,
+ data_interval_start: str | datetime = ...,
+ data_interval_end: str | datetime = ...,
+ clear_number: int = ...,
+ start_date: str | datetime = ...,
+ run_after: str | datetime = ...,
+ run_type: str = ...,
+ task_reschedule_count: int = ...,
+ conf: dict[str, Any] | None = ...,
+ ) -> TIRunContext: ...
+
+
+# Only needed in Airflow 3
@pytest.fixture
def make_ti_context() -> MakeTIContextCallable:
"""Factory for creating TIRunContext objects."""
diff --git a/pyproject.toml b/pyproject.toml
index 12028e1908e8a..0b11eb5163c0f 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -294,7 +294,7 @@ section-order = [
# Make sure we put the "dev" imports at the end, not as a third-party module
[tool.ruff.lint.isort.sections]
-testing = ["dev", "providers.tests", "task_sdk.tests", "tests_common", "tests"]
+testing = ["dev", "providers.tests", "tests_common", "tests"]
[tool.ruff.lint.extend-per-file-ignores]
"airflow/__init__.py" = ["F401", "TC004", "I002"]
@@ -916,6 +916,6 @@ members = [
"providers/ydb",
"providers/zendesk",
".",
- "task_sdk",
+ "task-sdk",
"devel-common",
]
diff --git a/scripts/ci/docker-compose/local.yml b/scripts/ci/docker-compose/local.yml
index f23a2a1f9c896..156f20a071c3f 100644
--- a/scripts/ci/docker-compose/local.yml
+++ b/scripts/ci/docker-compose/local.yml
@@ -89,8 +89,8 @@ services:
source: ../../../providers
target: /opt/airflow/providers
- type: bind
- source: ../../../task_sdk
- target: /opt/airflow/task_sdk
+ source: ../../../task-sdk
+ target: /opt/airflow/task-sdk
- type: bind
source: ../../../pyproject.toml
target: /opt/airflow/pyproject.toml
diff --git a/scripts/ci/docker-compose/providers-and-tests-sources.yml b/scripts/ci/docker-compose/providers-and-tests-sources.yml
index 91031e22cceae..2f5eff3835e31 100644
--- a/scripts/ci/docker-compose/providers-and-tests-sources.yml
+++ b/scripts/ci/docker-compose/providers-and-tests-sources.yml
@@ -30,10 +30,10 @@ services:
- ../../../empty:/opt/airflow/airflow
# But keep airflow/__init__.py to make hatchling happy when building empty package
- ../../../airflow/__init__.py:/opt/airflow/airflow/__init__.py
- # Remove task_sdk sources from container
- - ../../../empty:/opt/airflow/task_sdk/src
- # But keep task_sdk/a__init__.py to make hatchling happy when building empty package
- - ../../../task_sdk/src/airflow/sdk/__init__.py:/opt/airflow/task_sdk/src/airflow/sdk/__init__.py
+ # Remove task-sdk sources from container
+ - ../../../empty:/opt/airflow/task-sdk/src
+ # But keep task-sdk/a__init__.py to make hatchling happy when building empty package
+ - ../../../task-sdk/src/airflow/sdk/__init__.py:/opt/airflow/task-sdk/src/airflow/sdk/__init__.py
# And keep tests
- ../../../tests/:/opt/airflow/tests:cached
# Mount providers to make sure that we have the latest providers - both tests and sources
diff --git a/scripts/ci/docker-compose/remove-sources.yml b/scripts/ci/docker-compose/remove-sources.yml
index 5c31f0baee6f3..fd7b73c93dff1 100644
--- a/scripts/ci/docker-compose/remove-sources.yml
+++ b/scripts/ci/docker-compose/remove-sources.yml
@@ -22,8 +22,8 @@ services:
volumes:
# Removes airflow sources from container
- ../../../empty:/opt/airflow/airflow
- # Removes task_sdk sources from container
- - ../../../empty:/opt/airflow/task_sdk/src
+ # Removes task-sdk sources from container
+ - ../../../empty:/opt/airflow/task-sdk/src
# Remove all provider sources from container
# START automatically generated volumes by generate-volumes-for-sources pre-commit
- ../../../empty:/opt/airflow/providers/airbyte/src
diff --git a/scripts/ci/kubernetes/k8s_requirements.txt b/scripts/ci/kubernetes/k8s_requirements.txt
index 03908cd6818ef..0fbbf3827c3a0 100644
--- a/scripts/ci/kubernetes/k8s_requirements.txt
+++ b/scripts/ci/kubernetes/k8s_requirements.txt
@@ -1,4 +1,4 @@
-e ./providers/standard
--e ./task_sdk
+-e ./task-sdk
-e ./devel-common
-e ./providers/cncf/kubernetes
diff --git a/scripts/ci/pre_commit/base_operator_partial_arguments.py b/scripts/ci/pre_commit/base_operator_partial_arguments.py
index 070ffe767cccc..7b3b81bb838c4 100755
--- a/scripts/ci/pre_commit/base_operator_partial_arguments.py
+++ b/scripts/ci/pre_commit/base_operator_partial_arguments.py
@@ -27,9 +27,9 @@
ROOT_DIR = pathlib.Path(__file__).resolve().parents[3]
BASEOPERATOR_PY = ROOT_DIR.joinpath("airflow", "models", "baseoperator.py")
-SDK_BASEOPERATOR_PY = ROOT_DIR.joinpath("task_sdk", "src", "airflow", "sdk", "definitions", "baseoperator.py")
+SDK_BASEOPERATOR_PY = ROOT_DIR.joinpath("task-sdk", "src", "airflow", "sdk", "definitions", "baseoperator.py")
SDK_MAPPEDOPERATOR_PY = ROOT_DIR.joinpath(
- "task_sdk", "src", "airflow", "sdk", "definitions", "mappedoperator.py"
+ "task-sdk", "src", "airflow", "sdk", "definitions", "mappedoperator.py"
)
IGNORED = {
diff --git a/scripts/ci/pre_commit/mypy_folder.py b/scripts/ci/pre_commit/mypy_folder.py
index 09a643d45ccbd..ee4c35b076163 100755
--- a/scripts/ci/pre_commit/mypy_folder.py
+++ b/scripts/ci/pre_commit/mypy_folder.py
@@ -38,7 +38,7 @@
*[f"providers/{provider_id.replace('.', '/')}/src" for provider_id in get_all_provider_ids()],
"dev",
"docs",
- "task_sdk/src/airflow/sdk",
+ "task-sdk/src/airflow/sdk",
"all_providers",
]
@@ -69,8 +69,8 @@
elif mypy_folder.startswith("providers/"):
arguments.append(f"{Path(mypy_folder).parent.as_posix()}/tests")
namespace_packages = True
- if mypy_folder == "task_sdk/src/airflow/sdk":
- arguments.append("task_sdk/tests")
+ if mypy_folder == "task-sdk/src/airflow/sdk":
+ arguments.append("task-sdk/tests")
namespace_packages = True
if mypy_folder == "airflow":
arguments.append("tests")
diff --git a/scripts/ci/pre_commit/sync_init_decorator.py b/scripts/ci/pre_commit/sync_init_decorator.py
index 7b02136ead318..18efca02dd519 100755
--- a/scripts/ci/pre_commit/sync_init_decorator.py
+++ b/scripts/ci/pre_commit/sync_init_decorator.py
@@ -28,7 +28,7 @@
ROOT = pathlib.Path(__file__).resolve().parents[3]
PACKAGE_ROOT = ROOT.joinpath("airflow")
-SDK_DEFINITIONS_PKG = ROOT.joinpath("task_sdk", "src", "airflow", "sdk", "definitions")
+SDK_DEFINITIONS_PKG = ROOT.joinpath("task-sdk", "src", "airflow", "sdk", "definitions")
DAG_PY = SDK_DEFINITIONS_PKG.joinpath("dag.py")
TG_PY = SDK_DEFINITIONS_PKG.joinpath("taskgroup.py")
DECOS_TG_PY = PACKAGE_ROOT.joinpath("decorators", "task_group.py")
diff --git a/scripts/ci/pre_commit/template_context_key_sync.py b/scripts/ci/pre_commit/template_context_key_sync.py
index 270a3a695b263..f29f3a1a9c7af 100755
--- a/scripts/ci/pre_commit/template_context_key_sync.py
+++ b/scripts/ci/pre_commit/template_context_key_sync.py
@@ -27,9 +27,9 @@
ROOT_DIR = pathlib.Path(__file__).resolve().parents[3]
-TASKRUNNER_PY = ROOT_DIR.joinpath("task_sdk", "src", "airflow", "sdk", "execution_time", "task_runner.py")
+TASKRUNNER_PY = ROOT_DIR.joinpath("task-sdk", "src", "airflow", "sdk", "execution_time", "task_runner.py")
CONTEXT_PY = ROOT_DIR.joinpath("airflow", "utils", "context.py")
-CONTEXT_HINT = ROOT_DIR.joinpath("task_sdk", "src", "airflow", "sdk", "definitions", "context.py")
+CONTEXT_HINT = ROOT_DIR.joinpath("task-sdk", "src", "airflow", "sdk", "definitions", "context.py")
TEMPLATES_REF_RST = ROOT_DIR.joinpath("docs", "apache-airflow", "templates-ref.rst")
# These are only conditionally set
diff --git a/scripts/ci/testing/run_unit_tests.sh b/scripts/ci/testing/run_unit_tests.sh
index b602237ba7e52..cace8a1dfbfc1 100755
--- a/scripts/ci/testing/run_unit_tests.sh
+++ b/scripts/ci/testing/run_unit_tests.sh
@@ -135,7 +135,7 @@ function run_tests() {
core_tests
elif [[ "${TEST_GROUP}" == "providers" ]]; then
providers_tests
- elif [[ "${TEST_GROUP}" == "task_sdk" ]]; then
+ elif [[ "${TEST_GROUP}" == "task-sdk" ]]; then
task_sdk_tests
else
echo "Unknown test group: ${TEST_GROUP}"
diff --git a/scripts/docker/entrypoint_ci.sh b/scripts/docker/entrypoint_ci.sh
index e6abb50f73323..4caa698e24256 100755
--- a/scripts/docker/entrypoint_ci.sh
+++ b/scripts/docker/entrypoint_ci.sh
@@ -331,7 +331,7 @@ function check_force_lowest_dependencies() {
echo
fi
set -x
- uv pip install --python "$(which python)" --resolution lowest-direct --upgrade --editable ".${EXTRA}" --editable "./task_sdk" --editable "./devel-common"
+ uv pip install --python "$(which python)" --resolution lowest-direct --upgrade --editable ".${EXTRA}" --editable "./task-sdk" --editable "./devel-common"
set +x
}
diff --git a/scripts/docker/install_airflow.sh b/scripts/docker/install_airflow.sh
index d34ce3124ba36..5511c0d3afdef 100644
--- a/scripts/docker/install_airflow.sh
+++ b/scripts/docker/install_airflow.sh
@@ -49,7 +49,7 @@ function install_airflow() {
local installation_command_flags
if [[ ${AIRFLOW_INSTALLATION_METHOD} == "." ]]; then
# When installing from sources - we always use `--editable` mode
- installation_command_flags="--editable .[${AIRFLOW_EXTRAS}]${AIRFLOW_VERSION_SPECIFICATION} --editable ./task_sdk --editable ./devel-common"
+ installation_command_flags="--editable .[${AIRFLOW_EXTRAS}]${AIRFLOW_VERSION_SPECIFICATION} --editable ./task-sdk --editable ./devel-common"
while IFS= read -r -d '' pyproject_toml_file; do
project_folder=$(dirname ${pyproject_toml_file})
installation_command_flags="${installation_command_flags} --editable ${project_folder}"
diff --git a/scripts/in_container/install_airflow_and_providers.py b/scripts/in_container/install_airflow_and_providers.py
index 6892d99651e3b..33032a011c777 100755
--- a/scripts/in_container/install_airflow_and_providers.py
+++ b/scripts/in_container/install_airflow_and_providers.py
@@ -249,7 +249,7 @@ def find_installation_spec(
)
if airflow_extras:
airflow_package_spec += airflow_extras
- # We always install latest task_sdk - it's independent from Airflow
+ # We always install latest task-sdk - it's independent from Airflow
airflow_task_sdk_spec = find_airflow_task_sdk_package(extension)
if airflow_task_sdk_spec:
airflow_task_sdk_constraints_location = get_airflow_constraints_location(
@@ -287,7 +287,7 @@ def find_installation_spec(
github_repository=github_repository,
python_version=python_version,
)
- console.print(f"\nInstalling airflow task_sdk from remote spec {use_airflow_version}\n")
+ console.print(f"\nInstalling airflow task-sdk from remote spec {use_airflow_version}\n")
airflow_task_sdk_package = f"apache-airflow-task-sdk @ {use_airflow_version}"
airflow_constraints_location = get_airflow_constraints_location(
airflow_skip_constraints=airflow_skip_constraints,
@@ -318,7 +318,7 @@ def find_installation_spec(
python_version=python_version,
)
console.print(
- "\nDo not install airflow task_sdk. It should be installed automatically if needed "
+ "\nDo not install airflow task-sdk. It should be installed automatically if needed "
"by providers."
)
airflow_task_sdk_package = None
diff --git a/task_sdk/README.md b/task-sdk/README.md
similarity index 100%
rename from task_sdk/README.md
rename to task-sdk/README.md
diff --git a/task_sdk/dev/generate_models.py b/task-sdk/dev/generate_models.py
similarity index 100%
rename from task_sdk/dev/generate_models.py
rename to task-sdk/dev/generate_models.py
diff --git a/task_sdk/pyproject.toml b/task-sdk/pyproject.toml
similarity index 80%
rename from task_sdk/pyproject.toml
rename to task-sdk/pyproject.toml
index d980c0aae8259..8f983c3e6b5c1 100644
--- a/task_sdk/pyproject.toml
+++ b/task-sdk/pyproject.toml
@@ -141,3 +141,42 @@ custom-formatters = ['dev.datamodel_code_formatter',]
url = 'http://0.0.0.0:8080/execution/openapi.json'
output = 'src/airflow/sdk/api/datamodels/_generated.py'
+
+## pytest settings ##
+[tool.pytest.ini_options]
+addopts = [
+ "--tb=short",
+ "-rasl",
+ "--verbosity=2",
+ # Disable `flaky` plugin for pytest. This plugin conflicts with `rerunfailures` because provide the same marker.
+ "-p", "no:flaky",
+ # Disable `nose` builtin plugin for pytest. This feature is deprecated in 7.2 and will be removed in pytest>=8
+ "-p", "no:nose",
+ # Disable support of a legacy `LocalPath` in favor of stdlib `pathlib.Path`.
+ "-p", "no:legacypath",
+ # Disable warnings summary, because we use our warning summary.
+ "--disable-warnings",
+ "--asyncio-mode=strict",
+]
+
+norecursedirs = [
+ ".eggs",
+]
+log_level = "INFO"
+filterwarnings = [
+ "error::pytest.PytestCollectionWarning",
+ "error::pytest.PytestReturnNotNoneWarning",
+]
+python_files = [
+ "test_*.py",
+]
+testpaths = [
+ "tests",
+]
+asyncio_default_fixture_loop_scope = "function"
+
+pythonpath = "tests"
+
+# Keep temporary directories (created by `tmp_path`) for 2 recent runs only failed tests.
+tmp_path_retention_count = "2"
+tmp_path_retention_policy = "failed"
diff --git a/task_sdk/src/airflow/__init__.py b/task-sdk/src/airflow/__init__.py
similarity index 100%
rename from task_sdk/src/airflow/__init__.py
rename to task-sdk/src/airflow/__init__.py
diff --git a/task_sdk/src/airflow/sdk/__init__.py b/task-sdk/src/airflow/sdk/__init__.py
similarity index 100%
rename from task_sdk/src/airflow/sdk/__init__.py
rename to task-sdk/src/airflow/sdk/__init__.py
diff --git a/task_sdk/__init__.py b/task-sdk/src/airflow/sdk/api/__init__.py
similarity index 100%
rename from task_sdk/__init__.py
rename to task-sdk/src/airflow/sdk/api/__init__.py
diff --git a/task_sdk/src/airflow/sdk/api/client.py b/task-sdk/src/airflow/sdk/api/client.py
similarity index 100%
rename from task_sdk/src/airflow/sdk/api/client.py
rename to task-sdk/src/airflow/sdk/api/client.py
diff --git a/task_sdk/src/airflow/sdk/api/__init__.py b/task-sdk/src/airflow/sdk/api/datamodels/__init__.py
similarity index 100%
rename from task_sdk/src/airflow/sdk/api/__init__.py
rename to task-sdk/src/airflow/sdk/api/datamodels/__init__.py
diff --git a/task_sdk/src/airflow/sdk/api/datamodels/_generated.py b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py
similarity index 100%
rename from task_sdk/src/airflow/sdk/api/datamodels/_generated.py
rename to task-sdk/src/airflow/sdk/api/datamodels/_generated.py
diff --git a/task_sdk/src/airflow/sdk/api/datamodels/activities.py b/task-sdk/src/airflow/sdk/api/datamodels/activities.py
similarity index 100%
rename from task_sdk/src/airflow/sdk/api/datamodels/activities.py
rename to task-sdk/src/airflow/sdk/api/datamodels/activities.py
diff --git a/task_sdk/src/airflow/sdk/api/datamodels/__init__.py b/task-sdk/src/airflow/sdk/definitions/__init__.py
similarity index 100%
rename from task_sdk/src/airflow/sdk/api/datamodels/__init__.py
rename to task-sdk/src/airflow/sdk/definitions/__init__.py
diff --git a/task_sdk/src/airflow/sdk/definitions/__init__.py b/task-sdk/src/airflow/sdk/definitions/_internal/__init__.py
similarity index 100%
rename from task_sdk/src/airflow/sdk/definitions/__init__.py
rename to task-sdk/src/airflow/sdk/definitions/_internal/__init__.py
diff --git a/task_sdk/src/airflow/sdk/definitions/_internal/abstractoperator.py b/task-sdk/src/airflow/sdk/definitions/_internal/abstractoperator.py
similarity index 100%
rename from task_sdk/src/airflow/sdk/definitions/_internal/abstractoperator.py
rename to task-sdk/src/airflow/sdk/definitions/_internal/abstractoperator.py
diff --git a/task_sdk/src/airflow/sdk/definitions/_internal/contextmanager.py b/task-sdk/src/airflow/sdk/definitions/_internal/contextmanager.py
similarity index 100%
rename from task_sdk/src/airflow/sdk/definitions/_internal/contextmanager.py
rename to task-sdk/src/airflow/sdk/definitions/_internal/contextmanager.py
diff --git a/task_sdk/src/airflow/sdk/definitions/_internal/dag_parsing_context.py b/task-sdk/src/airflow/sdk/definitions/_internal/dag_parsing_context.py
similarity index 100%
rename from task_sdk/src/airflow/sdk/definitions/_internal/dag_parsing_context.py
rename to task-sdk/src/airflow/sdk/definitions/_internal/dag_parsing_context.py
diff --git a/task_sdk/src/airflow/sdk/definitions/_internal/decorators.py b/task-sdk/src/airflow/sdk/definitions/_internal/decorators.py
similarity index 100%
rename from task_sdk/src/airflow/sdk/definitions/_internal/decorators.py
rename to task-sdk/src/airflow/sdk/definitions/_internal/decorators.py
diff --git a/task_sdk/src/airflow/sdk/definitions/_internal/expandinput.py b/task-sdk/src/airflow/sdk/definitions/_internal/expandinput.py
similarity index 100%
rename from task_sdk/src/airflow/sdk/definitions/_internal/expandinput.py
rename to task-sdk/src/airflow/sdk/definitions/_internal/expandinput.py
diff --git a/task_sdk/src/airflow/sdk/definitions/_internal/mixins.py b/task-sdk/src/airflow/sdk/definitions/_internal/mixins.py
similarity index 100%
rename from task_sdk/src/airflow/sdk/definitions/_internal/mixins.py
rename to task-sdk/src/airflow/sdk/definitions/_internal/mixins.py
diff --git a/task_sdk/src/airflow/sdk/definitions/_internal/node.py b/task-sdk/src/airflow/sdk/definitions/_internal/node.py
similarity index 100%
rename from task_sdk/src/airflow/sdk/definitions/_internal/node.py
rename to task-sdk/src/airflow/sdk/definitions/_internal/node.py
diff --git a/task_sdk/src/airflow/sdk/definitions/_internal/templater.py b/task-sdk/src/airflow/sdk/definitions/_internal/templater.py
similarity index 100%
rename from task_sdk/src/airflow/sdk/definitions/_internal/templater.py
rename to task-sdk/src/airflow/sdk/definitions/_internal/templater.py
diff --git a/task_sdk/src/airflow/sdk/definitions/_internal/types.py b/task-sdk/src/airflow/sdk/definitions/_internal/types.py
similarity index 100%
rename from task_sdk/src/airflow/sdk/definitions/_internal/types.py
rename to task-sdk/src/airflow/sdk/definitions/_internal/types.py
diff --git a/task_sdk/src/airflow/sdk/definitions/asset/__init__.py b/task-sdk/src/airflow/sdk/definitions/asset/__init__.py
similarity index 100%
rename from task_sdk/src/airflow/sdk/definitions/asset/__init__.py
rename to task-sdk/src/airflow/sdk/definitions/asset/__init__.py
diff --git a/task_sdk/src/airflow/sdk/definitions/asset/decorators.py b/task-sdk/src/airflow/sdk/definitions/asset/decorators.py
similarity index 100%
rename from task_sdk/src/airflow/sdk/definitions/asset/decorators.py
rename to task-sdk/src/airflow/sdk/definitions/asset/decorators.py
diff --git a/task_sdk/src/airflow/sdk/definitions/asset/metadata.py b/task-sdk/src/airflow/sdk/definitions/asset/metadata.py
similarity index 100%
rename from task_sdk/src/airflow/sdk/definitions/asset/metadata.py
rename to task-sdk/src/airflow/sdk/definitions/asset/metadata.py
diff --git a/task_sdk/src/airflow/sdk/definitions/baseoperator.py b/task-sdk/src/airflow/sdk/definitions/baseoperator.py
similarity index 100%
rename from task_sdk/src/airflow/sdk/definitions/baseoperator.py
rename to task-sdk/src/airflow/sdk/definitions/baseoperator.py
diff --git a/task_sdk/src/airflow/sdk/definitions/connection.py b/task-sdk/src/airflow/sdk/definitions/connection.py
similarity index 100%
rename from task_sdk/src/airflow/sdk/definitions/connection.py
rename to task-sdk/src/airflow/sdk/definitions/connection.py
diff --git a/task_sdk/src/airflow/sdk/definitions/context.py b/task-sdk/src/airflow/sdk/definitions/context.py
similarity index 100%
rename from task_sdk/src/airflow/sdk/definitions/context.py
rename to task-sdk/src/airflow/sdk/definitions/context.py
diff --git a/task_sdk/src/airflow/sdk/definitions/dag.py b/task-sdk/src/airflow/sdk/definitions/dag.py
similarity index 100%
rename from task_sdk/src/airflow/sdk/definitions/dag.py
rename to task-sdk/src/airflow/sdk/definitions/dag.py
diff --git a/task_sdk/src/airflow/sdk/definitions/edges.py b/task-sdk/src/airflow/sdk/definitions/edges.py
similarity index 100%
rename from task_sdk/src/airflow/sdk/definitions/edges.py
rename to task-sdk/src/airflow/sdk/definitions/edges.py
diff --git a/task_sdk/src/airflow/sdk/definitions/macros.py b/task-sdk/src/airflow/sdk/definitions/macros.py
similarity index 100%
rename from task_sdk/src/airflow/sdk/definitions/macros.py
rename to task-sdk/src/airflow/sdk/definitions/macros.py
diff --git a/task_sdk/src/airflow/sdk/definitions/mappedoperator.py b/task-sdk/src/airflow/sdk/definitions/mappedoperator.py
similarity index 100%
rename from task_sdk/src/airflow/sdk/definitions/mappedoperator.py
rename to task-sdk/src/airflow/sdk/definitions/mappedoperator.py
diff --git a/task_sdk/src/airflow/sdk/definitions/param.py b/task-sdk/src/airflow/sdk/definitions/param.py
similarity index 100%
rename from task_sdk/src/airflow/sdk/definitions/param.py
rename to task-sdk/src/airflow/sdk/definitions/param.py
diff --git a/task_sdk/src/airflow/sdk/definitions/taskgroup.py b/task-sdk/src/airflow/sdk/definitions/taskgroup.py
similarity index 100%
rename from task_sdk/src/airflow/sdk/definitions/taskgroup.py
rename to task-sdk/src/airflow/sdk/definitions/taskgroup.py
diff --git a/task_sdk/src/airflow/sdk/definitions/template.py b/task-sdk/src/airflow/sdk/definitions/template.py
similarity index 100%
rename from task_sdk/src/airflow/sdk/definitions/template.py
rename to task-sdk/src/airflow/sdk/definitions/template.py
diff --git a/task_sdk/src/airflow/sdk/definitions/variable.py b/task-sdk/src/airflow/sdk/definitions/variable.py
similarity index 100%
rename from task_sdk/src/airflow/sdk/definitions/variable.py
rename to task-sdk/src/airflow/sdk/definitions/variable.py
diff --git a/task_sdk/src/airflow/sdk/definitions/xcom_arg.py b/task-sdk/src/airflow/sdk/definitions/xcom_arg.py
similarity index 100%
rename from task_sdk/src/airflow/sdk/definitions/xcom_arg.py
rename to task-sdk/src/airflow/sdk/definitions/xcom_arg.py
diff --git a/task_sdk/src/airflow/sdk/exceptions.py b/task-sdk/src/airflow/sdk/exceptions.py
similarity index 100%
rename from task_sdk/src/airflow/sdk/exceptions.py
rename to task-sdk/src/airflow/sdk/exceptions.py
diff --git a/task_sdk/src/airflow/sdk/execution_time/__init__.py b/task-sdk/src/airflow/sdk/execution_time/__init__.py
similarity index 100%
rename from task_sdk/src/airflow/sdk/execution_time/__init__.py
rename to task-sdk/src/airflow/sdk/execution_time/__init__.py
diff --git a/task_sdk/src/airflow/sdk/execution_time/comms.py b/task-sdk/src/airflow/sdk/execution_time/comms.py
similarity index 100%
rename from task_sdk/src/airflow/sdk/execution_time/comms.py
rename to task-sdk/src/airflow/sdk/execution_time/comms.py
diff --git a/task_sdk/src/airflow/sdk/execution_time/context.py b/task-sdk/src/airflow/sdk/execution_time/context.py
similarity index 100%
rename from task_sdk/src/airflow/sdk/execution_time/context.py
rename to task-sdk/src/airflow/sdk/execution_time/context.py
diff --git a/task_sdk/src/airflow/sdk/execution_time/execute_workload.py b/task-sdk/src/airflow/sdk/execution_time/execute_workload.py
similarity index 100%
rename from task_sdk/src/airflow/sdk/execution_time/execute_workload.py
rename to task-sdk/src/airflow/sdk/execution_time/execute_workload.py
diff --git a/task_sdk/src/airflow/sdk/execution_time/lazy_sequence.py b/task-sdk/src/airflow/sdk/execution_time/lazy_sequence.py
similarity index 100%
rename from task_sdk/src/airflow/sdk/execution_time/lazy_sequence.py
rename to task-sdk/src/airflow/sdk/execution_time/lazy_sequence.py
diff --git a/task_sdk/src/airflow/sdk/execution_time/secrets_masker.py b/task-sdk/src/airflow/sdk/execution_time/secrets_masker.py
similarity index 100%
rename from task_sdk/src/airflow/sdk/execution_time/secrets_masker.py
rename to task-sdk/src/airflow/sdk/execution_time/secrets_masker.py
diff --git a/task_sdk/src/airflow/sdk/execution_time/supervisor.py b/task-sdk/src/airflow/sdk/execution_time/supervisor.py
similarity index 100%
rename from task_sdk/src/airflow/sdk/execution_time/supervisor.py
rename to task-sdk/src/airflow/sdk/execution_time/supervisor.py
diff --git a/task_sdk/src/airflow/sdk/execution_time/task_runner.py b/task-sdk/src/airflow/sdk/execution_time/task_runner.py
similarity index 100%
rename from task_sdk/src/airflow/sdk/execution_time/task_runner.py
rename to task-sdk/src/airflow/sdk/execution_time/task_runner.py
diff --git a/task_sdk/src/airflow/sdk/log.py b/task-sdk/src/airflow/sdk/log.py
similarity index 100%
rename from task_sdk/src/airflow/sdk/log.py
rename to task-sdk/src/airflow/sdk/log.py
diff --git a/task_sdk/src/airflow/sdk/py.typed b/task-sdk/src/airflow/sdk/py.typed
similarity index 100%
rename from task_sdk/src/airflow/sdk/py.typed
rename to task-sdk/src/airflow/sdk/py.typed
diff --git a/task_sdk/src/airflow/sdk/types.py b/task-sdk/src/airflow/sdk/types.py
similarity index 100%
rename from task_sdk/src/airflow/sdk/types.py
rename to task-sdk/src/airflow/sdk/types.py
diff --git a/task_sdk/src/airflow/sdk/definitions/_internal/__init__.py b/task-sdk/tests/__init__.py
similarity index 100%
rename from task_sdk/src/airflow/sdk/definitions/_internal/__init__.py
rename to task-sdk/tests/__init__.py
diff --git a/task_sdk/tests/conftest.py b/task-sdk/tests/conftest.py
similarity index 99%
rename from task_sdk/tests/conftest.py
rename to task-sdk/tests/conftest.py
index b075d2e73ad65..4d2560106835a 100644
--- a/task_sdk/tests/conftest.py
+++ b/task-sdk/tests/conftest.py
@@ -89,7 +89,7 @@ def __call__(self, _: WrappedLogger, method_name: str, event: EventDict | bytes)
@pytest.fixture
def test_dags_dir():
- return Path(__file__).parent.joinpath("dags")
+ return Path(__file__).parent.joinpath("task_sdk", "dags")
@pytest.fixture
diff --git a/task-sdk/tests/task_sdk/__init__.py b/task-sdk/tests/task_sdk/__init__.py
new file mode 100644
index 0000000000000..059fa06ad2846
--- /dev/null
+++ b/task-sdk/tests/task_sdk/__init__.py
@@ -0,0 +1,45 @@
+# 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 httpx
+
+from airflow.sdk.api.client import Client
+from airflow.sdk.execution_time.comms import BundleInfo
+
+FAKE_BUNDLE = BundleInfo(name="anything", version="any")
+
+
+def make_client(transport: httpx.MockTransport) -> Client:
+ """Get a client with a custom transport."""
+ return Client(base_url="test://server", token="", transport=transport)
+
+
+def make_client_w_dry_run() -> Client:
+ """Get a client with dry_run enabled."""
+ return Client(base_url=None, dry_run=True, token="")
+
+
+def make_client_w_responses(responses: list[httpx.Response]) -> Client:
+ """Get a client with custom responses."""
+
+ def handle_request(request: httpx.Request) -> httpx.Response:
+ return responses.pop(0)
+
+ return Client(
+ base_url=None, dry_run=True, token="", mounts={"'http://": httpx.MockTransport(handle_request)}
+ )
diff --git a/task_sdk/tests/__init__.py b/task-sdk/tests/task_sdk/api/__init__.py
similarity index 100%
rename from task_sdk/tests/__init__.py
rename to task-sdk/tests/task_sdk/api/__init__.py
diff --git a/task_sdk/tests/api/test_client.py b/task-sdk/tests/task_sdk/api/test_client.py
similarity index 97%
rename from task_sdk/tests/api/test_client.py
rename to task-sdk/tests/task_sdk/api/test_client.py
index abadb7ef6832f..bcfe92d481c9b 100644
--- a/task_sdk/tests/api/test_client.py
+++ b/task-sdk/tests/task_sdk/api/test_client.py
@@ -23,8 +23,9 @@
import httpx
import pytest
import uuid6
+from task_sdk import make_client, make_client_w_dry_run, make_client_w_responses
-from airflow.sdk.api.client import Client, RemoteValidationError, ServerResponseError
+from airflow.sdk.api.client import RemoteValidationError, ServerResponseError
from airflow.sdk.api.datamodels._generated import ConnectionResponse, VariableResponse, XComResponse
from airflow.sdk.exceptions import ErrorType
from airflow.sdk.execution_time.comms import DeferTask, ErrorResponse, RescheduleTask
@@ -32,27 +33,6 @@
from airflow.utils.state import TerminalTIState
-def make_client(transport: httpx.MockTransport) -> Client:
- """Get a client with a custom transport"""
- return Client(base_url="test://server", token="", transport=transport)
-
-
-def make_client_w_dry_run() -> Client:
- """Get a client with dry_run enabled"""
- return Client(base_url=None, dry_run=True, token="")
-
-
-def make_client_w_responses(responses: list[httpx.Response]) -> Client:
- """Helper fixture to create a mock client with custom responses."""
-
- def handle_request(request: httpx.Request) -> httpx.Response:
- return responses.pop(0)
-
- return Client(
- base_url=None, dry_run=True, token="", mounts={"'http://": httpx.MockTransport(handle_request)}
- )
-
-
class TestClient:
@pytest.mark.parametrize(
["path", "json_response"],
diff --git a/task_sdk/tests/dags/dag_parsing_context.py b/task-sdk/tests/task_sdk/dags/dag_parsing_context.py
similarity index 100%
rename from task_sdk/tests/dags/dag_parsing_context.py
rename to task-sdk/tests/task_sdk/dags/dag_parsing_context.py
diff --git a/task_sdk/tests/dags/super_basic.py b/task-sdk/tests/task_sdk/dags/super_basic.py
similarity index 100%
rename from task_sdk/tests/dags/super_basic.py
rename to task-sdk/tests/task_sdk/dags/super_basic.py
diff --git a/task_sdk/tests/dags/super_basic_deferred_run.py b/task-sdk/tests/task_sdk/dags/super_basic_deferred_run.py
similarity index 100%
rename from task_sdk/tests/dags/super_basic_deferred_run.py
rename to task-sdk/tests/task_sdk/dags/super_basic_deferred_run.py
diff --git a/task_sdk/tests/dags/super_basic_run.py b/task-sdk/tests/task_sdk/dags/super_basic_run.py
similarity index 100%
rename from task_sdk/tests/dags/super_basic_run.py
rename to task-sdk/tests/task_sdk/dags/super_basic_run.py
diff --git a/task_sdk/tests/api/__init__.py b/task-sdk/tests/task_sdk/definitions/__init__.py
similarity index 100%
rename from task_sdk/tests/api/__init__.py
rename to task-sdk/tests/task_sdk/definitions/__init__.py
diff --git a/task_sdk/tests/definitions/__init__.py b/task-sdk/tests/task_sdk/definitions/_internal/__init__.py
similarity index 100%
rename from task_sdk/tests/definitions/__init__.py
rename to task-sdk/tests/task_sdk/definitions/_internal/__init__.py
diff --git a/task_sdk/tests/definitions/_internal/test_templater.py b/task-sdk/tests/task_sdk/definitions/_internal/test_templater.py
similarity index 100%
rename from task_sdk/tests/definitions/_internal/test_templater.py
rename to task-sdk/tests/task_sdk/definitions/_internal/test_templater.py
diff --git a/task_sdk/tests/definitions/conftest.py b/task-sdk/tests/task_sdk/definitions/conftest.py
similarity index 100%
rename from task_sdk/tests/definitions/conftest.py
rename to task-sdk/tests/task_sdk/definitions/conftest.py
diff --git a/task_sdk/tests/definitions/test_asset.py b/task-sdk/tests/task_sdk/definitions/test_asset.py
similarity index 100%
rename from task_sdk/tests/definitions/test_asset.py
rename to task-sdk/tests/task_sdk/definitions/test_asset.py
diff --git a/task_sdk/tests/definitions/test_asset_decorators.py b/task-sdk/tests/task_sdk/definitions/test_asset_decorators.py
similarity index 100%
rename from task_sdk/tests/definitions/test_asset_decorators.py
rename to task-sdk/tests/task_sdk/definitions/test_asset_decorators.py
diff --git a/task_sdk/tests/definitions/test_baseoperator.py b/task-sdk/tests/task_sdk/definitions/test_baseoperator.py
similarity index 100%
rename from task_sdk/tests/definitions/test_baseoperator.py
rename to task-sdk/tests/task_sdk/definitions/test_baseoperator.py
diff --git a/task_sdk/tests/definitions/test_connections.py b/task-sdk/tests/task_sdk/definitions/test_connections.py
similarity index 100%
rename from task_sdk/tests/definitions/test_connections.py
rename to task-sdk/tests/task_sdk/definitions/test_connections.py
diff --git a/task_sdk/tests/definitions/test_context.py b/task-sdk/tests/task_sdk/definitions/test_context.py
similarity index 100%
rename from task_sdk/tests/definitions/test_context.py
rename to task-sdk/tests/task_sdk/definitions/test_context.py
diff --git a/task_sdk/tests/definitions/test_dag.py b/task-sdk/tests/task_sdk/definitions/test_dag.py
similarity index 100%
rename from task_sdk/tests/definitions/test_dag.py
rename to task-sdk/tests/task_sdk/definitions/test_dag.py
diff --git a/task_sdk/tests/definitions/test_macros.py b/task-sdk/tests/task_sdk/definitions/test_macros.py
similarity index 100%
rename from task_sdk/tests/definitions/test_macros.py
rename to task-sdk/tests/task_sdk/definitions/test_macros.py
diff --git a/task_sdk/tests/definitions/test_mappedoperator.py b/task-sdk/tests/task_sdk/definitions/test_mappedoperator.py
similarity index 100%
rename from task_sdk/tests/definitions/test_mappedoperator.py
rename to task-sdk/tests/task_sdk/definitions/test_mappedoperator.py
diff --git a/task_sdk/tests/definitions/test_mixins.py b/task-sdk/tests/task_sdk/definitions/test_mixins.py
similarity index 100%
rename from task_sdk/tests/definitions/test_mixins.py
rename to task-sdk/tests/task_sdk/definitions/test_mixins.py
diff --git a/task_sdk/tests/definitions/test_param.py b/task-sdk/tests/task_sdk/definitions/test_param.py
similarity index 100%
rename from task_sdk/tests/definitions/test_param.py
rename to task-sdk/tests/task_sdk/definitions/test_param.py
diff --git a/task_sdk/tests/definitions/test_secrets_masker.py b/task-sdk/tests/task_sdk/definitions/test_secrets_masker.py
similarity index 100%
rename from task_sdk/tests/definitions/test_secrets_masker.py
rename to task-sdk/tests/task_sdk/definitions/test_secrets_masker.py
diff --git a/task_sdk/tests/definitions/test_template.py b/task-sdk/tests/task_sdk/definitions/test_template.py
similarity index 100%
rename from task_sdk/tests/definitions/test_template.py
rename to task-sdk/tests/task_sdk/definitions/test_template.py
diff --git a/task_sdk/tests/definitions/test_xcom_arg.py b/task-sdk/tests/task_sdk/definitions/test_xcom_arg.py
similarity index 100%
rename from task_sdk/tests/definitions/test_xcom_arg.py
rename to task-sdk/tests/task_sdk/definitions/test_xcom_arg.py
diff --git a/task_sdk/tests/definitions/_internal/__init__.py b/task-sdk/tests/task_sdk/execution_time/__init__.py
similarity index 100%
rename from task_sdk/tests/definitions/_internal/__init__.py
rename to task-sdk/tests/task_sdk/execution_time/__init__.py
diff --git a/task_sdk/tests/execution_time/conftest.py b/task-sdk/tests/task_sdk/execution_time/conftest.py
similarity index 100%
rename from task_sdk/tests/execution_time/conftest.py
rename to task-sdk/tests/task_sdk/execution_time/conftest.py
diff --git a/task_sdk/tests/execution_time/test_context.py b/task-sdk/tests/task_sdk/execution_time/test_context.py
similarity index 100%
rename from task_sdk/tests/execution_time/test_context.py
rename to task-sdk/tests/task_sdk/execution_time/test_context.py
diff --git a/task_sdk/tests/execution_time/test_supervisor.py b/task-sdk/tests/task_sdk/execution_time/test_supervisor.py
similarity index 99%
rename from task_sdk/tests/execution_time/test_supervisor.py
rename to task-sdk/tests/task_sdk/execution_time/test_supervisor.py
index ea639d9bdff71..66af26105fc7a 100644
--- a/task_sdk/tests/execution_time/test_supervisor.py
+++ b/task-sdk/tests/task_sdk/execution_time/test_supervisor.py
@@ -37,6 +37,7 @@
import pytest
import tenacity
from pytest_unordered import unordered
+from task_sdk import FAKE_BUNDLE, make_client
from uuid6 import uuid7
from airflow.executors.workloads import BundleInfo
@@ -80,9 +81,6 @@
from airflow.sdk.execution_time.task_runner import CommsDecoder
from airflow.utils import timezone, timezone as tz
-from task_sdk.tests.api.test_client import make_client
-from task_sdk.tests.execution_time.test_task_runner import FAKE_BUNDLE
-
if TYPE_CHECKING:
import kgb
diff --git a/task_sdk/tests/execution_time/test_task_runner.py b/task-sdk/tests/task_sdk/execution_time/test_task_runner.py
similarity index 99%
rename from task_sdk/tests/execution_time/test_task_runner.py
rename to task-sdk/tests/task_sdk/execution_time/test_task_runner.py
index aee592122f8a8..dd223d69bfb0e 100644
--- a/task_sdk/tests/execution_time/test_task_runner.py
+++ b/task-sdk/tests/task_sdk/execution_time/test_task_runner.py
@@ -29,6 +29,7 @@
from unittest.mock import patch
import pytest
+from task_sdk import FAKE_BUNDLE
from uuid6 import uuid7
from airflow.decorators import task as task_decorator
@@ -98,8 +99,6 @@
if TYPE_CHECKING:
from kgb import SpyAgency
-FAKE_BUNDLE = BundleInfo(name="anything", version="any")
-
def get_inline_dag(dag_id: str, task: BaseOperator) -> DAG:
"""Creates an inline dag and returns it based on dag_id and task."""
diff --git a/task_sdk/tests/test_log.py b/task-sdk/tests/task_sdk/log/test_log.py
similarity index 100%
rename from task_sdk/tests/test_log.py
rename to task-sdk/tests/task_sdk/log/test_log.py
diff --git a/task_sdk/tests/execution_time/__init__.py b/task_sdk/tests/execution_time/__init__.py
deleted file mode 100644
index 13a83393a9124..0000000000000
--- a/task_sdk/tests/execution_time/__init__.py
+++ /dev/null
@@ -1,16 +0,0 @@
-# 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.