feat(airflow): drop Airflow < 2.3 support + make plugin v2 the default

hsheth2 · hsheth2 · commit f111699bbefe · 2024-12-06T18:49:52.000-05:00
diff --git a/.github/workflows/airflow-plugin.yml b/.github/workflows/airflow-plugin.yml
@@ -34,29 +34,20 @@ jobs:
         include:
           # Note: this should be kept in sync with tox.ini.
           - python-version: "3.8"
-            extra_pip_requirements: "apache-airflow~=2.1.4"
-            extra_pip_extras: plugin-v1
-          - python-version: "3.8"
-            extra_pip_requirements: "apache-airflow~=2.2.4"
-            extra_pip_extras: plugin-v1
+            extra_pip_requirements: "apache-airflow~=2.3.4 -c https://raw.githubusercontent.com/apache/airflow/constraints-2.3.4/constraints-3.8.txt"
           - python-version: "3.10"
             extra_pip_requirements: "apache-airflow~=2.4.3"
-            extra_pip_extras: plugin-v2,test-airflow24
+            extra_pip_extras: test-airflow24
           - python-version: "3.10"
             extra_pip_requirements: "apache-airflow~=2.6.3 -c https://raw.githubusercontent.com/apache/airflow/constraints-2.6.3/constraints-3.10.txt"
-            extra_pip_extras: plugin-v2
           - python-version: "3.10"
             extra_pip_requirements: "apache-airflow~=2.7.3 -c https://raw.githubusercontent.com/apache/airflow/constraints-2.7.3/constraints-3.10.txt"
-            extra_pip_extras: plugin-v2
           - python-version: "3.10"
             extra_pip_requirements: "apache-airflow~=2.8.1 -c https://raw.githubusercontent.com/apache/airflow/constraints-2.8.1/constraints-3.10.txt"
-            extra_pip_extras: plugin-v2
           - python-version: "3.11"
             extra_pip_requirements: "apache-airflow~=2.9.3 -c https://raw.githubusercontent.com/apache/airflow/constraints-2.9.3/constraints-3.11.txt"
-            extra_pip_extras: plugin-v2
           - python-version: "3.11"
-            extra_pip_requirements: "apache-airflow~=2.10.2 -c https://raw.githubusercontent.com/apache/airflow/constraints-2.10.2/constraints-3.11.txt"
-            extra_pip_extras: plugin-v2
+            extra_pip_requirements: "apache-airflow~=2.10.3 -c https://raw.githubusercontent.com/apache/airflow/constraints-2.10.3/constraints-3.11.txt"
       fail-fast: false
     steps:
       - name: Set up JDK 17
diff --git a/docs/how/updating-datahub.md b/docs/how/updating-datahub.md
@@ -35,6 +35,8 @@ This file documents any backwards-incompatible changes in DataHub and assists pe
 
 - #11701: The Fivetran `sources_to_database` field is deprecated in favor of setting directly within `sources_to_platform_instance.<key>.database`.
 - #11742: For PowerBi ingestion, `use_powerbi_email` is now enabled by default when extracting ownership information.
+- #TODO: The DataHub Airflow plugin no longer supports Airflow 2.1 and Airflow 2.2.
+- #TODO: The DataHub Airflow plugin now defaults to the v2 plugin implementation.
 
 ### Breaking Changes
 
diff --git a/metadata-ingestion-modules/airflow-plugin/build.gradle b/metadata-ingestion-modules/airflow-plugin/build.gradle
@@ -13,7 +13,7 @@ if (!project.hasProperty("extra_pip_requirements")) {
     ext.extra_pip_requirements = ""
 }
 if (!project.hasProperty("extra_pip_extras")) {
-    ext.extra_pip_extras = "plugin-v2"
+    ext.extra_pip_extras = ""
 }
 // If extra_pip_extras is non-empty, we need to add a comma to the beginning of the string.
 if (extra_pip_extras != "") {
diff --git a/metadata-ingestion-modules/airflow-plugin/setup.py b/metadata-ingestion-modules/airflow-plugin/setup.py
@@ -24,8 +24,8 @@ def get_long_description():
 
 base_requirements = {
     f"acryl-datahub[datahub-rest]{_self_pin}",
-    # Actual dependencies.
-    "apache-airflow >= 2.0.2",
+    # We require Airflow 2.3.x, since we need the new DAG listener API.
+    "apache-airflow>=2.3.0",
 }
 
 plugins: Dict[str, Set[str]] = {
@@ -48,8 +48,9 @@ def get_long_description():
     },
 }
 
