From 229a3446c79a8b2eef93f16d3a882fbd30521a1e Mon Sep 17 00:00:00 2001 From: pankajastro Date: Wed, 9 Apr 2025 14:31:05 +0530 Subject: [PATCH 1/3] Improve AF3 tooling --- cosmos/operators/base.py | 12 +++++++++++- scripts/airflow3/env.sh | 6 ++++-- scripts/airflow3/requirements.txt | 12 +++++------- scripts/airflow3/setup.sh | 17 ++++------------- 4 files changed, 24 insertions(+), 23 deletions(-) diff --git a/cosmos/operators/base.py b/cosmos/operators/base.py index 552172f746..bd107853e3 100644 --- a/cosmos/operators/base.py +++ b/cosmos/operators/base.py @@ -7,9 +7,19 @@ from pathlib import Path from typing import Any, Sequence, Tuple +import airflow import yaml from airflow.utils.context import Context, context_merge -from airflow.utils.operator_helpers import context_to_airflow_vars +from packaging.version import Version + +from cosmos.constants import _AIRFLOW3_VERSION + +AIRFLOW_VERSION = Version(airflow.__version__) + +if AIRFLOW_VERSION < _AIRFLOW3_VERSION: + from airflow.utils.operator_helpers import context_to_airflow_vars +else: + from airflow.sdk.execution_time.context import context_to_airflow_vars # type: ignore from airflow.utils.strings import to_boolean from cosmos.dbt.executable import get_system_dbt diff --git a/scripts/airflow3/env.sh b/scripts/airflow3/env.sh index 5f1ecf5791..b87b0b28d0 100644 --- a/scripts/airflow3/env.sh +++ b/scripts/airflow3/env.sh @@ -1,8 +1,10 @@ #!/bin/bash -set -e -set -x +#set -e +#set -x +PYTHONPATH="$PWD" +export PYTHONPATH AIRFLOW_HOME="$PWD/scripts/airflow3" export AIRFLOW_HOME export AIRFLOW__LOGGING__BASE_LOG_FOLDER="$AIRFLOW_HOME/logs" diff --git a/scripts/airflow3/requirements.txt b/scripts/airflow3/requirements.txt index 7ce5ec81b1..decc3f9ed8 100644 --- a/scripts/airflow3/requirements.txt +++ b/scripts/airflow3/requirements.txt @@ -1,9 +1,7 @@ ---find-links https://dist.apache.org/repos/dist/dev/airflow/3.0.0b2/ -apache-airflow==3.0.0b2 -apache-airflow-task-sdk==1.0.0b2 -apache-airflow-providers-standard==1.0.0b2 -apache-airflow-providers-fab==2.0.0b2 -apache-airflow-providers-celery==3.11.0b2 -apache-airflow-providers-cncf-kubernetes==10.4.0b2 +apache-airflow +apache-airflow-task-sdk +apache-airflow-providers-standard +apache-airflow-providers-fab psycopg2 asyncpg +dbt-postgres diff --git a/scripts/airflow3/setup.sh b/scripts/airflow3/setup.sh index 315d186598..72e2147de5 100644 --- a/scripts/airflow3/setup.sh +++ b/scripts/airflow3/setup.sh @@ -4,8 +4,8 @@ set -e # Create a UV virtual environment named 'env' (you can change this as needed) -echo "Creating UV virtual environment at $(pwd)/tools" -uv venv venv --directory "$(pwd)/scripts/airflow3" +echo "Creating virtual environment at $(pwd)/tools" +python3 -m venv "$(pwd)/scripts/airflow3/venv" # Activate the virtual environment echo "Activating virtual environment..." @@ -13,17 +13,8 @@ source "$(pwd)/scripts/airflow3/venv/bin/activate" # Install dependencies in the virtual environment echo "Installing dependencies..." -uv pip install -r "$(pwd)/scripts/airflow3/requirements.txt" +pip3 install --pre -r "$(pwd)/scripts/airflow3/requirements.txt" -# Install Cosmos -uv pip install build -rm -rf "$(pwd)/dist/" -python3 -m build -for wheel in "$(pwd)"/dist/*.whl; do - uv pip install "$wheel" -done - -uv pip install dbt-core -uv pip install dbt-postgres +pip3 install ".[test]" echo "UV virtual environment setup and dependencies installed successfully!" From f04c5e7fd5126da32734506f3657c4ed8dafb9ad Mon Sep 17 00:00:00 2001 From: pankajastro Date: Wed, 9 Apr 2025 16:02:50 +0530 Subject: [PATCH 2/3] Apply review suggestion --- cosmos/operators/base.py | 11 +++-------- scripts/airflow3/README.md | 4 ++-- scripts/airflow3/setup.sh | 4 ++-- 3 files changed, 7 insertions(+), 12 deletions(-) diff --git a/cosmos/operators/base.py b/cosmos/operators/base.py index bd107853e3..04b5955649 100644 --- a/cosmos/operators/base.py +++ b/cosmos/operators/base.py @@ -7,19 +7,14 @@ from pathlib import Path from typing import Any, Sequence, Tuple -import airflow import yaml from airflow.utils.context import Context, context_merge -from packaging.version import Version -from cosmos.constants import _AIRFLOW3_VERSION - -AIRFLOW_VERSION = Version(airflow.__version__) - -if AIRFLOW_VERSION < _AIRFLOW3_VERSION: +try: from airflow.utils.operator_helpers import context_to_airflow_vars -else: +except ImportError: # pragma: no cover from airflow.sdk.execution_time.context import context_to_airflow_vars # type: ignore + from airflow.utils.strings import to_boolean from cosmos.dbt.executable import get_system_dbt diff --git a/scripts/airflow3/README.md b/scripts/airflow3/README.md index 02c37f4380..dfe8f174ba 100644 --- a/scripts/airflow3/README.md +++ b/scripts/airflow3/README.md @@ -49,7 +49,7 @@ sh scripts/airflow3/setup.sh Activate the virtual env created in previous step and run airflow ```commandline -source "$(pwd)/scripts/airflow3/venv/bin/activate" +source "$(pwd)/scripts/airflow3/venv-af3/bin/activate" airflow standalone ``` @@ -67,7 +67,7 @@ Once Airflow is running, you can also run tests. ```commandline source scripts/airflow3/env.sh -source "$(pwd)/scripts/airflow3/venv/bin/activate" +source "$(pwd)/scripts/airflow3/venv-af3/bin/activate" sh scripts/airflow3/tests.sh ``` diff --git a/scripts/airflow3/setup.sh b/scripts/airflow3/setup.sh index 72e2147de5..79cc7f2eb8 100644 --- a/scripts/airflow3/setup.sh +++ b/scripts/airflow3/setup.sh @@ -5,11 +5,11 @@ set -e # Create a UV virtual environment named 'env' (you can change this as needed) echo "Creating virtual environment at $(pwd)/tools" -python3 -m venv "$(pwd)/scripts/airflow3/venv" +python3 -m venv "$(pwd)/scripts/airflow3/venv-af3" # Activate the virtual environment echo "Activating virtual environment..." -source "$(pwd)/scripts/airflow3/venv/bin/activate" +source "$(pwd)/scripts/airflow3/venv-af3/bin/activate" # Install dependencies in the virtual environment echo "Installing dependencies..." From 1404886eb10418c342dd5a069136403360236687 Mon Sep 17 00:00:00 2001 From: pankajastro Date: Wed, 9 Apr 2025 16:03:37 +0530 Subject: [PATCH 3/3] Apply review suggestion --- scripts/airflow3/env.sh | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/scripts/airflow3/env.sh b/scripts/airflow3/env.sh index b87b0b28d0..0631049c82 100644 --- a/scripts/airflow3/env.sh +++ b/scripts/airflow3/env.sh @@ -1,7 +1,6 @@ #!/bin/bash -#set -e -#set -x +set -e PYTHONPATH="$PWD" export PYTHONPATH