uranusjr commented on code in PR #73015:
URL: https://github.com/apache/airflow/pull/73015#discussion_r4509965705
##########
providers/amazon/src/airflow/providers/amazon/aws/bundles/s3.py:
##########
@@ -163,7 +143,221 @@ def view_url_template(self) -> str | None:
if self.s3_hook.region_name:
url += f".{self.s3_hook.region_name}"
url += ".amazonaws.com"
- if self.prefix:
- url += f"/{self.prefix}"
+ key = self._view_url_key()
+ if key:
+ url += f"/{key}"
return url
+
+ def _reject_version_refresh(self) -> None:
+ if self.version:
+ raise AirflowException("Refreshing a specific version is not
supported")
+
+ def refresh(self) -> None:
+ """Refresh the Dag bundle by updating the local Dag files from S3."""
+ self._reject_version_refresh()
+ with self.lock():
+ self._refresh_locked()
+
+ def _refresh_locked(self) -> None:
+ """Stage the bundle while the bundle lock is held."""
+ raise NotImplementedError
+
+
+class S3DagBundle(_S3DagBundleBase):
+ """
+ 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.
+ A trailing slash is optional.
+ 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:
+ self.prefix = prefix
+ super().__init__(aws_conn_id=aws_conn_id, bucket_name=bucket_name,
**kwargs)
+
+ def _get_logger(self):
+ """Return a logger bound with this prefix bundle's context."""
+ return structlog.get_logger(__name__).bind(**self._log_bindings(),
prefix=self.prefix)
+
+ def _view_url_key(self) -> str:
+ return self.prefix
+
+ def _validate_remote(self) -> None:
+ 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.")
+
+ 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 _refresh_locked(self) -> None:
+ self._log.debug(
+ "Downloading Dags from s3://%s/%s to %s", self.bucket_name,
self.prefix, self.s3_dags_dir
+ )
+ sync_prefix = self.prefix
+ if sync_prefix and not sync_prefix.endswith("/"):
+ sync_prefix += "/"
+ self.s3_hook.sync_to_local_dir(
+ bucket_name=self.bucket_name,
+ s3_prefix=sync_prefix,
+ local_dir=self.s3_dags_dir,
+ delete_stale=True,
+ )
+
+
+class S3ArchiveDagBundle(_S3DagBundleBase):
+ """
+ S3 Dag bundle staged from a single ``.tar.gz`` object.
+
+ The archive is the only source of the Dags. Members are extracted relative
to the
+ bundle root, for example ``tar -C <dags_dir> -czf dags.tar.gz .``.
Requires Python
+ 3.11.4 or later.
+
+ :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 archive.
+ :param archive_key: S3 key of the ``.tar.gz`` archive holding the Dag
files.
+ """
+
+ supports_versioning = False
+
+ def __init__(
+ self,
+ *,
+ aws_conn_id: str = AwsBaseHook.default_conn_name,
+ bucket_name: str,
+ archive_key: str,
+ **kwargs,
+ ) -> None:
+ _require_tar_data_filter()
+ self.archive_key = archive_key
+ super().__init__(aws_conn_id=aws_conn_id, bucket_name=bucket_name,
**kwargs)
+
+ def _get_logger(self):
+ """Return a logger bound with this archive bundle's context."""
+ return structlog.get_logger(__name__).bind(**self._log_bindings(),
archive_key=self.archive_key)
+
+ def _view_url_key(self) -> str:
+ return self.archive_key
+
+ def __repr__(self):
+ return (
+ f"<S3ArchiveDagBundle("
+ f"name={self.name!r}, "
+ f"bucket_name={self.bucket_name!r}, "
+ f"archive_key={self.archive_key!r}, "
+ f"version={self.version!r}"
+ f")>"
+ )
+
+ def _get_archive_state_dir(self) -> Path:
+ return self.s3_dags_dir.parent / f".{self.name}"
+
+ def _get_archive_etag_path(self) -> Path:
+ """Return the ETag marker path inside the staged Dag tree."""
+ return self.s3_dags_dir / ".archive_etag"
+
+ def _bundle_dir_has_files(self) -> bool:
+ """Return whether the staged tree has files other than the ETag
marker."""
+ marker = self._get_archive_etag_path()
+ return self.s3_dags_dir.is_dir() and any(
+ path.is_file() and path != marker for path in
self.s3_dags_dir.rglob("*")
Review Comment:
I believe the `rglob("*")` is no longer needed now that the marker lives in
the tree.
##########
providers/amazon/src/airflow/providers/amazon/aws/bundles/s3.py:
##########
@@ -163,7 +143,221 @@ def view_url_template(self) -> str | None:
if self.s3_hook.region_name:
url += f".{self.s3_hook.region_name}"
url += ".amazonaws.com"
- if self.prefix:
- url += f"/{self.prefix}"
+ key = self._view_url_key()
+ if key:
+ url += f"/{key}"
return url
+
+ def _reject_version_refresh(self) -> None:
+ if self.version:
+ raise AirflowException("Refreshing a specific version is not
supported")
+
+ def refresh(self) -> None:
+ """Refresh the Dag bundle by updating the local Dag files from S3."""
+ self._reject_version_refresh()
+ with self.lock():
+ self._refresh_locked()
+
+ def _refresh_locked(self) -> None:
+ """Stage the bundle while the bundle lock is held."""
+ raise NotImplementedError
Review Comment:
Since this is an ABC, `@abstractmethod` handles this better. (existing
pattern can be found in the Airflow code base elsewhere)
##########
providers/amazon/src/airflow/providers/amazon/aws/bundles/s3.py:
##########
@@ -27,46 +30,59 @@
from airflow.providers.common.compat.sdk import AirflowException
-class S3DagBundle(BaseDagBundle):
- """
- S3 Dag bundle - exposes a directory in S3 as a Dag bundle.
+def _require_tar_data_filter() -> None:
+ if hasattr(tarfile, "data_filter"):
+ return
+ raise RuntimeError("Archive Dag bundles require Python 3.11.4 or later
(tarfile filter='data').")
- 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.
- A trailing slash is optional.
- If None, Dags are assumed to be at the root of the bucket
(Optional).
+def _extract_archive(archive_path: Path, unpack_dir: Path) -> None:
+ """Unpack a bundle archive without writing outside ``unpack_dir``."""
+ with tarfile.open(archive_path, "r:*") as tar:
+ tar.extractall(unpack_dir, filter="data")
+
+
+class _S3DagBundleBase(BaseDagBundle):
"""
+ S3 connection, local directory, and console URL shared by the S3 bundle
classes.
- supports_versioning = False
+ :meta private:
+ """
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._log = self._get_logger()
self._s3_hook: S3Hook | None = None
+ def _log_bindings(self) -> dict[str, str | None]:
+ """Return log context shared by every S3 Dag bundle."""
+ return {
+ "bundle_name": self.name,
+ "version": self.version,
+ "bucket_name": self.bucket_name,
+ "aws_conn_id": self.aws_conn_id,
+ }
+
+ def _get_logger(self):
Review Comment:
This should have a return annotation (same for subclasses)
##########
providers/amazon/src/airflow/providers/amazon/aws/bundles/s3.py:
##########
@@ -163,7 +143,221 @@ def view_url_template(self) -> str | None:
if self.s3_hook.region_name:
url += f".{self.s3_hook.region_name}"
url += ".amazonaws.com"
- if self.prefix:
- url += f"/{self.prefix}"
+ key = self._view_url_key()
+ if key:
+ url += f"/{key}"
return url
+
+ def _reject_version_refresh(self) -> None:
Review Comment:
This function can now be inlined since it’s only called in one place (and
the logic is very trivial)
##########
providers/amazon/src/airflow/providers/amazon/aws/bundles/s3.py:
##########
@@ -163,7 +143,221 @@ def view_url_template(self) -> str | None:
if self.s3_hook.region_name:
url += f".{self.s3_hook.region_name}"
url += ".amazonaws.com"
- if self.prefix:
- url += f"/{self.prefix}"
+ key = self._view_url_key()
+ if key:
+ url += f"/{key}"
return url
+
+ def _reject_version_refresh(self) -> None:
+ if self.version:
+ raise AirflowException("Refreshing a specific version is not
supported")
+
+ def refresh(self) -> None:
+ """Refresh the Dag bundle by updating the local Dag files from S3."""
+ self._reject_version_refresh()
+ with self.lock():
+ self._refresh_locked()
+
+ def _refresh_locked(self) -> None:
+ """Stage the bundle while the bundle lock is held."""
+ raise NotImplementedError
+
+
+class S3DagBundle(_S3DagBundleBase):
+ """
+ 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.
+ A trailing slash is optional.
+ 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:
+ self.prefix = prefix
+ super().__init__(aws_conn_id=aws_conn_id, bucket_name=bucket_name,
**kwargs)
+
+ def _get_logger(self):
+ """Return a logger bound with this prefix bundle's context."""
+ return structlog.get_logger(__name__).bind(**self._log_bindings(),
prefix=self.prefix)
+
+ def _view_url_key(self) -> str:
+ return self.prefix
+
+ def _validate_remote(self) -> None:
+ 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.")
+
+ 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 _refresh_locked(self) -> None:
+ self._log.debug(
+ "Downloading Dags from s3://%s/%s to %s", self.bucket_name,
self.prefix, self.s3_dags_dir
+ )
+ sync_prefix = self.prefix
+ if sync_prefix and not sync_prefix.endswith("/"):
+ sync_prefix += "/"
+ self.s3_hook.sync_to_local_dir(
+ bucket_name=self.bucket_name,
+ s3_prefix=sync_prefix,
+ local_dir=self.s3_dags_dir,
+ delete_stale=True,
+ )
+
+
+class S3ArchiveDagBundle(_S3DagBundleBase):
+ """
+ S3 Dag bundle staged from a single ``.tar.gz`` object.
+
+ The archive is the only source of the Dags. Members are extracted relative
to the
+ bundle root, for example ``tar -C <dags_dir> -czf dags.tar.gz .``.
Requires Python
+ 3.11.4 or later.
+
+ :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 archive.
+ :param archive_key: S3 key of the ``.tar.gz`` archive holding the Dag
files.
+ """
+
+ supports_versioning = False
+
+ def __init__(
+ self,
+ *,
+ aws_conn_id: str = AwsBaseHook.default_conn_name,
+ bucket_name: str,
+ archive_key: str,
+ **kwargs,
+ ) -> None:
+ _require_tar_data_filter()
+ self.archive_key = archive_key
+ super().__init__(aws_conn_id=aws_conn_id, bucket_name=bucket_name,
**kwargs)
+
+ def _get_logger(self):
+ """Return a logger bound with this archive bundle's context."""
+ return structlog.get_logger(__name__).bind(**self._log_bindings(),
archive_key=self.archive_key)
+
+ def _view_url_key(self) -> str:
+ return self.archive_key
+
+ def __repr__(self):
+ return (
+ f"<S3ArchiveDagBundle("
+ f"name={self.name!r}, "
+ f"bucket_name={self.bucket_name!r}, "
+ f"archive_key={self.archive_key!r}, "
+ f"version={self.version!r}"
+ f")>"
+ )
+
+ def _get_archive_state_dir(self) -> Path:
+ return self.s3_dags_dir.parent / f".{self.name}"
Review Comment:
This would still collide e.g. there are two bundles, one named `foo` and the
other `.foo`. Core puts no limit on bundle names, so no prefix character is
safe, even nesting isn’t either. However, GitDagBundle already nests under
`base_dir`, so it’s probably the safest in practice. Let’s keep everything
under the bundle's own folder instead.
Also, no longer relevant once we switch to nesting, but `path.parent / name`
has a shorthand `path.with_name(name)`.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
