diff --git a/docs/docs/pypaimon/catalogs.mdx b/docs/docs/pypaimon/catalogs.mdx index 45f035b23654..e5354c76968b 100644 --- a/docs/docs/pypaimon/catalogs.mdx +++ b/docs/docs/pypaimon/catalogs.mdx @@ -170,6 +170,17 @@ PyPaimon supports filesystem, JDBC, and REST catalogs. See [Catalog](../concepts Use this catalog for the database and table operations below. +## OSS and custom S3 checksums + +OSS rejects PyArrow 22+ optional checksum trailers (checksums sent after the data). +For PyArrow-backed OSS and explicit S3 endpoints (including AWS), PyPaimon sets +`AWS_REQUEST_CHECKSUM_CALCULATION=WHEN_REQUIRED` to disable optional checksums. +This also affects later AWS clients in the process. Explicit environment values take +precedence; required checksums and older PyArrow versions are unchanged. + +Set `fs.s3.checksum-compatibility.auto-configure=false` on all relevant catalogs before +creating clients to disable this setup. Existing process settings are not cleared. + ## Create Database Tables belong to a database. Create the database before creating its tables. diff --git a/paimon-python/pypaimon/common/options/config.py b/paimon-python/pypaimon/common/options/config.py index e67030957954..f6d7026789b4 100644 --- a/paimon-python/pypaimon/common/options/config.py +++ b/paimon-python/pypaimon/common/options/config.py @@ -46,6 +46,9 @@ class OssOptions: class S3Options: + CHECKSUM_COMPATIBILITY_AUTO_CONFIGURE = ( + ConfigOptions.key("fs.s3.checksum-compatibility.auto-configure").boolean_type().default_value(True) + .with_description("Apply a process-wide optional-checksum default for PyArrow OSS/custom S3 clients.")) S3_ACCESS_KEY_ID = ConfigOptions.key("fs.s3.accessKeyId").string_type().no_default_value().with_description( "S3 access key ID") S3_ACCESS_KEY_SECRET = ConfigOptions.key("fs.s3.accessKeySecret").string_type().no_default_value().with_description( diff --git a/paimon-python/pypaimon/filesystem/pyarrow_file_io.py b/paimon-python/pypaimon/filesystem/pyarrow_file_io.py index ad5d3c8a4ea9..446bb64ed5f2 100644 --- a/paimon-python/pypaimon/filesystem/pyarrow_file_io.py +++ b/paimon-python/pypaimon/filesystem/pyarrow_file_io.py @@ -54,15 +54,15 @@ class PyArrowFileIO(FileIO): def __init__(self, path: str, catalog_options: Options): self.properties = catalog_options self.logger = logging.getLogger(__name__) - self._pyarrow_gte_8 = parse(pyarrow.__version__) >= parse("8.0.0") - # force_virtual_addressing landed in PyArrow 16; below it the OSS bucket - # goes into endpoint_override, so keys must omit it (init + path share - # this flag so they can't drift). - self._pyarrow_gte_16 = parse(pyarrow.__version__) >= parse("16.0.0") - self._oss_bucket_in_endpoint = not self._pyarrow_gte_16 + self._set_pyarrow_version() scheme, netloc, _ = self.parse_location(path) self.uri_reader_factory = UriReaderFactory(catalog_options) self._is_oss = scheme in {"oss"} + self._is_s3 = scheme in {"s3", "s3a", "s3n"} + self._s3_endpoint = ( + self._get_s3_property("endpoint", S3Options.S3_ENDPOINT.key()) + if self._is_s3 else None + ) self._oss_bucket = None _oss_impl = self.properties.get(OssOptions.OSS_IMPL) self._use_jindo = False @@ -86,7 +86,7 @@ def __init__(self, path: str, catalog_options: Options): "Falling back to legacy PyArrow S3FileSystem implementation. " "Install pyjindosdk for better performance: pip install pyjindosdk") self.filesystem = self._initialize_oss_fs(path) - elif scheme in {"s3", "s3a", "s3n"}: + elif self._is_s3: self.filesystem = self._initialize_s3_fs() elif scheme in {"hdfs", "viewfs"}: self.filesystem = self._initialize_hdfs_fs(scheme, netloc) @@ -95,15 +95,52 @@ def __init__(self, path: str, catalog_options: Options): else: raise ValueError(f"Unrecognized filesystem type in URI: {scheme}") + def _set_pyarrow_version(self): + self._pyarrow_gte_8 = parse(pyarrow.__version__) >= parse("8.0.0") + # force_virtual_addressing landed in PyArrow 16; below it the OSS bucket + # goes into endpoint_override, so keys must omit it (init + path share + # this flag so they can't drift). + self._pyarrow_gte_16 = parse(pyarrow.__version__) >= parse("16.0.0") + self._oss_bucket_in_endpoint = not self._pyarrow_gte_16 + + def _uses_s3_compatibility(self) -> bool: + return (not self._use_jindo + and (self._is_oss or bool(self._s3_endpoint))) + + def _configure_s3_checksums(self): + if parse(pyarrow.__version__) < parse("22.0.0"): + return + if self._uses_s3_compatibility() and self.properties.get(S3Options.CHECKSUM_COMPATIBILITY_AUTO_CONFIGURE): + # Process-wide default; preserve explicit settings and do not restore it. + os.environ.setdefault("AWS_REQUEST_CHECKSUM_CALCULATION", "WHEN_REQUIRED") + def __getstate__(self): state = self.__dict__.copy() # threading.Lock cannot be pickled; recreated in __setstate__. state.pop("_legacy_bucket_lock", None) + state.pop("logger", None) + # Recreate S3-compatible clients with the worker's AWS SDK settings. + if self._uses_s3_compatibility(): + state.pop("filesystem", None) return state def __setstate__(self, state): self.__dict__.update(state) + self.logger = logging.getLogger(__name__) + self._set_pyarrow_version() + if "_is_s3" not in state: + self._is_s3 = (not self._is_oss + and isinstance(self.filesystem, pafs.S3FileSystem)) + if "_s3_endpoint" not in state: + self._s3_endpoint = ( + self._get_s3_property("endpoint", S3Options.S3_ENDPOINT.key()) + if self._is_s3 else None) self._legacy_bucket_lock = threading.Lock() + if self._uses_s3_compatibility(): + self.filesystem = ( + self._initialize_oss_fs(None) + if self._is_oss else self._initialize_s3_fs() + ) @staticmethod def parse_location(location: str): @@ -193,6 +230,7 @@ def _initialize_jindo_fs(self, path) -> FileSystem: return pafs.PyFileSystem(fs_handler) def _initialize_oss_fs(self, path) -> FileSystem: + self._configure_s3_checksums() if self.properties.get(OssOptions.OSS_ACCESS_KEY_ID): # When explicit credentials are provided, disable the EC2 Instance Metadata # Service (IMDS) probe to avoid multi-second timeouts in non-AWS environments. @@ -222,6 +260,7 @@ def _initialize_oss_fs(self, path) -> FileSystem: return pafs.S3FileSystem(**client_kwargs) def _initialize_s3_fs(self) -> FileSystem: + self._configure_s3_checksums() access_key = self._get_property( S3Options.S3_ACCESS_KEY_ID.key(), *self._s3_key_variants("access-key", "access.key")) diff --git a/paimon-python/pypaimon/tests/s3_client_compatibility_test.py b/paimon-python/pypaimon/tests/s3_client_compatibility_test.py new file mode 100644 index 000000000000..3e5456779747 --- /dev/null +++ b/paimon-python/pypaimon/tests/s3_client_compatibility_test.py @@ -0,0 +1,359 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +"""Unit tests for PyArrow-backed OSS and S3-compatible storage. + +No real OSS access is required. +""" + +import multiprocessing +import os +import pickle +import unittest +import threading +from itertools import product +from http.server import BaseHTTPRequestHandler, HTTPServer +from socketserver import ThreadingMixIn +from unittest import mock + +import pyarrow +import pyarrow.fs as pafs +from packaging.version import parse + +from pypaimon.common.options import Options +from pypaimon.common.options.config import OssOptions, S3Options +from pypaimon.filesystem.pyarrow_file_io import PyArrowFileIO + + +def _restore_s3_file_io(connection): + os.environ.pop("AWS_REQUEST_CHECKSUM_CALCULATION", None) + connection.send("ready") + payload = connection.recv_bytes() + client = object() + settings = [] + + def create_client(**kwargs): + settings.append(os.environ.get("AWS_REQUEST_CHECKSUM_CALCULATION")) + return client + + with mock.patch("pyarrow.fs.S3FileSystem", side_effect=create_client) as s3: + restored = pickle.loads(payload) + connection.send(( + os.environ.get("AWS_REQUEST_CHECKSUM_CALCULATION"), + settings, + s3.call_count, + restored.filesystem is client, + )) + connection.close() + + +def _write_in_worker(connection): + os.environ.pop("AWS_REQUEST_CHECKSUM_CALCULATION", None) + connection.send("ready") + try: + file_io = pickle.loads(connection.recv_bytes()) + with file_io.filesystem.open_output_stream("test-bucket/file") as stream: + stream.write(b"x") + connection.send("written") + except Exception as error: + connection.send(repr(error)) + finally: + connection.close() + + +class _HTTPServer(ThreadingMixIn, HTTPServer): + daemon_threads = True + + +class _UploadHandler(BaseHTTPRequestHandler): + def _respond(self, body=b""): + self.send_response(200) + self.send_header("Content-Length", str(len(body))) + self.send_header("ETag", '"etag"') + self.send_header("Connection", "close") + self.end_headers() + self.wfile.write(body) + self.close_connection = True + + def do_POST(self): + if "uploads" in self.path: + self._respond(b"test-bucket" + b"filetest-id" + b"") + else: + self._respond(b"test-bucket" + b"fileetag") + + def do_PUT(self): + self.server.put_headers.append(dict(self.headers)) + self._respond() + + def log_message(self, *args): + pass + + +class S3ClientCompatibilityTest(unittest.TestCase): + def setUp(self): + environment = mock.patch.dict(os.environ) + environment.start() + self.addCleanup(environment.stop) + os.environ.pop("AWS_REQUEST_CHECKSUM_CALCULATION", None) + + def _expected_checksum_default(self): + return "WHEN_REQUIRED" if parse(pyarrow.__version__) >= parse("22.0.0") else None + + def test_checksum_default_version_boundary(self): + for version in ("19.0.1", "20.0.0", "21.0.0", "22.0.0", "23.0.0"): + for scheme in ("oss", "s3", "s3a", "s3n"): + with self.subTest(version=version, scheme=scheme), \ + mock.patch.dict(os.environ), \ + mock.patch.object(pyarrow, "__version__", version): + os.environ.pop("AWS_REQUEST_CHECKSUM_CALCULATION", None) + self._new_file_io(scheme) + expected = "WHEN_REQUIRED" if version in ("22.0.0", "23.0.0") else None + self.assertEqual(expected, os.environ.get("AWS_REQUEST_CHECKSUM_CALCULATION")) + + def test_checksum_compatibility_can_be_disabled(self): + with mock.patch("pyarrow.fs.S3FileSystem", return_value=mock.Mock()): + PyArrowFileIO("s3://test-bucket/", Options({ + "fs.s3.endpoint": "http://minio:9000", + "fs.s3.checksum-compatibility.auto-configure": "false", + })) + self.assertNotIn("AWS_REQUEST_CHECKSUM_CALCULATION", os.environ) + + @mock.patch.object(pyarrow, "__version__", "23.0.0") + def test_disabling_auto_configuration_preserves_existing_process_default(self): + settings = [] + + def create_client(**kwargs): + settings.append(os.environ.get("AWS_REQUEST_CHECKSUM_CALCULATION")) + return mock.Mock() + + with mock.patch("pyarrow.fs.S3FileSystem", side_effect=create_client): + for auto_configure in ("true", "false"): + PyArrowFileIO("s3://test-bucket/", Options({ + "fs.s3.endpoint": "http://minio:9000", + "fs.s3.checksum-compatibility.auto-configure": auto_configure, + })) + self.assertEqual(["WHEN_REQUIRED", "WHEN_REQUIRED"], settings) + self.assertEqual("WHEN_REQUIRED", os.environ["AWS_REQUEST_CHECKSUM_CALCULATION"]) + + def _new_file_io(self, scheme="s3"): + if scheme == "oss": + options = Options({ + OssOptions.OSS_IMPL.key(): "legacy", + OssOptions.OSS_ENDPOINT.key(): "oss-cn-test.example.com", + OssOptions.OSS_ACCESS_KEY_ID.key(): "ak", + OssOptions.OSS_ACCESS_KEY_SECRET.key(): "sk", + }) + else: + options = Options({S3Options.S3_ENDPOINT.key(): "http://minio:9000"}) + with mock.patch("pyarrow.fs.S3FileSystem", return_value=pafs.LocalFileSystem()): + file_io = PyArrowFileIO(scheme + "://test-bucket/", options) + return file_io + + def test_process_default_retained_after_client_creation_failure(self): + for previous in (None, "WHEN_SUPPORTED"): + with self.subTest(previous=previous), mock.patch.dict(os.environ, {}, clear=True): + if previous is not None: + os.environ["AWS_REQUEST_CHECKSUM_CALCULATION"] = previous + with mock.patch("pyarrow.fs.S3FileSystem", side_effect=RuntimeError("failed")): + with self.assertRaisesRegex(RuntimeError, "failed"): + self._new_file_io_failure() + self.assertEqual(previous or self._expected_checksum_default(), + os.environ.get("AWS_REQUEST_CHECKSUM_CALCULATION")) + + def _new_file_io_failure(self): + PyArrowFileIO("s3://test-bucket/", Options({"fs.s3.endpoint": "http://minio:9000"})) + + def test_pickle_rebuilds_client_in_started_worker(self): + methods = [name for name in ("spawn", "fork") + if name in multiprocessing.get_all_start_methods()] + for scheme, method in product(("oss", "s3", "s3a", "s3n"), methods): + with self.subTest(scheme=scheme, method=method): + context = multiprocessing.get_context(method) + parent, child = context.Pipe() + process = context.Process(target=_restore_s3_file_io, args=(child,)) + process.start() + child.close() + try: + self.assertTrue(parent.poll(20)) + self.assertEqual("ready", parent.recv()) + file_io = self._new_file_io(scheme) + self.assertNotIn("filesystem", file_io.__getstate__()) + parent.send_bytes(pickle.dumps(file_io)) + self.assertTrue(parent.poll(20)) + expected = self._expected_checksum_default() + self.assertEqual((expected, [expected], 1, True), parent.recv()) + finally: + parent.close() + process.join(20) + if process.is_alive(): + process.terminate() + process.join() + self.assertEqual(0, process.exitcode) + + def test_oss_initialization_preserves_process_setting(self): + options = Options({ + OssOptions.OSS_ACCESS_KEY_ID.key(): "ak", + OssOptions.OSS_ACCESS_KEY_SECRET.key(): "sk", + OssOptions.OSS_ENDPOINT.key(): "oss-cn-test.example.com", + OssOptions.OSS_REGION.key(): "cn-test", + OssOptions.OSS_IMPL.key(): "legacy", + }) + settings = [] + + def create_client(**kwargs): + settings.append(os.environ.get("AWS_REQUEST_CHECKSUM_CALCULATION")) + return mock.Mock() + + with mock.patch.dict("os.environ", { + "AWS_REQUEST_CHECKSUM_CALCULATION": "WHEN_SUPPORTED"}, clear=True), \ + mock.patch("pyarrow.fs.S3FileSystem", side_effect=create_client): + PyArrowFileIO("oss://test-bucket/", options) + self.assertEqual( + "WHEN_SUPPORTED", + os.environ["AWS_REQUEST_CHECKSUM_CALCULATION"]) + self.assertEqual(["WHEN_SUPPORTED"], settings) + + def test_initialization_sets_default_for_all_s3_schemes(self): + options = Options({ + S3Options.S3_ENDPOINT.key(): "http://minio:9000", + }) + for scheme in ("s3", "s3a", "s3n"): + settings = [] + + def create_client(**kwargs): + settings.append(os.environ.get("AWS_REQUEST_CHECKSUM_CALCULATION")) + return mock.Mock() + + with self.subTest(scheme=scheme), \ + mock.patch.dict("os.environ", {}, clear=True), \ + mock.patch("pyarrow.fs.S3FileSystem", side_effect=create_client): + PyArrowFileIO( + "{}://test-bucket/warehouse".format(scheme), options) + self.assertEqual(self._expected_checksum_default(), + os.environ.get("AWS_REQUEST_CHECKSUM_CALCULATION")) + self.assertEqual([self._expected_checksum_default()], settings) + + def test_native_s3_does_not_change_checksum_setting(self): + with mock.patch.dict("os.environ", {}, clear=True), \ + mock.patch("pyarrow.fs.S3FileSystem", return_value=mock.Mock()): + PyArrowFileIO("s3://test-bucket/warehouse", Options({})) + self.assertNotIn( + "AWS_REQUEST_CHECKSUM_CALCULATION", os.environ) + + def test_worker_recomputes_pyarrow_version(self): + state = self._new_file_io().__getstate__() + state.update({ + "_pyarrow_gte_8": False, + "_pyarrow_gte_16": False, + "_oss_bucket_in_endpoint": True, + }) + restored = object.__new__(PyArrowFileIO) + with mock.patch.object( + PyArrowFileIO, "_initialize_s3_fs", return_value=object()): + restored.__setstate__(state) + + version = parse(pyarrow.__version__) + self.assertEqual(version >= parse("8.0.0"), restored._pyarrow_gte_8) + self.assertEqual(version >= parse("16.0.0"), restored._pyarrow_gte_16) + self.assertEqual(version < parse("16.0.0"), restored._oss_bucket_in_endpoint) + + def test_worker_accepts_old_non_oss_pickle(self): + state = self._new_file_io().__dict__.copy() + for key in ("_legacy_bucket_lock", "_is_s3", + "_s3_endpoint"): + state.pop(key) + state["filesystem"] = mock.Mock(spec=pafs.S3FileSystem) + restored = object.__new__(PyArrowFileIO) + with mock.patch.object( + PyArrowFileIO, "_initialize_s3_fs", return_value=object() + ) as initialize: + restored.__setstate__(state) + self.assertTrue(restored._is_s3) + self.assertEqual("http://minio:9000", restored._s3_endpoint) + initialize.assert_called_once() + + state["filesystem"] = pafs.LocalFileSystem() + restored = object.__new__(PyArrowFileIO) + restored.__setstate__(state) + self.assertFalse(restored._is_s3) + self.assertIsNone(restored._s3_endpoint) + + @unittest.skipUnless( + parse(pyarrow.__version__) >= parse("22.0.0"), + "requires PyArrow 22+ optional request checksums", + ) + def test_default_applies_to_parent_and_worker(self): + server = _HTTPServer( + ("127.0.0.1", 0), _UploadHandler) + server.requests = [] + server.put_headers = [] + server.bucket_objects = {"test-bucket": set()} + server_thread = threading.Thread(target=server.serve_forever) + server_thread.start() + try: + endpoint = "http://127.0.0.1:{}".format(server.server_port) + options = Options({ + S3Options.S3_ACCESS_KEY_ID.key(): "ak", + S3Options.S3_ACCESS_KEY_SECRET.key(): "sk", + S3Options.S3_ENDPOINT.key(): endpoint, + S3Options.S3_REGION.key(): "us-east-1", + "fs.s3.path.style.access": "true", + }) + with mock.patch.dict(os.environ, { + "NO_PROXY": "127.0.0.1,localhost", + "no_proxy": "127.0.0.1,localhost", + }): + compatible = PyArrowFileIO("s3://test-bucket/", options) + native = pafs.S3FileSystem( + access_key="ak", secret_key="sk", region="us-east-1", + endpoint_override=endpoint) + self.assertEqual("WHEN_REQUIRED", + os.environ["AWS_REQUEST_CHECKSUM_CALCULATION"]) + with native.open_output_stream("test-bucket/file") as stream: + stream.write(b"x") + with compatible.filesystem.open_output_stream( + "test-bucket/file") as stream: + stream.write(b"x") + + context = multiprocessing.get_context("spawn") + parent, child = context.Pipe() + process = context.Process(target=_write_in_worker, args=(child,)) + process.start() + child.close() + try: + self.assertTrue(parent.poll(20)) + self.assertEqual("ready", parent.recv()) + parent.send_bytes(pickle.dumps(compatible)) + self.assertTrue(parent.poll(20)) + self.assertEqual("written", parent.recv()) + finally: + parent.close() + process.join(20) + if process.is_alive(): + process.terminate() + process.join() + self.assertEqual(0, process.exitcode) + self.assertEqual(3, len(server.put_headers)) + for headers in server.put_headers: + self.assertNotIn("x-amz-trailer", {key.lower() for key in headers}) + finally: + server.shutdown() + server.server_close() + server_thread.join()