aoelvp94 commented on code in PR #73015:
URL: https://github.com/apache/airflow/pull/73015#discussion_r4522099435
##########
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:
Good point. Everything now stays under the bundle's own `base_dir`: the Dag
tree is `base_dir/dags` (what `bundle.path` returns) and staging dirs are
created next to it, so there's no sibling naming anymore and the `with_name`
suggestion no longer applies. Added a test asserting a bundle named `.test`
next to `test` is left untouched.
https://github.com/aoelvp94/airflow/commit/7caa4348297bed4370af1c6baf3154afed425980
--
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]
