chrislu
2 weeks ago
4 changed files with 340 additions and 3 deletions
-
11.github/workflows/s3-parquet-tests.yml
-
57test/s3/parquet/Makefile
-
14test/s3/parquet/README.md
-
261test/s3/parquet/test_sse_s3_compatibility.py
@ -0,0 +1,261 @@ |
|||
#!/usr/bin/env python3 |
|||
""" |
|||
Test script for SSE-S3 compatibility with PyArrow native S3 filesystem. |
|||
|
|||
This test specifically targets the SSE-S3 multipart upload bug where |
|||
SeaweedFS panics with "bad IV length" when reading multipart uploads |
|||
that were encrypted with bucket-default SSE-S3. |
|||
|
|||
Requirements: |
|||
- pyarrow>=10.0.0 |
|||
- boto3>=1.28.0 |
|||
|
|||
Environment Variables: |
|||
S3_ENDPOINT_URL: S3 endpoint (default: localhost:8333) |
|||
S3_ACCESS_KEY: S3 access key (default: some_access_key1) |
|||
S3_SECRET_KEY: S3 secret key (default: some_secret_key1) |
|||
BUCKET_NAME: S3 bucket name (default: test-parquet-bucket) |
|||
|
|||
Usage: |
|||
# Start SeaweedFS with SSE-S3 enabled |
|||
make start-seaweedfs-ci ENABLE_SSE_S3=true |
|||
|
|||
# Run the test |
|||
python3 test_sse_s3_compatibility.py |
|||
""" |
|||
|
|||
import os |
|||
import secrets |
|||
import sys |
|||
import logging |
|||
from datetime import datetime |
|||
from typing import Optional |
|||
|
|||
import pyarrow as pa |
|||
import pyarrow.dataset as pads |
|||
import pyarrow.fs as pafs |
|||
import pyarrow.parquet as pq |
|||
|
|||
try: |
|||
import boto3 |
|||
from botocore.exceptions import ClientError |
|||
HAS_BOTO3 = True |
|||
except ImportError: |
|||
HAS_BOTO3 = False |
|||
logging.error("boto3 is required for this test") |
|||
sys.exit(1) |
|||
|
|||
logging.basicConfig(level=logging.INFO, format="%(message)s") |
|||
|
|||
# Configuration |
|||
S3_ENDPOINT_URL = os.environ.get("S3_ENDPOINT_URL", "localhost:8333") |
|||
S3_ACCESS_KEY = os.environ.get("S3_ACCESS_KEY", "some_access_key1") |
|||
S3_SECRET_KEY = os.environ.get("S3_SECRET_KEY", "some_secret_key1") |
|||
BUCKET_NAME = os.getenv("BUCKET_NAME", "test-parquet-bucket") |
|||
|
|||
TEST_RUN_ID = secrets.token_hex(8) |
|||
TEST_DIR = f"sse-s3-tests/{TEST_RUN_ID}" |
|||
|
|||
# Test sizes designed to trigger multipart uploads |
|||
# PyArrow typically uses 5MB chunks, so these sizes should trigger multipart |
|||
TEST_SIZES = { |
|||
"tiny": 10, # Single part |
|||
"small": 1_000, # Single part |
|||
"medium": 50_000, # Single part (~1.5MB) |
|||
"large": 200_000, # Multiple parts (~6MB) |
|||
"very_large": 500_000, # Multiple parts (~15MB) |
|||
} |
|||
|
|||
|
|||
def create_sample_table(num_rows: int = 5) -> pa.Table: |
|||
"""Create a sample PyArrow table for testing.""" |
|||
return pa.table( |
|||
{ |
|||
"id": pa.array(range(num_rows), type=pa.int64()), |
|||
"name": pa.array([f"user_{i}" for i in range(num_rows)], type=pa.string()), |
|||
"value": pa.array([float(i) * 1.5 for i in range(num_rows)], type=pa.float64()), |
|||
"flag": pa.array([i % 2 == 0 for i in range(num_rows)], type=pa.bool_()), |
|||
} |
|||
) |
|||
|
|||
|
|||
def init_s3_filesystem() -> tuple[Optional[pafs.S3FileSystem], str, str]: |
|||
"""Initialize PyArrow's native S3 filesystem.""" |
|||
try: |
|||
logging.info("Initializing PyArrow S3FileSystem...") |
|||
|
|||
# Determine scheme from endpoint |
|||
if S3_ENDPOINT_URL.startswith("http://"): |
|||
scheme = "http" |
|||
endpoint = S3_ENDPOINT_URL[7:] |
|||
elif S3_ENDPOINT_URL.startswith("https://"): |
|||
scheme = "https" |
|||
endpoint = S3_ENDPOINT_URL[8:] |
|||
else: |
|||
scheme = "http" |
|||
endpoint = S3_ENDPOINT_URL |
|||
|
|||
s3 = pafs.S3FileSystem( |
|||
access_key=S3_ACCESS_KEY, |
|||
secret_key=S3_SECRET_KEY, |
|||
endpoint_override=endpoint, |
|||
scheme=scheme, |
|||
allow_bucket_creation=True, |
|||
allow_bucket_deletion=True, |
|||
) |
|||
|
|||
logging.info("✓ PyArrow S3FileSystem initialized\n") |
|||
return s3, scheme, endpoint |
|||
except Exception as e: |
|||
logging.exception("✗ Failed to initialize PyArrow S3FileSystem") |
|||
return None, "", "" |
|||
|
|||
|
|||
def ensure_bucket_exists(scheme: str, endpoint: str) -> bool: |
|||
"""Ensure the test bucket exists using boto3.""" |
|||
try: |
|||
endpoint_url = f"{scheme}://{endpoint}" |
|||
s3_client = boto3.client( |
|||
's3', |
|||
endpoint_url=endpoint_url, |
|||
aws_access_key_id=S3_ACCESS_KEY, |
|||
aws_secret_access_key=S3_SECRET_KEY, |
|||
region_name='us-east-1', |
|||
) |
|||
|
|||
try: |
|||
s3_client.head_bucket(Bucket=BUCKET_NAME) |
|||
logging.info(f"✓ Bucket exists: {BUCKET_NAME}") |
|||
except ClientError as e: |
|||
error_code = e.response['Error']['Code'] |
|||
if error_code == '404': |
|||
logging.info(f"Creating bucket: {BUCKET_NAME}") |
|||
s3_client.create_bucket(Bucket=BUCKET_NAME) |
|||
logging.info(f"✓ Bucket created: {BUCKET_NAME}") |
|||
|
|||
# Note: SeaweedFS doesn't support GetBucketEncryption API |
|||
# so we can't verify if SSE-S3 is enabled via API |
|||
# We assume it's configured correctly in the s3.json config file |
|||
logging.info(f"✓ Assuming SSE-S3 is configured in s3.json") |
|||
return True |
|||
|
|||
except Exception as e: |
|||
logging.exception(f"✗ Failed to check bucket: {e}") |
|||
return False |
|||
|
|||
|
|||
def test_write_read_with_sse( |
|||
s3: pafs.S3FileSystem, |
|||
test_name: str, |
|||
num_rows: int |
|||
) -> tuple[bool, str, int]: |
|||
"""Test writing and reading with SSE-S3 encryption.""" |
|||
try: |
|||
table = create_sample_table(num_rows) |
|||
filename = f"{BUCKET_NAME}/{TEST_DIR}/{test_name}/data.parquet" |
|||
|
|||
logging.info(f" Writing {num_rows:,} rows...") |
|||
pads.write_dataset( |
|||
table, |
|||
filename, |
|||
filesystem=s3, |
|||
format="parquet", |
|||
) |
|||
|
|||
logging.info(f" Reading back...") |
|||
table_read = pq.read_table(filename, filesystem=s3) |
|||
|
|||
if table_read.num_rows != num_rows: |
|||
return False, f"Row count mismatch: {table_read.num_rows} != {num_rows}", 0 |
|||
|
|||
return True, "Success", table_read.num_rows |
|||
|
|||
except Exception as e: |
|||
error_msg = f"{type(e).__name__}: {str(e)}" |
|||
logging.error(f" ✗ Failed: {error_msg}") |
|||
return False, error_msg, 0 |
|||
|
|||
|
|||
def main(): |
|||
"""Run SSE-S3 compatibility tests.""" |
|||
print("=" * 80) |
|||
print("SSE-S3 Compatibility Tests for PyArrow Native S3") |
|||
print("Testing Multipart Upload Encryption") |
|||
print("=" * 80 + "\n") |
|||
|
|||
print("Configuration:") |
|||
print(f" S3 Endpoint: {S3_ENDPOINT_URL}") |
|||
print(f" Bucket: {BUCKET_NAME}") |
|||
print(f" Test Directory: {TEST_DIR}") |
|||
print(f" PyArrow Version: {pa.__version__}") |
|||
print() |
|||
|
|||
# Initialize |
|||
s3, scheme, endpoint = init_s3_filesystem() |
|||
if s3 is None: |
|||
print("Cannot proceed without S3 connection") |
|||
return 1 |
|||
|
|||
# Check bucket and SSE-S3 |
|||
if not ensure_bucket_exists(scheme, endpoint): |
|||
print("\n⚠ WARNING: SSE-S3 is not enabled on the bucket!") |
|||
print("This test requires SSE-S3 encryption to be enabled.") |
|||
print("Please start SeaweedFS with: make start-seaweedfs-ci ENABLE_SSE_S3=true") |
|||
return 1 |
|||
|
|||
print() |
|||
results = [] |
|||
|
|||
# Test all sizes |
|||
for size_name, num_rows in TEST_SIZES.items(): |
|||
print(f"\n{'='*80}") |
|||
print(f"Testing {size_name} dataset ({num_rows:,} rows)") |
|||
print(f"{'='*80}") |
|||
|
|||
success, message, rows_read = test_write_read_with_sse( |
|||
s3, size_name, num_rows |
|||
) |
|||
results.append((size_name, num_rows, success, message, rows_read)) |
|||
|
|||
if success: |
|||
print(f" ✓ SUCCESS: Read {rows_read:,} rows") |
|||
else: |
|||
print(f" ✗ FAILED: {message}") |
|||
|
|||
# Summary |
|||
print("\n" + "=" * 80) |
|||
print("SUMMARY") |
|||
print("=" * 80) |
|||
|
|||
passed = sum(1 for _, _, success, _, _ in results if success) |
|||
total = len(results) |
|||
print(f"\nTotal: {passed}/{total} tests passed\n") |
|||
|
|||
print(f"{'Size':<15} {'Rows':>10} {'Status':<10} {'Message':<40}") |
|||
print("-" * 80) |
|||
for size_name, num_rows, success, message, rows_read in results: |
|||
status = "✓ PASS" if success else "✗ FAIL" |
|||
print(f"{size_name:<15} {num_rows:>10,} {status:<10} {message[:40]}") |
|||
|
|||
print("\n" + "=" * 80) |
|||
if passed == total: |
|||
print("✓ ALL TESTS PASSED WITH SSE-S3!") |
|||
print("\nThis means:") |
|||
print(" - SSE-S3 encryption is working correctly") |
|||
print(" - PyArrow native S3 filesystem is compatible") |
|||
print(" - Multipart uploads are handled properly") |
|||
else: |
|||
print(f"✗ {total - passed} test(s) failed") |
|||
print("\nPossible issues:") |
|||
print(" - SSE-S3 multipart upload bug with empty IV") |
|||
print(" - Encryption/decryption mismatch") |
|||
print(" - File corruption during upload") |
|||
|
|||
print("=" * 80 + "\n") |
|||
|
|||
return 0 if passed == total else 1 |
|||
|
|||
|
|||
if __name__ == "__main__": |
|||
sys.exit(main()) |
|||
|
|||
Write
Preview
Loading…
Cancel
Save
Reference in new issue