Improve suffix handling for provider-generated dependencies (#38029)

In order to better address chicken-egg provider problem (in the
light of upcoming release of 2.9.0 and separated FAB provider, we
should make sure that providers released together with Airflow in
the same pre-release version, but also with another providers they
refer to, have the same suffix for version.

For example:

* FAB released with dev0 should refer to airflow >= x.y.z.dev0
* FAB released with b0 (normalized from beta0) should refer to
  airflow>=x.y.zb0
* Postgres released with rc1 should refer to common-sql >= x.y.zrc1

This will only work when the packages released are released together
in a single wave (which we should do now for Airlfow and FAB.

Later on - this will be less of a problem, because this is the only
time where we do not have the X.Y.Z version of Airflow released and
mandatory provider (FAB) has >=X.Y.Z. At the same time airflow has
fab set as mandatory provider so we need to make sure both of them
are released in the same wave.

By releasing them together, we solve the chicken-egg provider
problem for mandatory providers.

airflow/providers/fab/provider.yaml

@@ -37,7 +37,7 @@ versions:
   - 1.0.0
 
 dependencies:
-  - apache-airflow>=2.9.0.dev0
+  - apache-airflow>=2.9.0
   - flask>=2.2,<2.3
   # We are tightly coupled with FAB version as we vendored-in part of FAB code related to security manager
   # This is done as part of preparation to removing FAB as dependency, but we are not ready for it yet

dev/breeze/src/airflow_breeze/utils/packages.py

@@ -392,6 +392,25 @@ def get_wheel_package_name(provider_id: str) -> str:
     return "apache_airflow_providers_" + provider_id.replace(".", "_")
 
 
+def apply_version_suffix(install_clause: str, version_suffix: str) -> str:
+    if install_clause.startswith("apache-airflow") and ">=" in install_clause and version_suffix:
+        # Applies version suffix to the apache-airflow and provider package dependencies to make
+        # sure that pre-release versions have correct limits - this address the issue with how
+        # pip handles pre-release versions when packages are pre-release and refer to each other - we
+        # need to make sure that all our >= references for all apache-airflow packages in pre-release
+        # versions of providers contain the same suffix as the provider itself.
+        # For example `apache-airflow-providers-fab==2.0.0.dev0` should refer to
+        # `apache-airflow>=2.9.0.dev0` and not `apache-airflow>=2.9.0` because both packages are
+        # released together and >= 2.9.0 is not correct reference for 2.9.0.dev0 version of Airflow.
+        prefix, version = install_clause.split(">=")
+        from packaging.version import Version
+
+        base_version = Version(version).base_version
+        target_version = Version(str(base_version) + "." + version_suffix)
+        return prefix + ">=" + str(target_version)
+    return install_clause
+
+
 def get_install_requirements(provider_id: str, version_suffix: str) -> str:
     """
     Returns install requirements for the package.
@@ -401,32 +420,15 @@ def get_install_requirements(provider_id: str, version_suffix: str) -> str:
 
     :return: install requirements of the package
     """
-
-    def apply_version_suffix(install_clause: str) -> str:
-        if (
-            install_clause.startswith("apache-airflow")
-            and ">=" in install_clause
-            and version_suffix != ""
-            and not install_clause.endswith(version_suffix)
-        ):
-            # This is workaround for `pip` way of handling `--pre` installation switch. It apparently does
-            # not modify the meaning of `install_requires` to include also pre-releases, so we need to
-            # modify our internal provider and airflow package version references to include all pre-releases
-            # including all development releases. When you specify dependency as >= X.Y.Z, and you
-            # have packages X.Y.Zdev0 or X.Y.Zrc1 in a local file, such package is not considered
-            # as fulfilling the requirement even if `--pre` switch is used.
-            return install_clause + ".dev0"
-        return install_clause
-
     if provider_id in get_removed_provider_ids():
         dependencies = get_provider_requirements(provider_id)
     else:
         dependencies = PROVIDER_DEPENDENCIES.get(provider_id)["deps"]
-    install_requires = [apply_version_suffix(clause) for clause in dependencies]
+    install_requires = [apply_version_suffix(clause, version_suffix) for clause in dependencies]
     return "".join(f'\n    "{ir}",' for ir in install_requires)
 
 
-def get_package_extras(provider_id: str) -> dict[str, list[str]]:
+def get_package_extras(provider_id: str, version_suffix: str) -> dict[str, list[str]]:
     """
     Finds extras for the package specified.
 
@@ -458,6 +460,8 @@ def get_package_extras(provider_id: str) -> dict[str, list[str]]:
                     extras_dict[name].append(new_dependency)
             else:
                 extras_dict[name] = dependencies
+    for extra, dependencies in extras_dict.items():
+        extras_dict[extra] = [apply_version_suffix(clause, version_suffix) for clause in dependencies]
     return extras_dict
 
 
@@ -572,7 +576,9 @@ def get_provider_jinja_context(
         "INSTALL_REQUIREMENTS": get_install_requirements(
             provider_id=provider_details.provider_id, version_suffix=version_suffix
         ),
-        "EXTRAS_REQUIREMENTS": get_package_extras(provider_id=provider_details.provider_id),
+        "EXTRAS_REQUIREMENTS": get_package_extras(
+            provider_id=provider_details.provider_id, version_suffix=version_suffix
+        ),
         "CHANGELOG_RELATIVE_PATH": os.path.relpath(
             provider_details.source_provider_package_path,
             provider_details.documentation_provider_package_path,

dev/breeze/tests/test_packages.py

@@ -158,37 +158,150 @@ def test_get_documentation_package_path():
     assert get_documentation_package_path("apache.hdfs") == DOCS_ROOT / "apache-airflow-providers-apache-hdfs"
 
 
-def test_get_install_requirements():
-    assert (
-        get_install_requirements("asana", "").strip()
-        == """
+@pytest.mark.parametrize(
+    "provider, version_suffix, expected",
+    [
+        pytest.param(
+            "fab",
+            "",
+            """
+    "apache-airflow>=2.9.0",
+    "flask-appbuilder==4.3.11",
+    "flask-login>=0.6.2",
+    "flask>=2.2,<2.3",
+    "google-re2>=1.0",
+    """,
+            id="No suffix fab",
+        ),
+        pytest.param(
+            "fab",
+            "dev0",
+            """
+    "apache-airflow>=2.9.0.dev0",
+    "flask-appbuilder==4.3.11",
+    "flask-login>=0.6.2",
+    "flask>=2.2,<2.3",
+    "google-re2>=1.0",
+    """,
+            id="dev0 suffix fab",
+        ),
+        pytest.param(
+            "fab",
+            "beta0",
+            """
+    "apache-airflow>=2.9.0b0",
+    "flask-appbuilder==4.3.11",
+    "flask-login>=0.6.2",
+    "flask>=2.2,<2.3",
+    "google-re2>=1.0",
+    """,
+            id="beta0 suffix fab",
+        ),
+        pytest.param(
+            "postgres",
+            "beta0",
+            """
+    "apache-airflow-providers-common-sql>=1.3.1b0",
+    "apache-airflow>=2.6.0b0",
+    "psycopg2-binary>=2.8.0",
+    """,
+            id="beta0 suffix postgres",
+        ),
+        pytest.param(
+            "postgres",
+            "",
+            """
+    "apache-airflow-providers-common-sql>=1.3.1",
     "apache-airflow>=2.6.0",
-    "asana>=0.10,<4.0.0",
-""".strip()
-    )
+    "psycopg2-binary>=2.8.0",
+    """,
+            id="No suffix postgres",
+        ),
+    ],
+)
+def test_get_install_requirements(provider: str, version_suffix: str, expected: str):
+    assert get_install_requirements(provider, version_suffix).strip() == expected.strip()
 
 
-def test_get_package_extras():
-    assert get_package_extras("google") == {
-        "amazon": ["apache-airflow-providers-amazon>=2.6.0"],
-        "apache.beam": ["apache-airflow-providers-apache-beam", "apache-beam[gcp]"],
-        "apache.cassandra": ["apache-airflow-providers-apache-cassandra"],
-        "cncf.kubernetes": ["apache-airflow-providers-cncf-kubernetes>=7.2.0"],
-        "common.sql": ["apache-airflow-providers-common-sql"],
-        "facebook": ["apache-airflow-providers-facebook>=2.2.0"],
-        "leveldb": ["plyvel"],
-        "microsoft.azure": ["apache-airflow-providers-microsoft-azure"],
-        "microsoft.mssql": ["apache-airflow-providers-microsoft-mssql"],
-        "mysql": ["apache-airflow-providers-mysql"],
-        "openlineage": ["apache-airflow-providers-openlineage"],
-        "oracle": ["apache-airflow-providers-oracle>=3.1.0"],
-        "postgres": ["apache-airflow-providers-postgres"],
-        "presto": ["apache-airflow-providers-presto"],
-        "salesforce": ["apache-airflow-providers-salesforce"],
-        "sftp": ["apache-airflow-providers-sftp"],
-        "ssh": ["apache-airflow-providers-ssh"],
-        "trino": ["apache-airflow-providers-trino"],
-    }
+@pytest.mark.parametrize(
+    "version_suffix, expected",
+    [
+        pytest.param(
+            "",
+            {
+                "amazon": ["apache-airflow-providers-amazon>=2.6.0"],
+                "apache.beam": ["apache-airflow-providers-apache-beam", "apache-beam[gcp]"],
+                "apache.cassandra": ["apache-airflow-providers-apache-cassandra"],
+                "cncf.kubernetes": ["apache-airflow-providers-cncf-kubernetes>=7.2.0"],
+                "common.sql": ["apache-airflow-providers-common-sql"],
+                "facebook": ["apache-airflow-providers-facebook>=2.2.0"],
+                "leveldb": ["plyvel"],
+                "microsoft.azure": ["apache-airflow-providers-microsoft-azure"],
+                "microsoft.mssql": ["apache-airflow-providers-microsoft-mssql"],
+                "mysql": ["apache-airflow-providers-mysql"],
+                "openlineage": ["apache-airflow-providers-openlineage"],
+                "oracle": ["apache-airflow-providers-oracle>=3.1.0"],
+                "postgres": ["apache-airflow-providers-postgres"],
+                "presto": ["apache-airflow-providers-presto"],
+                "salesforce": ["apache-airflow-providers-salesforce"],
+                "sftp": ["apache-airflow-providers-sftp"],
+                "ssh": ["apache-airflow-providers-ssh"],
+                "trino": ["apache-airflow-providers-trino"],
+            },
+            id="No suffix",
+        ),
+        pytest.param(
+            "dev0",
+            {
+                "amazon": ["apache-airflow-providers-amazon>=2.6.0.dev0"],
+                "apache.beam": ["apache-airflow-providers-apache-beam", "apache-beam[gcp]"],
+                "apache.cassandra": ["apache-airflow-providers-apache-cassandra"],
+                "cncf.kubernetes": ["apache-airflow-providers-cncf-kubernetes>=7.2.0.dev0"],
+                "common.sql": ["apache-airflow-providers-common-sql"],
+                "facebook": ["apache-airflow-providers-facebook>=2.2.0.dev0"],
+                "leveldb": ["plyvel"],
+                "microsoft.azure": ["apache-airflow-providers-microsoft-azure"],
+                "microsoft.mssql": ["apache-airflow-providers-microsoft-mssql"],
+                "mysql": ["apache-airflow-providers-mysql"],
+                "openlineage": ["apache-airflow-providers-openlineage"],
+                "oracle": ["apache-airflow-providers-oracle>=3.1.0.dev0"],
+                "postgres": ["apache-airflow-providers-postgres"],
+                "presto": ["apache-airflow-providers-presto"],
+                "salesforce": ["apache-airflow-providers-salesforce"],
+                "sftp": ["apache-airflow-providers-sftp"],
+                "ssh": ["apache-airflow-providers-ssh"],
+                "trino": ["apache-airflow-providers-trino"],
+            },
+            id="With dev0 suffix",
+        ),
+        pytest.param(
+            "beta0",
+            {
+                "amazon": ["apache-airflow-providers-amazon>=2.6.0b0"],
+                "apache.beam": ["apache-airflow-providers-apache-beam", "apache-beam[gcp]"],
+                "apache.cassandra": ["apache-airflow-providers-apache-cassandra"],
+                "cncf.kubernetes": ["apache-airflow-providers-cncf-kubernetes>=7.2.0b0"],
+                "common.sql": ["apache-airflow-providers-common-sql"],
+                "facebook": ["apache-airflow-providers-facebook>=2.2.0b0"],
+                "leveldb": ["plyvel"],
+                "microsoft.azure": ["apache-airflow-providers-microsoft-azure"],
+                "microsoft.mssql": ["apache-airflow-providers-microsoft-mssql"],
+                "mysql": ["apache-airflow-providers-mysql"],
+                "openlineage": ["apache-airflow-providers-openlineage"],
+                "oracle": ["apache-airflow-providers-oracle>=3.1.0b0"],
+                "postgres": ["apache-airflow-providers-postgres"],
+                "presto": ["apache-airflow-providers-presto"],
+                "salesforce": ["apache-airflow-providers-salesforce"],
+                "sftp": ["apache-airflow-providers-sftp"],
+                "ssh": ["apache-airflow-providers-ssh"],
+                "trino": ["apache-airflow-providers-trino"],
+            },
+            id="With beta0 suffix normalized automatically to b0 (PEP 440)",
+        ),
+    ],
+)
+def test_get_package_extras(version_suffix: str, expected: dict[str, list[str]]):
+    assert get_package_extras("google", version_suffix=version_suffix) == expected
 
 
 def test_get_provider_details():

generated/provider_dependencies.json

@@ -467,7 +467,7 @@
   },
   "fab": {
     "deps": [
-      "apache-airflow>=2.9.0.dev0",
+      "apache-airflow>=2.9.0",
       "flask-appbuilder==4.3.11",
       "flask-login>=0.6.2",
       "flask>=2.2,<2.3",