-# Include datahub-rest in the base requirements.
+# Require some plugins by default.
 base_requirements.update(plugins["datahub-rest"])
+base_requirements.update(plugins["plugin-v2"])
 
 
 mypy_stubs = {
diff --git a/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/datahub_listener.py b/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/datahub_listener.py
@@ -74,7 +74,7 @@ def hookimpl(f: _F) -> _F:  # type: ignore[misc] # noqa: F811
     "1",
 )
 _RUN_IN_THREAD_TIMEOUT = float(
-    os.getenv("DATAHUB_AIRFLOW_PLUGIN_RUN_IN_THREAD_TIMEOUT", 15)
+    os.getenv("DATAHUB_AIRFLOW_PLUGIN_RUN_IN_THREAD_TIMEOUT", 10)
 )
 _DATAHUB_CLEANUP_DAG = "Datahub_Cleanup"
 
@@ -102,6 +102,7 @@ def get_airflow_plugin_listener() -> Optional["DataHubListener"]:
                     "capture_tags": plugin_config.capture_tags_info,
                     "capture_ownership": plugin_config.capture_ownership_info,
                     "enable_extractors": plugin_config.enable_extractors,
+                    "render_templates": plugin_config.render_templates,
                     "disable_openlineage_plugin": plugin_config.disable_openlineage_plugin,
                 },
             )
diff --git a/metadata-ingestion-modules/airflow-plugin/tests/integration/test_plugin.py b/metadata-ingestion-modules/airflow-plugin/tests/integration/test_plugin.py
@@ -12,6 +12,7 @@
 import time
 from typing import Any, Iterator, Sequence
 
+import packaging.version
 import pytest
 import requests
 import tenacity
