uranusjr commented on code in PR #73015:
URL: https://github.com/apache/airflow/pull/73015#discussion_r4534144214
##########
providers/amazon/tests/unit/amazon/aws/bundles/test_s3.py:
##########
@@ -262,3 +282,296 @@ def test_refresh_uses_directory_prefix(self, s3_client,
prefix, extra_suffix):
assert repr(bundle) == original_repr
assert bundle.view_url_template() == original_url
assert s3_client.get_object(Bucket=S3_BUCKET_NAME,
Key=extra_key)["Body"].read() == b"outside"
+
+
[email protected](not airflow.version.version.strip().startswith("3"),
reason="Airflow >=3.0.0 test")
+class TestS3ArchiveDagBundle:
+ def _archive_bundle(self) -> S3ArchiveDagBundle:
+ return S3ArchiveDagBundle(
+ name="test",
+ aws_conn_id=AWS_CONN_ID_WITH_REGION,
+ archive_key=S3_ARCHIVE_KEY,
+ bucket_name=S3_BUCKET_NAME,
+ )
+
+ def test_view_url_template_points_at_archive(self):
+ bundle = self._archive_bundle()
+ assert (
+ bundle.view_url_template()
+ ==
f"https://{S3_BUCKET_NAME}.s3.{AWS_CONN_ID_REGION}.amazonaws.com/{S3_ARCHIVE_KEY}";
+ )
+
+ def test_supports_versioning(self):
+ bundle = self._archive_bundle()
+ assert S3ArchiveDagBundle.supports_versioning is False
+ bundle.version = "test_version"
+ with pytest.raises(AirflowException, match="Refreshing a specific
version is not supported"):
+ bundle.refresh()
+ with pytest.raises(AirflowException, match="S3 url with version is not
supported"):
+ bundle.view_url("test_version")
+
+ def test_refresh_from_archive(self, s3_bucket, s3_client):
+ archive = _make_archive({"dag_01.py": b"test data",
"subproject1/dag_a.py": b"test data"})
+ s3_client.put_object(Bucket=S3_BUCKET_NAME, Key=S3_ARCHIVE_KEY,
Body=archive)
+
+ bundle = self._archive_bundle()
+ bundle.initialize()
+
+ assert bundle.path == bundle.base_dir / "dags"
+ assert (bundle.path / "dag_01.py").is_file()
+ assert (bundle.path / "subproject1" / "dag_a.py").is_file()
+ assert bundle._get_archive_etag_path() == bundle.path / ".archive_etag"
+ assert bundle._get_archive_etag_path().is_file()
+ assert sorted(
+ path.relative_to(bundle.path).as_posix()
+ for path in bundle.path.rglob("*")
+ if path != bundle._get_archive_etag_path()
+ ) == [
+ "dag_01.py",
+ "subproject1",
+ "subproject1/dag_a.py",
+ ]
+ assert not (bundle.path / S3_BUCKET_PREFIX / "dag_02.py").exists()
+ assert not (bundle.path / "dag_02.py").exists()
+
+ def test_refresh_from_archive_skips_staging_when_etag_unchanged(self,
s3_bucket, s3_client):
+ archive = _make_archive({"dag_01.py": b"test data"})
+ s3_client.put_object(Bucket=S3_BUCKET_NAME, Key=S3_ARCHIVE_KEY,
Body=archive)
+
+ bundle = self._archive_bundle()
+ bundle.initialize()
+
+ sentinel = bundle.path / "sentinel.txt"
+ sentinel.write_text("still here")
+ bundle.refresh()
+ assert sentinel.is_file()
+
+ updated = _make_archive({"dag_new.py": b"test data"})
+ s3_client.put_object(Bucket=S3_BUCKET_NAME, Key=S3_ARCHIVE_KEY,
Body=updated)
+ bundle.refresh()
+ assert (bundle.path / "dag_new.py").is_file()
+ assert not (bundle.path / "dag_01.py").exists()
+ assert not sentinel.exists()
+
+ def
test_initialize_restages_when_bundle_dir_cleared_and_etag_matches(self,
s3_bucket, s3_client):
+ archive = _make_archive({"dag_01.py": b"test data"})
+ s3_client.put_object(Bucket=S3_BUCKET_NAME, Key=S3_ARCHIVE_KEY,
Body=archive)
+
+ bundle = self._archive_bundle()
+ bundle.initialize()
+ marker = bundle._get_archive_etag_path()
+ etag = marker.read_text()
+ assert marker == bundle.path / ".archive_etag"
+ assert (bundle.path / "dag_01.py").read_bytes() == b"test data"
+
+ shutil.rmtree(bundle.path)
+
+ client = bundle.s3_hook.get_conn()
+ recorded: list[dict] = []
+ original = client.get_object
+
+ def record_get_object(**kwargs):
+ recorded.append(kwargs)
+ return original(**kwargs)
+
+ with mock.patch.object(client, "get_object", record_get_object):
+ bundle.initialize()
+
+ assert recorded
+ assert (bundle.path / "dag_01.py").read_bytes() == b"test data"
+ assert bundle._get_archive_etag_path().read_text() == etag
+
+ def test_refresh_restages_when_bundle_dir_emptied(self, s3_bucket,
s3_client):
+ archive = _make_archive({"dag_01.py": b"test data"})
+ s3_client.put_object(Bucket=S3_BUCKET_NAME, Key=S3_ARCHIVE_KEY,
Body=archive)
+
+ bundle = self._archive_bundle()
+ bundle.initialize()
+ for child in bundle.path.iterdir():
+ child.unlink()
+ assert bundle.path.is_dir()
+ assert not bundle._get_archive_etag_path().exists()
Review Comment:
These are not needed; they only assert the setup in this test.
##########
providers/amazon/tests/unit/amazon/aws/bundles/test_s3.py:
##########
@@ -262,3 +282,296 @@ def test_refresh_uses_directory_prefix(self, s3_client,
prefix, extra_suffix):
assert repr(bundle) == original_repr
assert bundle.view_url_template() == original_url
assert s3_client.get_object(Bucket=S3_BUCKET_NAME,
Key=extra_key)["Body"].read() == b"outside"
+
+
[email protected](not airflow.version.version.strip().startswith("3"),
reason="Airflow >=3.0.0 test")
+class TestS3ArchiveDagBundle:
+ def _archive_bundle(self) -> S3ArchiveDagBundle:
+ return S3ArchiveDagBundle(
+ name="test",
+ aws_conn_id=AWS_CONN_ID_WITH_REGION,
+ archive_key=S3_ARCHIVE_KEY,
+ bucket_name=S3_BUCKET_NAME,
+ )
+
+ def test_view_url_template_points_at_archive(self):
+ bundle = self._archive_bundle()
+ assert (
+ bundle.view_url_template()
+ ==
f"https://{S3_BUCKET_NAME}.s3.{AWS_CONN_ID_REGION}.amazonaws.com/{S3_ARCHIVE_KEY}";
+ )
+
+ def test_supports_versioning(self):
+ bundle = self._archive_bundle()
+ assert S3ArchiveDagBundle.supports_versioning is False
+ bundle.version = "test_version"
+ with pytest.raises(AirflowException, match="Refreshing a specific
version is not supported"):
+ bundle.refresh()
+ with pytest.raises(AirflowException, match="S3 url with version is not
supported"):
+ bundle.view_url("test_version")
+
+ def test_refresh_from_archive(self, s3_bucket, s3_client):
+ archive = _make_archive({"dag_01.py": b"test data",
"subproject1/dag_a.py": b"test data"})
+ s3_client.put_object(Bucket=S3_BUCKET_NAME, Key=S3_ARCHIVE_KEY,
Body=archive)
+
+ bundle = self._archive_bundle()
+ bundle.initialize()
+
+ assert bundle.path == bundle.base_dir / "dags"
+ assert (bundle.path / "dag_01.py").is_file()
+ assert (bundle.path / "subproject1" / "dag_a.py").is_file()
+ assert bundle._get_archive_etag_path() == bundle.path / ".archive_etag"
+ assert bundle._get_archive_etag_path().is_file()
+ assert sorted(
+ path.relative_to(bundle.path).as_posix()
+ for path in bundle.path.rglob("*")
+ if path != bundle._get_archive_etag_path()
+ ) == [
+ "dag_01.py",
+ "subproject1",
+ "subproject1/dag_a.py",
+ ]
+ assert not (bundle.path / S3_BUCKET_PREFIX / "dag_02.py").exists()
+ assert not (bundle.path / "dag_02.py").exists()
+
+ def test_refresh_from_archive_skips_staging_when_etag_unchanged(self,
s3_bucket, s3_client):
+ archive = _make_archive({"dag_01.py": b"test data"})
+ s3_client.put_object(Bucket=S3_BUCKET_NAME, Key=S3_ARCHIVE_KEY,
Body=archive)
+
+ bundle = self._archive_bundle()
+ bundle.initialize()
+
+ sentinel = bundle.path / "sentinel.txt"
+ sentinel.write_text("still here")
+ bundle.refresh()
+ assert sentinel.is_file()
+
+ updated = _make_archive({"dag_new.py": b"test data"})
+ s3_client.put_object(Bucket=S3_BUCKET_NAME, Key=S3_ARCHIVE_KEY,
Body=updated)
+ bundle.refresh()
+ assert (bundle.path / "dag_new.py").is_file()
+ assert not (bundle.path / "dag_01.py").exists()
+ assert not sentinel.exists()
+
+ def
test_initialize_restages_when_bundle_dir_cleared_and_etag_matches(self,
s3_bucket, s3_client):
+ archive = _make_archive({"dag_01.py": b"test data"})
+ s3_client.put_object(Bucket=S3_BUCKET_NAME, Key=S3_ARCHIVE_KEY,
Body=archive)
+
+ bundle = self._archive_bundle()
+ bundle.initialize()
+ marker = bundle._get_archive_etag_path()
+ etag = marker.read_text()
+ assert marker == bundle.path / ".archive_etag"
+ assert (bundle.path / "dag_01.py").read_bytes() == b"test data"
+
+ shutil.rmtree(bundle.path)
+
+ client = bundle.s3_hook.get_conn()
+ recorded: list[dict] = []
+ original = client.get_object
+
+ def record_get_object(**kwargs):
+ recorded.append(kwargs)
+ return original(**kwargs)
+
+ with mock.patch.object(client, "get_object", record_get_object):
+ bundle.initialize()
+
+ assert recorded
+ assert (bundle.path / "dag_01.py").read_bytes() == b"test data"
+ assert bundle._get_archive_etag_path().read_text() == etag
+
+ def test_refresh_restages_when_bundle_dir_emptied(self, s3_bucket,
s3_client):
+ archive = _make_archive({"dag_01.py": b"test data"})
+ s3_client.put_object(Bucket=S3_BUCKET_NAME, Key=S3_ARCHIVE_KEY,
Body=archive)
+
+ bundle = self._archive_bundle()
+ bundle.initialize()
+ for child in bundle.path.iterdir():
+ child.unlink()
+ assert bundle.path.is_dir()
+ assert not bundle._get_archive_etag_path().exists()
+
+ client = bundle.s3_hook.get_conn()
+ recorded: list[dict] = []
+ original = client.get_object
+
+ def record_get_object(**kwargs):
+ recorded.append(kwargs)
+ return original(**kwargs)
+
+ with mock.patch.object(client, "get_object", record_get_object):
+ bundle.refresh()
+
+ assert recorded
+ assert (bundle.path / "dag_01.py").read_bytes() == b"test data"
+
+ def test_prefix_refresh_drops_archive_etag_marker(self, s3_bucket,
s3_client):
+ archive = _make_archive({"archive_only.py": b"from archive"})
+ s3_client.put_object(Bucket=S3_BUCKET_NAME, Key=S3_ARCHIVE_KEY,
Body=archive)
+
+ archive_bundle = self._archive_bundle()
+ archive_bundle.initialize()
+ assert (archive_bundle.path / ".archive_etag").is_file()
+ assert (archive_bundle.path / "archive_only.py").is_file()
+
+ prefix_bundle = S3DagBundle(
+ name="test",
+ aws_conn_id=AWS_CONN_ID_WITH_REGION,
+ prefix=S3_BUCKET_PREFIX,
+ bucket_name=S3_BUCKET_NAME,
+ )
+ prefix_bundle.initialize()
+ assert archive_bundle.path == prefix_bundle.path / "dags"
+ assert not (archive_bundle.path / ".archive_etag").exists()
+ assert not (prefix_bundle.path / ".archive_etag").exists()
+ assert (prefix_bundle.path / "dag_01.py").is_file()
+ assert not (prefix_bundle.path / "archive_only.py").exists()
+
+ archive_bundle = self._archive_bundle()
+ archive_bundle.initialize()
+ assert (archive_bundle.path / "archive_only.py").read_bytes() ==
b"from archive"
+ assert not (archive_bundle.path / "dag_01.py").exists()
+ assert (archive_bundle.path / ".archive_etag").is_file()
+
+ def test_initialize_fails_when_archive_is_missing(self, s3_bucket,
s3_client):
+ bundle = self._archive_bundle()
+ with pytest.raises(ClientError) as exc_info:
+ bundle.initialize()
+ assert exc_info.value.response["Error"]["Code"] == "404"
+
+ def test_refresh_raises_and_keeps_bundle_when_archive_is_corrupt(self,
mocked_s3_resource, s3_client):
+ mocked_s3_resource.create_bucket(Bucket=S3_BUCKET_NAME)
+ s3_client.put_object(
+ Bucket=S3_BUCKET_NAME, Key=S3_ARCHIVE_KEY,
Body=_make_archive({"dag_01.py": b"test data"})
+ )
+
+ bundle = self._archive_bundle()
+ bundle.initialize()
+ assert (bundle.path / "dag_01.py").is_file()
+
+ s3_client.put_object(Bucket=S3_BUCKET_NAME, Key=S3_ARCHIVE_KEY,
Body=b"this is not a tarball")
+ with pytest.raises(tarfile.ReadError):
+ bundle.refresh()
+
+ assert (bundle.path / "dag_01.py").is_file()
+ assert not list(bundle.base_dir.glob("archive-staging-*"))
+
+ def test_archive_staging_restores_previous_bundle_when_swap_fails(self,
mocked_s3_resource, s3_client):
+ mocked_s3_resource.create_bucket(Bucket=S3_BUCKET_NAME)
+ s3_client.put_object(
+ Bucket=S3_BUCKET_NAME, Key=S3_ARCHIVE_KEY,
Body=_make_archive({"dag_01.py": b"old"})
+ )
+
+ bundle = self._archive_bundle()
+ bundle.initialize()
+ assert (bundle.path / "dag_01.py").read_text() == "old"
+
+ s3_client.put_object(
+ Bucket=S3_BUCKET_NAME, Key=S3_ARCHIVE_KEY,
Body=_make_archive({"dag_01.py": b"new"})
+ )
+
+ real_rename = Path.rename
+
+ def fail_swap_into_place(self: Path, target):
+ if self.name == "unpacked" and Path(target) == bundle.path:
+ raise OSError("swap failed")
+ return real_rename(self, target)
+
+ with (
+ mock.patch.object(Path, "rename", fail_swap_into_place),
+ pytest.raises(OSError, match="swap failed"),
+ ):
+ bundle.refresh()
+
+ assert (bundle.path / "dag_01.py").read_text() == "old"
+ assert not list(bundle.base_dir.glob("archive-staging-*"))
+
+ def test_refresh_removes_only_this_bundles_stale_staging_dir(self,
s3_bucket, s3_client):
+ archive = _make_archive({"dag_01.py": b"test data"})
+ s3_client.put_object(Bucket=S3_BUCKET_NAME, Key=S3_ARCHIVE_KEY,
Body=archive)
+ bundle = self._archive_bundle()
+ bundle.initialize()
+
+ storage_root = bundle.base_dir.parent
+ stale = bundle.base_dir / "archive-staging-dead"
+ stale.mkdir()
+ (stale / "orphan.txt").write_text("left by a killed process")
+ dotted_sibling = storage_root / f".{bundle.name}" /
"archive-staging-dead"
+ dotted_sibling.mkdir(parents=True)
+ other = storage_root / "other" / "archive-staging-dead"
+ other.mkdir(parents=True)
+
+ bundle.refresh()
+
+ assert not stale.exists()
+ assert dotted_sibling.is_dir()
+ assert other.is_dir()
+ assert (bundle.path / "dag_01.py").is_file()
+
+ def test_refresh_get_object_uses_head_etag(self, s3_bucket, s3_client):
+ s3_client.put_object(
+ Bucket=S3_BUCKET_NAME, Key=S3_ARCHIVE_KEY,
Body=_make_archive({"dag_01.py": b"test data"})
+ )
+ bundle = self._archive_bundle()
+ bundle.initialize()
+ marker = bundle._get_archive_etag_path()
+ etag = marker.read_text()
+ marker.unlink()
+
+ recorded: list[dict] = []
+ client = bundle.s3_hook.get_conn()
+ original = client.get_object
+
+ def record_get_object(**kwargs):
+ recorded.append(kwargs)
+ return original(**kwargs)
+
+ with mock.patch.object(client, "get_object", record_get_object):
+ bundle.refresh()
+
+ assert recorded[0]["IfMatch"] == etag
+ assert (bundle.path / "dag_01.py").is_file()
+
+ def test_refresh_keeps_bundle_when_archive_etag_no_longer_matches(self,
s3_bucket, s3_client):
+ s3_client.put_object(
+ Bucket=S3_BUCKET_NAME, Key=S3_ARCHIVE_KEY,
Body=_make_archive({"dag_01.py": b"test data"})
+ )
+ bundle = self._archive_bundle()
+ bundle.initialize()
+ bundle._get_archive_etag_path().unlink()
+
+ error = ClientError(
+ {"Error": {"Code": "PreconditionFailed", "Message": "precondition
failed"}},
+ "GetObject",
+ )
+ client = bundle.s3_hook.get_conn()
+ with mock.patch.object(client, "get_object", side_effect=error),
pytest.raises(ClientError):
+ bundle.refresh()
+
+ assert (bundle.path / "dag_01.py").is_file()
+ assert not list(bundle.base_dir.glob("archive-staging-*"))
+
+ def test_refresh_rejects_archive_member_that_escapes_bundle(self,
mocked_s3_resource, s3_client):
+ mocked_s3_resource.create_bucket(Bucket=S3_BUCKET_NAME)
+ s3_client.put_object(
+ Bucket=S3_BUCKET_NAME, Key=S3_ARCHIVE_KEY,
Body=_make_archive({"dag_01.py": b"kept"})
+ )
+ bundle = self._archive_bundle()
+ bundle.initialize()
+
+ s3_client.put_object(
+ Bucket=S3_BUCKET_NAME,
+ Key=S3_ARCHIVE_KEY,
+ Body=_make_archive({"../escape.py": b"nope"}),
+ )
+ with pytest.raises(tarfile.TarError):
+ bundle.refresh()
+
+ assert (bundle.path / "dag_01.py").read_bytes() == b"kept"
+ assert not (bundle.base_dir / "escape.py").exists()
Review Comment:
This would never fail because we don’t write `escape.py` at this location.
It’s in the staging directory instead.
--
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]
