Repository: apache/airflow
Pull request: Add basic support for git and local dag bundles
This is a first pass at git and local dag bundles, and the bundle interface. These will certainly be expanded on (e.g. to support non-public repos).

Related: #42286
Closes: #42287 
Closes: #42288

Problem statement:
Create Local DAG bundle backend
### Body

We need a local DAG bundle backend, which will be used for backcompat with the DAG folder and for local dev. It will not support versioning.

### Committer

- [x] I acknowledge that I am a maintainer/committer of the Apache Airflow project.

Create git DAG bundle backend
### Body

Create a git DAG bundle backend, with versioning.

### Committer

- [x] I acknowledge that I am a maintainer/committer of the Apache Airflow project.

Hints from issue discussion:

Commit message: Remove `cleanup` method; update docstrings a bit

We will add cleanup back when we implement that functionality.

Patch to review:
diff --git a/airflow/config_templates/config.yml b/airflow/config_templates/config.yml
index cfadcb16fd28a..27f6955006e23 100644
--- a/airflow/config_templates/config.yml
+++ b/airflow/config_templates/config.yml
@@ -28,6 +28,14 @@ core:
       type: string
       example: ~
       default: "{AIRFLOW_HOME}/dags"
+    dag_bundle_storage_path:
+      description: |
+        The folder where Airflow will store DAG bundles locally.
+        By default, this will be ``tempfile.gettempdir()/airflow``. This path must be absolute.
+      version_added: 3.0.0
+      type: string
+      example: "{AIRFLOW_HOME}/dag_bundles"
+      default: ~
     hostname_callable:
       description: |
         Hostname by providing a path to a callable, which will resolve the hostname.
