From d2dd9e8759a38127191e045a964476b9f8bf500d Mon Sep 17 00:00:00 2001 From: fritz-astronomer <80706212+fritz-astronomer@users.noreply.github.com> Date: Thu, 12 Sep 2024 17:51:55 -0400 Subject: [PATCH] Airflow 2.0.2 Compatibility (#109) * cicd: bump act versions, add codecov token * fix: move Context import to TYPECHECKING * fix: rm Context import, test down to 2.0.2 * fix: add compat for 2.0, bump version --- .github/workflows/checks.yaml | 6 ++-- .github/workflows/deploy.yaml | 2 +- astronomer_starship/__init__.py | 2 +- .../compat/starship_compatability.py | 28 ++++++++++++++++++- .../providers/starship/operators/starship.py | 8 +++--- tests/validation_test.py | 2 +- 6 files changed, 38 insertions(+), 10 deletions(-) diff --git a/.github/workflows/checks.yaml b/.github/workflows/checks.yaml index 84e6415..3189e26 100644 --- a/.github/workflows/checks.yaml +++ b/.github/workflows/checks.yaml @@ -11,7 +11,7 @@ jobs: with: python-version: '3.10' cache: 'pip' - - uses: extractions/setup-just@v1 + - uses: extractions/setup-just@v2 - run: just install - run: | git fetch origin @@ -29,7 +29,9 @@ jobs: with: python-version: ${{ matrix.python-version }} cache: 'pip' - - uses: extractions/setup-just@v1 + - uses: extractions/setup-just@v2 - run: just install - run: just test-with-coverage - uses: codecov/codecov-action@v4 + with: + token: ${{ secrets.CODECOV_TOKEN }} diff --git a/.github/workflows/deploy.yaml b/.github/workflows/deploy.yaml index c304afe..ffc90d8 100644 --- a/.github/workflows/deploy.yaml +++ b/.github/workflows/deploy.yaml @@ -18,7 +18,7 @@ jobs: with: python-version: '3.10' cache: 'pip' - - uses: extractions/setup-just@v1 + - uses: extractions/setup-just@v2 - run: just build - uses: pypa/gh-action-pypi-publish@release/v1 with: diff --git a/astronomer_starship/__init__.py b/astronomer_starship/__init__.py index c8ac9c0..8ebf122 100644 --- a/astronomer_starship/__init__.py +++ b/astronomer_starship/__init__.py @@ -1,4 +1,4 @@ -__version__ = "2.0.5" +__version__ = "2.0.6" def get_provider_info(): diff --git a/astronomer_starship/compat/starship_compatability.py b/astronomer_starship/compat/starship_compatability.py index 0dbad63..8cc2756 100644 --- a/astronomer_starship/compat/starship_compatability.py +++ b/astronomer_starship/compat/starship_compatability.py @@ -994,6 +994,30 @@ def get_task_instances(self, dag_id: str, offset: int = 0, limit: int = 10): raise e +class StarshipAirflow20(StarshipAirflow21): + """ + - description does not exist in variables + - queued_at not on dag_run + """ + + def variable_attrs(self): + attrs = super().variable_attrs() + del attrs["description"] + return attrs + + def dag_runs_attrs(self): + attrs = super().dag_runs_attrs() + if "queued_at" in attrs["dag_runs"]["test_value"][0]: + del attrs["dag_runs"]["test_value"][0]["queued_at"] + return attrs + + def dag_run_attrs(self): + attrs = super().dag_run_attrs() + if "queued_at" in attrs: + del attrs["queued_at"] + return attrs + + class StarshipAirflow27(StarshipAirflow): """ - include_deferred is required in pools @@ -1078,8 +1102,10 @@ def __new__(cls, airflow_version: "Union[str, None]" = None) -> StarshipAirflow: return StarshipAirflow27() if int(minor) == 2: return StarshipAirflow22() - if int(minor) <= 1: + if int(minor) == 1: return StarshipAirflow21() + if int(minor) == 0: + return StarshipAirflow20() return StarshipAirflow() else: raise RuntimeError(f"Unsupported Airflow Version: {airflow_version}") diff --git a/astronomer_starship/providers/starship/operators/starship.py b/astronomer_starship/providers/starship/operators/starship.py index 629a88f..d7bffec 100644 --- a/astronomer_starship/providers/starship/operators/starship.py +++ b/astronomer_starship/providers/starship/operators/starship.py @@ -8,7 +8,6 @@ from airflow.decorators import task from airflow.exceptions import AirflowSkipException from airflow.models.baseoperator import BaseOperator -from airflow.utils.context import Context from airflow.utils.task_group import TaskGroup from astronomer_starship.providers.starship.hooks.starship import ( @@ -16,6 +15,7 @@ StarshipHttpHook, ) + # Compatability Notes: # - @task() is >=AF2.0 # - @task_group is >=AF2.1 @@ -37,7 +37,7 @@ def __init__(self, variable_key: Union[str, None] = None, **kwargs): super().__init__(**kwargs) self.variable_key = variable_key - def execute(self, context: Context) -> Any: + def execute(self, context) -> Any: logging.info("Getting Variable", self.variable_key) variables = self.source_hook.get_variables() variable: Union[dict, None] = ( @@ -90,7 +90,7 @@ def __init__(self, pool_name: Union[str, None] = None, **kwargs): super().__init__(**kwargs) self.pool_name = pool_name - def execute(self, context: Context) -> Any: + def execute(self, context) -> Any: logging.info("Getting Pool", self.pool_name) pool: Union[dict, None] = ( [v for v in self.source_hook.get_pools() if v["name"] == self.pool_name] @@ -140,7 +140,7 @@ def __init__(self, connection_id: Union[str, None] = None, **kwargs): super().__init__(**kwargs) self.connection_id = connection_id - def execute(self, context: Context) -> Any: + def execute(self, context) -> Any: logging.info("Getting Connection", self.connection_id) connection: Union[dict, None] = ( [ diff --git a/tests/validation_test.py b/tests/validation_test.py index 93a0938..3b5f305 100644 --- a/tests/validation_test.py +++ b/tests/validation_test.py @@ -35,7 +35,7 @@ "apache/airflow:2.3.4", "apache/airflow:2.2.4", "apache/airflow:2.1.3", - # "apache/airflow:2.0.0", + "apache/airflow:2.0.2", # # "apache/airflow:1.10.15", # # "apache/airflow:1.10.10", ]