@@ -20,6 +21,7 @@
 from datahub.testing.compare_metadata_json import assert_metadata_files_equal
 
 from datahub_airflow_plugin._airflow_shims import (
+    AIRFLOW_VERSION,
     HAS_AIRFLOW_DAG_LISTENER_API,
     HAS_AIRFLOW_LISTENER_API,
     HAS_AIRFLOW_STANDALONE_CMD,
@@ -36,6 +38,8 @@
 
 DAG_TO_SKIP_INGESTION = "dag_to_skip"
 
+TEST_V1_PLUGIN = AIRFLOW_VERSION < packaging.version.parse("2.4.0")
+
 
 @dataclasses.dataclass
 class AirflowInstance:
@@ -242,6 +246,7 @@ def _run_airflow(
         # Note that we could also disable the RUN_IN_THREAD entirely,
         # but I want to minimize the difference between CI and prod.
         "DATAHUB_AIRFLOW_PLUGIN_RUN_IN_THREAD_TIMEOUT": "30",
+        "DATAHUB_AIRFLOW_PLUGIN_USE_V1_PLUGIN": "true" if is_v1 else "false",
         # Convenience settings.
         "AIRFLOW__DATAHUB__LOG_LEVEL": "DEBUG",
         "AIRFLOW__DATAHUB__DEBUG_EMITTER": "True",
@@ -361,15 +366,14 @@ class DagTestCase:
 @pytest.mark.parametrize(
     ["golden_filename", "test_case", "is_v1"],
     [
-        # On Airflow <= 2.2, test plugin v1.
         *[
             pytest.param(
                 f"v1_{test_case.dag_id}",
                 test_case,
                 True,
                 id=f"v1_{test_case.dag_id}",
                 marks=pytest.mark.skipif(
-                    HAS_AIRFLOW_LISTENER_API,
+                    not TEST_V1_PLUGIN,
                     reason="Not testing plugin v1 on newer Airflow versions",
                 ),
             )
diff --git a/metadata-ingestion-modules/airflow-plugin/tox.ini b/metadata-ingestion-modules/airflow-plugin/tox.ini
@@ -4,17 +4,16 @@
 # and then run "tox" from this directory.
 
 [tox]
-envlist = py38-airflow21, py38-airflow22, py310-airflow24, py310-airflow26, py310-airflow27, py310-airflow28, py311-airflow29, py311-airflow210
+envlist = py38-airflow23, py310-airflow24, py310-airflow26, py310-airflow27, py310-airflow28, py311-airflow29, py311-airflow210
 
 [testenv]
 use_develop = true
-extras = dev,integration-tests,plugin-v1
+extras = dev,integration-tests,plugin-v1,plugin-v2
 deps =
     # This should be kept in sync with the Github Actions matrix.
     -e ../../metadata-ingestion/
     # Airflow version
-    airflow21: apache-airflow~=2.1.0
-    airflow22: apache-airflow~=2.2.0
+    airflow23: apache-airflow~=2.3.0
     airflow24: apache-airflow~=2.4.0
     airflow26: apache-airflow~=2.6.0
     airflow27: apache-airflow~=2.7.0
@@ -24,14 +23,15 @@ deps =
 
     # Respect the Airflow constraints files.
     # We can't make ourselves work with the constraints of Airflow < 2.3.
+    py38-airflow23: -c https://raw.githubusercontent.com/apache/airflow/constraints-2.3.4/constraints-3.8.txt
     # The Airflow 2.4 constraints file requires a version of the sqlite provider whose
     # hook type is missing the `conn_name_attr` property.
     ; py310-airflow24: -c https://raw.githubusercontent.com/apache/airflow/constraints-2.4.3/constraints-3.10.txt
     py310-airflow26: -c https://raw.githubusercontent.com/apache/airflow/constraints-2.6.3/constraints-3.10.txt
     py310-airflow27: -c https://raw.githubusercontent.com/apache/airflow/constraints-2.7.3/constraints-3.10.txt
     py310-airflow28: -c https://raw.githubusercontent.com/apache/airflow/constraints-2.8.1/constraints-3.10.txt
     py311-airflow29: -c https://raw.githubusercontent.com/apache/airflow/constraints-2.9.3/constraints-3.11.txt
-    py311-airflow210: -c https://raw.githubusercontent.com/apache/airflow/constraints-2.10.2/constraints-3.11.txt
+    py311-airflow210: -c https://raw.githubusercontent.com/apache/airflow/constraints-2.10.3/constraints-3.11.txt
 
     # Before pinning to the constraint files, we previously left the dependencies
     # more open. There were a number of packages for which this caused issues.
@@ -55,10 +55,6 @@ deps =
 commands =
     pytest --cov-append {posargs}
 
-# For Airflow 2.4+, add the plugin-v2 extra.
+# For Airflow 2.4, add a few extra requirements.
 [testenv:py310-airflow24]
-extras = dev,integration-tests,plugin-v2,test-airflow24
-
-[testenv:py3{10,11}-airflow{26,27,28,29,210}]
-extras = dev,integration-tests,plugin-v2
-
+extras = dev,integration-tests,plugin-v1,plugin-v2,test-airflow24

Original file line number	Diff line number	Diff line change
`@@ -13,7 +13,7 @@ if (!project.hasProperty("extra_pip_requirements")) {`
`13`	`13`	`ext.extra_pip_requirements = ""`
`14`	`14`	`}`
`15`	`15`	`if (!project.hasProperty("extra_pip_extras")) {`
`16`		`- ext.extra_pip_extras = "plugin-v2"`
	`16`	`+ ext.extra_pip_extras = ""`
`17`	`17`	`}`
`18`	`18`	`// If extra_pip_extras is non-empty, we need to add a comma to the beginning of the string.`
`19`	`19`	`if (extra_pip_extras != "") {`
Original file line number	Diff line number	Diff line change
`@@ -74,7 +74,7 @@ def hookimpl(f: _F) -> _F: # type: ignore[misc] # noqa: F811`
`74`	`74`	`"1",`
`75`	`75`	`)`
`76`	`76`	`_RUN_IN_THREAD_TIMEOUT = float(`
`77`		`- os.getenv("DATAHUB_AIRFLOW_PLUGIN_RUN_IN_THREAD_TIMEOUT", 15)`
	`77`	`+ os.getenv("DATAHUB_AIRFLOW_PLUGIN_RUN_IN_THREAD_TIMEOUT", 10)`
`78`	`78`	`)`
`79`	`79`	`_DATAHUB_CLEANUP_DAG = "Datahub_Cleanup"`
`80`	`80`
`@@ -102,6 +102,7 @@ def get_airflow_plugin_listener() -> Optional["DataHubListener"]:`
`102`	`102`	`"capture_tags": plugin_config.capture_tags_info,`
`103`	`103`	`"capture_ownership": plugin_config.capture_ownership_info,`
`104`	`104`	`"enable_extractors": plugin_config.enable_extractors,`
	`105`	`+ "render_templates": plugin_config.render_templates,`
`105`	`106`	`"disable_openlineage_plugin": plugin_config.disable_openlineage_plugin,`
`106`	`107`	`},`
`107`	`108`	`)`