aoelvp94 commented on code in PR #73015:
URL: https://github.com/apache/airflow/pull/73015#discussion_r4522074357
##########
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:
Right, since the marker now lives in the tree, wiping the tree also removes
the marker and forces a re-stage. Dropped the walk and the helper; the check is
now just the marker matching the HEAD ETag.
https://github.com/aoelvp94/airflow/commit/0427f4476d07d035f6a69f217fb9dea85a3a67fc
--
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]
