uranusjr commented on code in PR #73015:
URL: https://github.com/apache/airflow/pull/73015#discussion_r4410474677
##########
providers/amazon/src/airflow/providers/amazon/aws/bundles/s3.py:
##########
@@ -140,6 +193,95 @@ def refresh(self) -> None:
delete_stale=True,
)
+ def _archive_staging_prefix(self) -> str:
+ # The storage root is shared by every bundle. The lock is per bundle
name, so
+ # leftover cleanup must not match another bundle's in-progress staging
directory.
+ return f".{self.name}-archive-staging-"
+
+ def _remove_stale_archive_staging_dirs(self) -> None:
+ storage_root = self.s3_dags_dir.parent
+ if not storage_root.is_dir():
+ return
+ prefix = self._archive_staging_prefix()
+ for leftover in storage_root.iterdir():
+ if leftover.is_dir() and leftover.name.startswith(prefix):
+ shutil.rmtree(leftover, ignore_errors=True)
+
+ def _download_archive(self, client, archive_path: Path, etag: str) -> str:
+ """Download the archive object named by ``etag``."""
+ # download_file HEADs again and GETs whatever is current, so a replace
in between
+ # can store an ETag for a different object than the one unpacked.
+ request = {"Bucket": self.bucket_name, "Key": self.archive_key}
+ if etag:
+ request["IfMatch"] = etag
+ response = client.get_object(**request)
+ body = response["Body"]
+ try:
+ with archive_path.open("wb") as handle:
+ for chunk in body.iter_chunks():
+ handle.write(chunk)
+ finally:
+ body.close()
+ return response.get("ETag", "")
+
+ def _refresh_from_archive(self) -> None:
+ """Stage the Dag bundle by downloading and unpacking the single
archive object."""
+ # A killed process never reaches the finally that removes its staging
directory.
+ self._remove_stale_archive_staging_dirs()
+ client = self.s3_hook.get_conn()
+ head = client.head_object(Bucket=self.bucket_name,
Key=self.archive_key)
+ etag: str = head.get("ETag", "")
+
+ marker = self.s3_dags_dir / self.archive_etag_marker
+ if etag and marker.is_file() and marker.read_text() == etag:
+ self._log.debug(
+ "Dag bundle archive 's3://%s/%s' is unchanged (ETag %s),
skipping staging",
+ self.bucket_name,
+ self.archive_key,
+ etag,
+ )
+ return
+
+ staging_dir = Path(
+ tempfile.mkdtemp(dir=self.s3_dags_dir.parent,
prefix=self._archive_staging_prefix())
+ )
+ try:
+ archive_path = staging_dir / "_bundle_archive"
+ downloaded_etag = self._download_archive(client, archive_path,
etag)
+
+ unpack_dir = staging_dir / "unpacked"
+ unpack_dir.mkdir()
+ _extract_archive(archive_path, unpack_dir)
+ archive_path.unlink()
+
+ marker_etag = downloaded_etag or etag
+ if marker_etag:
+ (unpack_dir / self.archive_etag_marker).write_text(marker_etag)
+
+ # Swap the freshly unpacked tree into place so a partially staged
bundle is
+ # never observable at self.s3_dags_dir. The previous tree is kept
inside the
+ # staging directory: a failed swap restores it, and the staging
cleanup below
+ # removes it once the swap succeeded.
Review Comment:
Not accurate, as Niko already pointed out. Reword or drop.
--
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]
