Add support for S3 dag bundle

providers/amazon/src/airflow/providers/amazon/aws/bundles/__init__.py

@@ -0,0 +1,16 @@
+# 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.

providers/amazon/src/airflow/providers/amazon/aws/bundles/s3.py

@@ -0,0 +1,152 @@
+# 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 pathlib import Path
+
+import structlog
+
+from airflow.dag_processing.bundles.base import BaseDagBundle
+from airflow.exceptions import AirflowException
+from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
+from airflow.providers.amazon.aws.hooks.s3 import S3Hook
+
+
+class S3DagBundle(BaseDagBundle):
+    """
+    S3 DAG bundle - exposes a directory in S3 as a DAG bundle.
+
+    This allows Airflow to load DAGs directly from an S3 bucket.
+
+    :param aws_conn_id: Airflow connection ID for AWS.  Defaults to AwsBaseHook.default_conn_name.
+    :param bucket_name: The name of the S3 bucket containing the DAG files.
+    :param prefix:  Optional subdirectory within the S3 bucket where the DAGs are stored.
+                    If None, DAGs are assumed to be at the root of the bucket (Optional).
+    """
+
+    supports_versioning = False
+
+    def __init__(
+        self,
+        *,
+        aws_conn_id: str = AwsBaseHook.default_conn_name,
+        bucket_name: str,
+        prefix: str = "",
+        **kwargs,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.aws_conn_id = aws_conn_id
+        self.bucket_name = bucket_name
+        self.prefix = prefix
+        # Local path where S3 DAGs are downloaded
+        self.s3_dags_dir: Path = self.base_dir
+
+        log = structlog.get_logger(__name__)
+        self._log = log.bind(
+            bundle_name=self.name,
+            version=self.version,
+            bucket_name=self.bucket_name,
+            prefix=self.prefix,
+            aws_conn_id=self.aws_conn_id,
+        )
+        self._s3_hook: S3Hook | None = None
+
+    def _initialize(self):
+        with self.lock():
+            if not self.s3_dags_dir.exists():
+                self._log.info("Creating local DAGs directory: %s", self.s3_dags_dir)
+                os.makedirs(self.s3_dags_dir)
+
+            if not self.s3_dags_dir.is_dir():
+                raise AirflowException(f"Local DAGs path: {self.s3_dags_dir} is not a directory.")
+
+            if not self.s3_hook.check_for_bucket(bucket_name=self.bucket_name):
+                raise AirflowException(f"S3 bucket '{self.bucket_name}' does not exist.")
+
+            if self.prefix:
+                # don't check when prefix is ""
+                if not self.s3_hook.check_for_prefix(
+                    bucket_name=self.bucket_name, prefix=self.prefix, delimiter="/"
+                ):
+                    raise AirflowException(
+                        f"S3 prefix 's3://{self.bucket_name}/{self.prefix}' does not exist."
+                    )
+            self.refresh()
+
+    def initialize(self) -> None:
+        self._initialize()
+        super().initialize()
+
+    @property
+    def s3_hook(self):
+        if self._s3_hook is None:
+            try:
+                self._s3_hook: S3Hook = S3Hook(aws_conn_id=self.aws_conn_id)  # Initialize S3 hook.
+            except AirflowException as e:
+                self._log.warning("Could not create S3Hook for connection %s: %s", self.aws_conn_id, e)
+        return self._s3_hook
+
+    def __repr__(self):
+        return (
+            f"<S3DagBundle("
+            f"name={self.name!r}, "
+            f"bucket_name={self.bucket_name!r}, "
+            f"prefix={self.prefix!r}, "
+            f"version={self.version!r}"
+            f")>"
+        )
+
+    def get_current_version(self) -> str | None:
+        """Return the current version of the DAG bundle. Currently not supported."""
+        return None
+
+    @property
+    def path(self) -> Path:
+        """Return the local path to the DAG files."""
+        return self.s3_dags_dir  # Path where DAGs are downloaded.
+
+    def refresh(self) -> None:
+        """Refresh the DAG bundle by re-downloading the DAGs from S3."""
+        if self.version:
+            raise AirflowException("Refreshing a specific version is not supported")
+
+        with self.lock():
+            self._log.debug(
+                "Downloading DAGs from s3://%s/%s to %s", self.bucket_name, self.prefix, self.s3_dags_dir
+            )
+            self.s3_hook.sync_to_local_dir(
+                bucket_name=self.bucket_name,
+                s3_prefix=self.prefix,
+                local_dir=self.s3_dags_dir,
+                delete_stale=True,
+            )
+
+    def view_url(self, version: str | None = None) -> str | None:
+        """Return a URL for viewing the DAGs in S3. Currently, versioning is not supported."""
+        if self.version:
+            raise AirflowException("S3 url with version is not supported")
+
+        # https://<bucket-name>.s3.<region>.amazonaws.com/<object-key>
+        url = f"https://{self.bucket_name}.s3"
+        if self.s3_hook.region_name:
+            url += f".{self.s3_hook.region_name}"
+        url += ".amazonaws.com"
+        if self.prefix:
+            url += f"/{self.prefix}"
+
+        return url

providers/amazon/src/airflow/providers/amazon/aws/hooks/s3.py

@@ -1683,3 +1683,80 @@ def delete_bucket_tagging(self, bucket_name: str | None = None) -> None:
         """
         s3_client = self.get_conn()
         s3_client.delete_bucket_tagging(Bucket=bucket_name)
+
+    def _sync_to_local_dir_delete_stale_local_files(self, current_s3_objects: list[Path], local_dir: Path):
+        current_s3_keys = {key for key in current_s3_objects}
+
+        for item in local_dir.iterdir():
+            item: Path  # type: ignore[no-redef]
+            absolute_item_path = item.resolve()
+
+            if absolute_item_path not in current_s3_keys:
+                try:
+                    if item.is_file():
+                        item.unlink(missing_ok=True)
+                        self.log.debug("Deleted stale local file: %s", item)
+                    elif item.is_dir():
+                        # delete only when the folder is empty
+                        if not os.listdir(item):
+                            item.rmdir()
+                            self.log.debug("Deleted stale empty directory: %s", item)
+                    else:
+                        self.log.debug("Skipping stale item of unknown type: %s", item)
+                except OSError as e:
+                    self.log.error("Error deleting stale item %s: %s", item, e)
+                    raise e
+
+    def _sync_to_local_dir_if_changed(self, s3_bucket, s3_object, local_target_path: Path):
+        should_download = False
+        download_msg = ""
+        if not local_target_path.exists():
+            should_download = True
+            download_msg = f"Local file {local_target_path} does not exist."
+        else:
+            local_stats = local_target_path.stat()
+
+            if s3_object.size != local_stats.st_size:
+                should_download = True
+                download_msg = (
+                    f"S3 object size ({s3_object.size}) and local file size ({local_stats.st_size}) differ."
+                )
+
+            s3_last_modified = s3_object.last_modified
+            if local_stats.st_mtime < s3_last_modified.microsecond:
+                should_download = True
+                download_msg = f"S3 object last modified ({s3_last_modified.microsecond}) and local file last modified ({local_stats.st_mtime}) differ."
+
+        if should_download:
+            s3_bucket.download_file(s3_object.key, local_target_path)
+            self.log.debug(
+                "%s Downloaded %s to %s", download_msg, s3_object.key, local_target_path.as_posix()
+            )
+        else:
+            self.log.debug(
+                "Local file %s is up-to-date with S3 object %s. Skipping download.",
+                local_target_path.as_posix(),
+                s3_object.key,
+            )
+
+    def sync_to_local_dir(self, bucket_name: str, local_dir: Path, s3_prefix="", delete_stale: bool = True):
+        """Download S3 files from the S3 bucket to the local directory."""
+        self.log.debug("Downloading data from s3://%s/%s to %s", bucket_name, s3_prefix, local_dir)
+
+        local_s3_objects = []
+        s3_bucket = self.get_bucket(bucket_name)
+        for obj in s3_bucket.objects.filter(Prefix=s3_prefix):
+            obj_path = Path(obj.key)
+            local_target_path = local_dir.joinpath(obj_path.relative_to(s3_prefix))
+            if not local_target_path.parent.exists():
+                local_target_path.parent.mkdir(parents=True, exist_ok=True)
+                self.log.debug("Created local directory: %s", local_target_path.parent)
+            self._sync_to_local_dir_if_changed(
+                s3_bucket=s3_bucket, s3_object=obj, local_target_path=local_target_path
+            )
+            local_s3_objects.append(local_target_path)
+
+        if delete_stale:
+            self._sync_to_local_dir_delete_stale_local_files(
+                current_s3_objects=local_s3_objects, local_dir=local_dir
+            )

providers/amazon/tests/unit/amazon/aws/bundles/__init__.py

@@ -0,0 +1,16 @@
+# 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.