diff --git a/airflow/dag_processing/bundles/base.py b/airflow/dag_processing/bundles/base.py
new file mode 100644
index 0000000000000..34f3dd65c8b4e
--- /dev/null
+++ b/airflow/dag_processing/bundles/base.py
@@ -0,0 +1,82 @@
+# 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 tempfile
+from abc import ABC, abstractmethod
+from pathlib import Path
+
+from airflow.configuration import conf
+
+
+class BaseDagBundle(ABC):
+    """
+    Base class for DAG bundles.
+
+    DAG bundles are used both by the DAG processor and by a worker when running a task. These usage
+    patterns are different, however.
+
+    When running a task, we know what version of the bundle we need (assuming the bundle supports versioning).
+    And we likely only need to keep this specific bundle version around for as long as we have tasks running using
+    that bundle version. This also means, that on a single worker, it's possible that multiple versions of the same
+    bundle are used at the same time.
+
+    In contrast, the DAG processor uses a bundle to keep the DAGs from that bundle up to date. There will not be
+    multiple versions of the same bundle in use at the same time. The DAG processor will always use the latest version.
+
+    :param name: String identifier for the DAG bundle
+    :param version: Version of the DAG bundle (Optional)
+    """
+
+    supports_versioning: bool = False
+
+    def __init__(self, *, name: str, version: str | None = None) -> None:
+        self.name = name
+        self.version = version
+
+    @property
+    def _dag_bundle_root_storage_path(self) -> Path:
+        """
+        Where bundles can store DAGs on disk (if local disk is required).
+
+        This is the root path, shared by various bundles. Each bundle should have its own subdirectory.
+        """
+        if configured_location := conf.get("core", "dag_bundle_storage_path"):
+            return Path(configured_location)
+        return Path(tempfile.gettempdir(), "airflow", "dag_bundles")
+
+    @property
+    @abstractmethod
+    def path(self) -> Path:
+        """
+        Path for this bundle.
+
+        Airflow will use this path to load/execute the DAGs from the bundle.
+        """
+
+    @abstractmethod
+    def get_current_version(self) -> str:
+        """
+        Retrieve a string that represents the version of the DAG bundle.
+
+        Airflow can use this value to retrieve this same bundle version later.
+        """
+
+    @abstractmethod
+    def refresh(self) -> None:
+        """Retrieve the latest version of the files in the bundle."""
diff --git a/airflow/dag_processing/bundles/git.py b/airflow/dag_processing/bundles/git.py
new file mode 100644
index 0000000000000..1f00461b03ce4
--- /dev/null
+++ b/airflow/dag_processing/bundles/git.py
@@ -0,0 +1,113 @@
+# 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 os
+from typing import TYPE_CHECKING
+
+from git import Repo
+
+from airflow.dag_processing.bundles.base import BaseDagBundle
+from airflow.exceptions import AirflowException
+
+if TYPE_CHECKING:
+    from pathlib import Path
+
+
+class GitDagBundle(BaseDagBundle):
+    """
+    git DAG bundle - exposes a git repository as a DAG bundle.
+
+    Instead of cloning the repository every time, we clone the repository once into a bare repo from the source
+    and then do a clone for each version from there.
+
+    :param repo_url: URL of the git repository
+    :param head: Branch or tag for this DAG bundle
+    """
+
+    supports_versioning = True
+
+    def __init__(self, *, repo_url: str, head: str, **kwargs) -> None:
+        super().__init__(**kwargs)
+        self.repo_url = repo_url
+        self.head = head
+
+        self.bare_repo_path = self._dag_bundle_root_storage_path / "git" / self.name
+        self._ensure_bare_repo_exists()
+        self._ensure_version_in_bare_repo()
+        self._ensure_repo_exists()
+
+        if self.version:
+            if not self._has_version(self.repo, self.version):
+                self.repo.remotes.origin.fetch()
+
+            self.repo.head.set_reference(self.repo.commit(self.version))
+            self.repo.head.reset(index=True, working_tree=True)
+        else:
+            self.refresh()
+
+    def _ensure_repo_exists(self) -> None:
+        if not os.path.exists(self.path):
+            Repo.clone_from(
+                url=self.bare_repo_path,
+                to_path=self.path,
+            )
+        self.repo = Repo(self.path)
+
+    def _ensure_bare_repo_exists(self) -> None:
+        if not os.path.exists(self.bare_repo_path):
+            Repo.clone_from(
+                url=self.repo_url,
+                to_path=self.bare_repo_path,
+                bare=True,
+            )
+        self.bare_repo = Repo(self.bare_repo_path)
+
+    def _ensure_version_in_bare_repo(self) -> None:
+        if not self.version:
+            return
+        if not self._has_version(self.bare_repo, self.version):
+            self.bare_repo.remotes.origin.fetch("+refs/heads/*:refs/heads/*")
+            if not self._has_version(self.bare_repo, self.version):
+                raise AirflowException(f"Version {self.version} not found in the repository")
+
+    def __hash__(self) -> int:
+        return hash((self.name, self.get_current_version()))
+
+    def get_current_version(self) -> str:
+        return self.repo.head.commit.hexsha
+
+    @property
+    def path(self) -> Path:
+        location = self.version or self.head
+        return self._dag_bundle_root_storage_path / "git" / f"{self.name}+{location}"
+
+    @staticmethod
+    def _has_version(repo: Repo, version: str) -> bool:
+        try:
+            repo.commit(version)
+            return True
+        except Exception:
+            return False
+
+    def refresh(self) -> None:
+        if self.version:
+            raise AirflowException("Refreshing a specific version is not supported")
+
+        self.bare_repo.remotes.origin.fetch("+refs/heads/*:refs/heads/*")
+        self.repo.remotes.origin.pull()
diff --git a/airflow/dag_processing/bundles/local.py b/airflow/dag_processing/bundles/local.py
new file mode 100644
index 0000000000000..dcbac96911159
--- /dev/null
+++ b/airflow/dag_processing/bundles/local.py
@@ -0,0 +1,47 @@
+# 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
+
+from pathlib import Path
+
+from airflow.dag_processing.bundles.base import BaseDagBundle
+from airflow.exceptions import AirflowException
+
+
+class LocalDagBundle(BaseDagBundle):
+    """
+    Local DAG bundle - exposes a local directory as a DAG bundle.
+
+    :param local_folder: Local folder where the DAGs are stored
+    """
+
+    supports_versioning = False
+
+    def __init__(self, *, local_folder: str, **kwargs) -> None:
+        super().__init__(**kwargs)
+        self._path = Path(local_folder)
+
+    def get_current_version(self) -> str:
+        raise AirflowException("Not versioned!")
+
+    def refresh(self) -> None:
+        """Nothing to refresh - it's just a local directory."""
+
+    @property
+    def path(self) -> Path:
+        return self._path
diff --git a/hatch_build.py b/hatch_build.py
index 4c0e7f377112e..f9a0e9447b6c4 100644
--- a/hatch_build.py
+++ b/hatch_build.py
@@ -190,7 +190,6 @@
     ],
     "devel-devscripts": [
         "click>=8.0",
-        "gitpython>=3.1.40",
         "hatch>=1.9.1",
         # Incremental 24.7.0, 24.7.1 has broken `python -m virtualenv` command when run in /opt/airflow directory
         "incremental!=24.7.0,!=24.7.1,>=22.10.0",
@@ -385,6 +384,7 @@
     # We should remove the limitation after 2.3 is released and our dependencies are updated to handle it
     "flask>=2.2.1,<2.3",
     "fsspec>=2023.10.0",
+    "gitpython>=3.1.40",
     'google-re2>=1.0;python_version<"3.12"',
     'google-re2>=1.1;python_version>="3.12"',
     "gunicorn>=20.1.0",
diff --git a/tests/dag_processing/test_dag_bundles.py b/tests/dag_processing/test_dag_bundles.py
new file mode 100644
index 0000000000000..5ccee4ab4269b
--- /dev/null
+++ b/tests/dag_processing/test_dag_bundles.py
@@ -0,0 +1,165 @@
+# 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 tempfile
+from pathlib import Path
+
+import pytest
+from git import Repo
+
+from airflow.dag_processing.bundles.git import GitDagBundle
+from airflow.dag_processing.bundles.local import LocalDagBundle
+from airflow.exceptions import AirflowException
+
+from tests_common.test_utils.config import conf_vars
+
+
+@pytest.fixture(autouse=True)
+def bundle_temp_dir(tmp_path):
+    with conf_vars({("core", "dag_bundle_storage_path"): str(tmp_path)}):
+        yield tmp_path
+
+
+def test_default_dag_storage_path():
+    with conf_vars({("core", "dag_bundle_storage_path"): ""}):
+        bundle = LocalDagBundle(name="test", local_folder="/hello")
+        assert bundle._dag_bundle_root_storage_path == Path(tempfile.gettempdir(), "airflow", "dag_bundles")
+
+
+class TestLocalDagBundle:
+    def test_path(self):
+        bundle = LocalDagBundle(name="test", local_folder="/hello")
+        assert bundle.path == Path("/hello")
+
+    def test_does_not_support_versioning(self):
+        assert LocalDagBundle.supports_versioning is False
+
+        bundle = LocalDagBundle(name="test", local_folder="/hello")
+
+        with pytest.raises(AirflowException):
+            bundle.get_current_version()
+
+
+@pytest.fixture
+def git_repo(tmp_path_factory):
+    directory = tmp_path_factory.mktemp("repo")
+    repo = Repo.init(directory)
+    file_path = directory / "test_dag.py"
+    with open(file_path, "w") as f:
+        f.write("hello world")
+    repo.index.add([file_path])
+    repo.index.commit("Initial commit")
+    return (directory, repo)
+
+
+class TestGitDagBundle:
+    def test_supports_versioning(self):
+        assert GitDagBundle.supports_versioning is True
+
+    def test_uses_dag_bundle_root_storage_path(self, git_repo):
+        repo_path, repo = git_repo
+        bundle = GitDagBundle(name="test", repo_url=repo_path, head="master")
+        assert str(bundle._dag_bundle_root_storage_path) in str(bundle.path)
+
+    def test_get_current_version(self, git_repo):
+        repo_path, repo = git_repo
+        bundle = GitDagBundle(name="test", repo_url=repo_path, head="master")
+
+        assert bundle.get_current_version() == repo.head.commit.hexsha
+
+    def test_get_specific_version(self, git_repo):
+        repo_path, repo = git_repo
+        starting_commit = repo.head.commit
+
+        # Add new file to the repo
+        file_path = repo_path / "new_test.py"
+        with open(file_path, "w") as f:
+            f.write("hello world")
+        repo.index.add([file_path])
+        repo.index.commit("Another commit")
+
+        bundle = GitDagBundle(name="test", version=starting_commit.hexsha, repo_url=repo_path, head="master")
+
+        assert bundle.get_current_version() == starting_commit.hexsha
+
+        files_in_repo = {f.name for f in bundle.path.iterdir() if f.is_file()}
+        assert {"test_dag.py"} == files_in_repo
+
+    def test_get_tag_version(self, git_repo):
+        repo_path, repo = git_repo
+        starting_commit = repo.head.commit
+
+        # add tag
+        repo.create_tag("test")
+        print(repo.tags)
+
+        # Add new file to the repo
+        file_path = repo_path / "new_test.py"
+        with open(file_path, "w") as f:
+            f.write("hello world")
+        repo.index.add([file_path])
+        repo.index.commit("Another commit")
+
+        bundle = GitDagBundle(name="test", version="test", repo_url=repo_path, head="master")
+
+        assert bundle.get_current_version() == starting_commit.hexsha
+
+        files_in_repo = {f.name for f in bundle.path.iterdir() if f.is_file()}
+        assert {"test_dag.py"} == files_in_repo
+
+    def test_get_latest(self, git_repo):
+        repo_path, repo = git_repo
+        starting_commit = repo.head.commit
+
+        file_path = repo_path / "new_test.py"
+        with open(file_path, "w") as f:
+            f.write("hello world")
+        repo.index.add([file_path])
+        repo.index.commit("Another commit")
+
+        bundle = GitDagBundle(name="test", repo_url=repo_path, head="master")
+
+        assert bundle.get_current_version() != starting_commit.hexsha
+
+        files_in_repo = {f.name for f in bundle.path.iterdir() if f.is_file()}
+        assert {"test_dag.py", "new_test.py"} == files_in_repo
+
+    def test_refresh(self, git_repo):
+        repo_path, repo = git_repo
+        starting_commit = repo.head.commit
+
+        bundle = GitDagBundle(name="test", repo_url=repo_path, head="master")
+
+        assert bundle.get_current_version() == starting_commit.hexsha
+
+        files_in_repo = {f.name for f in bundle.path.iterdir() if f.is_file()}
+        assert {"test_dag.py"} == files_in_repo
+
+        file_path = repo_path / "new_test.py"
+        with open(file_path, "w") as f:
+            f.write("hello world")
+        repo.index.add([file_path])
+        commit = repo.index.commit("Another commit")
+
+        bundle.refresh()
+
+        assert bundle.get_current_version() == commit.hexsha
+
+        files_in_repo = {f.name for f in bundle.path.iterdir() if f.is_file()}
+        assert {"test_dag.py", "new_test.py"} == files_in_repo


Code review comments:
diff_hunk: @@ -28,6 +28,14 @@ core:
       type: string
       example: ~
       default: "{AIRFLOW_HOME}/dags"
+    dag_bundle_storage_path:
+      description: |
+        The folder where Airflow will store DAG bundles locally.
+        By default, this will be ``tempfile.gettempdir()/airflow``. This path must be absolute.
path: airflow/config_templates/config.yml
text: @user1:
i think this needs to be updated. logic looks different.

@author:
Updated.

diff_hunk: @@ -0,0 +1,113 @@
+# 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 os
+from typing import TYPE_CHECKING
+
+from git import Repo
+
+from airflow.dag_processing.bundles.base import BaseDagBundle
+from airflow.exceptions import AirflowException
+
+if TYPE_CHECKING:
+    from pathlib import Path
+
+
+class GitDagBundle(BaseDagBundle):
+    """
+    git DAG bundle - exposes a git repository as a DAG bundle.
+
+    Instead of cloning the repository every time, we clone the repository once into a bare repo from the source
+    and then do a clone for each version from there.
+
+    :param repo_url: URL of the git repository
+    :param head: Branch or tag for this DAG bundle
+    """
+
+    supports_versioning = True
+
+    def __init__(self, *, repo_url: str, head: str, **kwargs) -> None:
+        super().__init__(**kwargs)
+        self.repo_url = repo_url
+        self.head = head
+
+        self.bare_repo_path = self._dag_bundle_root_storage_path / "git" / self.name
+        self._ensure_bare_repo_exists()
+        self._ensure_version_in_bare_repo()
+        self._ensure_repo_exists()
+
+        if self.version:
+            if not self._has_version(self.repo, self.version):
+                self.repo.remotes.origin.fetch()
+
+            self.repo.head.set_reference(self.repo.commit(self.version))
+            self.repo.head.reset(index=True, working_tree=True)
+        else:
+            self.refresh()
+
+    def _ensure_repo_exists(self) -> None:
path: airflow/dag_processing/bundles/git.py
text: @user1:
nit: `_ensure_repo_exists` sounds more like a validation function than a function that actually instantiates the `repo` object. same with `_ensure_bare_repo_exists` below

@author:
Renamed them.