"""Firehose S3 delivery compression (issue #2328). The delivered S3 object must be compressed the way the stream's CompressionFormat says, with the key extension and Content-Encoding real AWS uses. In CI this suite also runs against the amd64 native image, where the two Snappy formats depend on a JNI library bundled per platform, so this test is the proof that the x86_64 library survived the native build. Both Snappy variants are container formats around raw Snappy blocks: ``Snappy`` is the xerial snappy-java stream format (not the official framing format) and ``HADOOP_SNAPPY`` is Hadoop's block framing. Neither has a Python reader without native system libraries, so the framing is parsed here and only the raw blocks go through cramjam. """ import gzip import io import logging import struct import time import uuid import zipfile import boto3 import cramjam import pytest from botocore.exceptions import ClientError logger = logging.getLogger(__name__) RECORD = '{"seq":1}\n' RECORD_COUNT = 5 EXPECTED_PAYLOAD = (RECORD * RECORD_COUNT).encode() # The buffering interval has to elapse before delivery happens, so it is kept short. # AWS accepts 0..900, so this is still a contract-faithful stream; against Floci the # flush lands on the next tick of its background flusher. BUFFERING_INTERVAL_SECONDS = 5 POLL_INTERVAL_SECONDS = 5 POLL_TIMEOUT_SECONDS = 240 ROLE_ARN = "arn:aws:iam::000000000000:role/firehose-role" XERIAL_MAGIC = b"\x82SNAPPY\x00" def _unzip(body): with zipfile.ZipFile(io.BytesIO(body)) as archive: names = archive.namelist() assert len(names) == 1, f"expected a single zip entry, got {names}" return archive.read(names[0]) def _raw_snappy(chunk): return bytes(cramjam.snappy.decompress_raw(chunk)) def _xerial_snappy(body): """snappy-java's SnappyOutputStream format. A 16 byte header (8 byte magic, 4 byte version, 4 byte compatible version), then chunks of a big-endian length and a raw Snappy block. A negative length marks an uncompressed chunk of that many bytes, which snappy-java emits when compression would not shrink a block. """ assert body[:8] == XERIAL_MAGIC, f"not a snappy-java stream: {body[:8]!r}" out = bytearray() offset = 16 while offset < len(body): (length,) = struct.unpack(">i", body[offset:offset + 4]) offset += 4 if length < 0: out += body[offset:offset - length] offset -= length else: out += _raw_snappy(body[offset:offset + length]) offset += length return bytes(out) def _hadoop_snappy(body): """Hadoop's block framing. Per block, a big-endian uncompressed block length followed by one or more chunks of a big-endian compressed length and a raw Snappy block. The leading length belongs to the first block, not the whole payload, so anything above the 32 KiB block size carries several blocks and must be looped over. """ out = bytearray() offset = 0 while offset < len(body): (block_length,) = struct.unpack(">i", body[offset:offset + 4]) offset += 4 produced = 0 while produced < block_length: (chunk_length,) = struct.unpack(">i", body[offset:offset + 4]) offset += 4 plain = _raw_snappy(body[offset:offset + chunk_length]) offset += chunk_length out += plain produced += len(plain) assert produced == block_length, f"block decoded to {produced} bytes, header said {block_length}" return bytes(out) # (CompressionFormat as the API spells it, S3 key extension, Content-Encoding, decoder). # HADOOP_SNAPPY shares the .snappy extension with Snappy and differs only in framing # and Content-Encoding: the S3 object name documentation tables .hsnappy, but the # service delivers .snappy (verified against real AWS in issue #2328). FORMATS = [ ("UNCOMPRESSED", "", None, bytes), ("GZIP", ".gz", "gzip", gzip.decompress), ("ZIP", ".zip", "zip", _unzip), ("Snappy", ".snappy", "snappy-java", _xerial_snappy), ("HADOOP_SNAPPY", ".snappy", "hadoop-snappy", _hadoop_snappy), ] def _slug(compression_format): return compression_format.lower().replace("_", "-") @pytest.fixture(scope="class") def delivered(aws_config, client_config): """One bucket and one stream per format, all created and filled up front. A single buffering interval is then shared by every format instead of being waited out once per test. """ firehose = boto3.client("firehose", config=client_config, **aws_config) s3 = boto3.client("s3", config=client_config, **aws_config) suffix = uuid.uuid4().hex[:8] bucket = f"firehose-compression-{suffix}" stream_prefix = f"sdk-compression-{suffix}" streams = {} s3.create_bucket(Bucket=bucket) for compression_format, _, _, _ in FORMATS: stream_name = f"{stream_prefix}-{_slug(compression_format)}" streams[compression_format] = stream_name firehose.create_delivery_stream( DeliveryStreamName=stream_name, DeliveryStreamType="DirectPut", ExtendedS3DestinationConfiguration={ "RoleARN": ROLE_ARN, "BucketARN": f"arn:aws:s3:::{bucket}", "Prefix": f"{_slug(compression_format)}/", "CompressionFormat": compression_format, "BufferingHints": { "IntervalInSeconds": BUFFERING_INTERVAL_SECONDS, "SizeInMBs": 1, }, }, ) for stream_name in streams.values(): response = firehose.put_record_batch( DeliveryStreamName=stream_name, Records=[{"Data": RECORD.encode()} for _ in range(RECORD_COUNT)], ) assert response["FailedPutCount"] == 0, response yield {"s3": s3, "bucket": bucket} for stream_name in streams.values(): try: firehose.delete_delivery_stream(DeliveryStreamName=stream_name) except ClientError as e: logger.warning("delete_delivery_stream %s failed: %s", stream_name, e) try: paginator = s3.get_paginator("list_objects_v2") for page in paginator.paginate(Bucket=bucket): for obj in page.get("Contents", []): s3.delete_object(Bucket=bucket, Key=obj["Key"]) s3.delete_bucket(Bucket=bucket) except ClientError as e: logger.warning("bucket cleanup %s failed: %s", bucket, e) def _wait_for_delivered_key(s3, bucket, prefix): deadline = time.time() + POLL_TIMEOUT_SECONDS while time.time() < deadline: contents = s3.list_objects_v2(Bucket=bucket, Prefix=prefix).get("Contents", []) if contents: return contents[0]["Key"] time.sleep(POLL_INTERVAL_SECONDS) raise AssertionError(f"no object delivered under {prefix} within {POLL_TIMEOUT_SECONDS}s") @pytest.mark.timeout(POLL_TIMEOUT_SECONDS + 60) class TestFirehoseCompressionDelivery: @pytest.mark.parametrize( "compression_format,extension,content_encoding,decode", FORMATS, ids=[_slug(f[0]) for f in FORMATS], ) def test_delivered_object_is_compressed_as_declared( self, delivered, compression_format, extension, content_encoding, decode ): s3 = delivered["s3"] bucket = delivered["bucket"] key = _wait_for_delivered_key(s3, bucket, f"{_slug(compression_format)}/") assert key.endswith(extension), f"{key} should end with {extension!r}" if not extension: assert not key.endswith((".gz", ".zip", ".snappy")), key response = s3.get_object(Bucket=bucket, Key=key) assert response["ContentType"] == "application/octet-stream" assert response.get("ContentEncoding") == content_encoding body = response["Body"].read() assert decode(body) == EXPECTED_PAYLOAD, "delivered body must decode to the records that were put"