[python] Use conditional S3 puts for atomic metadata writes

.github/workflows/ci-python.yml

@@ -122,7 +122,7 @@ jobs:
           if [[ "${{ matrix.python-version }}" == "3.6.15" ]]; then
             python -m pip install --upgrade pip==21.3.1
             python --version
-            python -m pip install --no-cache-dir pyroaring readerwriterlock==1.0.9 'fsspec==2021.10.1' 'cachetools==4.2.4' 'ossfs==2021.8.0' pyarrow==6.0.1 pandas==1.1.5 'polars==0.9.12' 'fastavro==1.4.7' zstandard==0.19.0 dataclasses==0.8.0 flake8 pytest py4j==0.10.9.9 requests parameterized==0.8.1 datasketches==4.1.0 2>&1 >/dev/null
+            python -m pip install --no-cache-dir pyroaring readerwriterlock==1.0.9 'fsspec==2021.10.1' 'cachetools==4.2.4' 'ossfs==2021.8.0' pyarrow==6.0.1 pandas==1.1.5 'polars==0.9.12' 'fastavro==1.4.7' zstandard==0.19.0 dataclasses==0.8.0 boto3 flake8 pytest py4j==0.10.9.9 requests parameterized==0.8.1 datasketches==4.1.0 2>&1 >/dev/null
             python -m pip install 'lumina-data>=${{ env.LUMINA_DATA_VERSION }}' -i https://pypi.org/simple/
           elif [[ "${{ matrix.python-version }}" == "3.7" ]]; then
             # 3.7 installs the version-pinned set declared for 3.7 in dev/requirements.txt.
@@ -134,7 +134,7 @@ jobs:
           else
             python -m pip install --upgrade pip
             pip install torch --index-url https://download.pytorch.org/whl/cpu
-            python -m pip install pyroaring readerwriterlock==1.0.9 fsspec==2024.3.1 cachetools==5.3.3 ossfs==2023.12.0 ray==2.54.0 fastavro==1.11.1 'isal>=1.8,<2' zstandard==0.24.0 polars==1.32.0 duckdb==1.3.2 pylance==0.39.0 cramjam pytest~=7.0 py4j==0.10.9.9 requests parameterized==0.9.0 'daft>=0.7.6' 'datafusion>=54,<55' datasketches
+            python -m pip install pyroaring readerwriterlock==1.0.9 fsspec==2024.3.1 cachetools==5.3.3 ossfs==2023.12.0 ray==2.54.0 fastavro==1.11.1 'isal>=1.8,<2' zstandard==0.24.0 polars==1.32.0 duckdb==1.3.2 pylance==0.39.0 cramjam boto3 pytest~=7.0 py4j==0.10.9.9 requests parameterized==0.9.0 'daft>=0.7.6' 'datafusion>=54,<55' datasketches
             if python -c "import sys; sys.exit(0 if sys.version_info >= (3, 12) else 1)"; then
               python -m pip install pyarrow==24.0.0 numpy==2.4.6 pandas==2.3.3 flake8==7.1.2
             else

paimon-python/dev/lint-python.sh

@@ -198,7 +198,7 @@ function pytest_check() {
     # 3.6/3.7 run a curated core subset (their dep ceiling rules out the
     # vector/index/multimodal/blob suites); 3.10+ run the full suite.
     if [ "$PYTHON_VERSION" = "3.6" ] || [ "$PYTHON_VERSION" = "3.7" ]; then
-        TEST_DIR="pypaimon/tests/py36 pypaimon/tests/file_io_test.py"
+        TEST_DIR="pypaimon/tests/py36 pypaimon/tests/file_io_test.py pypaimon/tests/s3_atomic_write_test.py"
         echo "Running core test subset for Python $PYTHON_VERSION: $TEST_DIR"
     else
         TEST_DIR="pypaimon/tests pypaimon/acceptance --ignore=pypaimon/tests/py36 --ignore=pypaimon/tests/e2e --ignore=pypaimon/tests/torch_read_test.py"

paimon-python/dev/requirements.txt

@@ -40,5 +40,6 @@ requests>=2.21.0,<3
 urllib3>=1.26,<3
 zstandard>=0.19,<1
 backports.zstd>=1.0.0,<1.4.0; python_version >= "3.9" and python_version < "3.14"
+boto3>=1.23.10,<2
 cramjam>=1.3.0,<3; python_version>="3.7"
 pyyaml>=5.4,<7

paimon-python/pypaimon/filesystem/pyarrow_file_io.py

@@ -62,12 +62,14 @@ def __init__(self, path: str, catalog_options: Options):
         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._oss_bucket = None
         _oss_impl = self.properties.get(OssOptions.OSS_IMPL)
         self._use_jindo = False
         self._legacy_bucket_checked = False
         self._legacy_bucket_error = None
         self._legacy_bucket_lock = threading.Lock()
+        self._s3_atomic_client = None
 
         if self._is_oss:
             self._oss_bucket = self._extract_oss_bucket(path)
@@ -98,11 +100,13 @@ def __getstate__(self):
         state = self.__dict__.copy()
         # threading.Lock cannot be pickled; recreated in __setstate__.
         state.pop("_legacy_bucket_lock", None)
+        state.pop("_s3_atomic_client", None)
         return state
 
     def __setstate__(self, state):
         self.__dict__.update(state)
         self._legacy_bucket_lock = threading.Lock()
+        self._s3_atomic_client = None
 
     @staticmethod
     def parse_location(location: str):
@@ -578,6 +582,29 @@ def delete_directory_quietly(self, directory: str):
             self.logger.warning(f"Exception occurs when deleting directory {directory}", exc_info=True)
 
     def try_to_write_atomic(self, path: str, content: str) -> bool:
+        if self._is_s3:
+            from botocore.exceptions import ClientError
+
+            uri = urlparse(path)
+            if uri.scheme and uri.scheme not in {"s3", "s3a", "s3n"}:
+                raise ValueError(f"Object path uses a different filesystem: {path}")
+            path_str = self.to_filesystem_path(path)
+            bucket, _, key = path_str.partition("/")
+            if not bucket or not key or key == ".":
+                raise ValueError(f"Invalid S3 object path: {path}")
+            if self._get_file_info(path_str).type == pafs.FileType.Directory:
+                return False
+
+            client = self._get_s3_atomic_client()
+            try:
+                client.put_object(Bucket=bucket, Key=key, Body=content.encode("utf-8"))
+                return True
+            except ClientError as error:
+                if error.response.get("Error", {}).get("Code") in (
+                        "PreconditionFailed", "ConditionalRequestConflict"):
+                    return False
+                raise
+
         if self.exists(path):
             path_str = self.to_filesystem_path(path)
             file_info = self._get_file_info(path_str)
@@ -594,6 +621,47 @@ def try_to_write_atomic(self, path: str, content: str) -> bool:
                 self.delete_quietly(temp_path)
         return success
 
+    def _get_s3_atomic_client(self):
+        if self._s3_atomic_client is not None:
+            return self._s3_atomic_client
+
+        import boto3
+        from botocore.config import Config
+
+        access_key = self._get_property(
+            S3Options.S3_ACCESS_KEY_ID.key(),
+            *self._s3_key_variants("access-key", "access.key"))
+        secret_key = self._get_property(
+            S3Options.S3_ACCESS_KEY_SECRET.key(),
+            *self._s3_key_variants("secret-key", "secret.key"))
+        session_token = self._get_property(
+            S3Options.S3_SECURITY_TOKEN.key(),
+            *self._s3_key_variants("session-token", "session.token",
+                                   "security-token", "security.token"))
+        endpoint = self._get_s3_property("endpoint", S3Options.S3_ENDPOINT.key())
+        region = self._get_s3_property("region", S3Options.S3_REGION.key())
+        path_style = (self._get_s3_boolean_property("path-style-access") or
+                      self._get_s3_boolean_property("path.style.access"))
+        addressing_style = "path" if path_style else (
+            "virtual" if self._pyarrow_gte_16 else "auto")
+
+        if endpoint and "://" not in endpoint:
+            endpoint = "https://" + endpoint
+        config_kwargs = {"s3": {"addressing_style": addressing_style}}
+        if "request_checksum_calculation" in Config.OPTION_DEFAULTS:
+            config_kwargs["request_checksum_calculation"] = "when_required"
+        client = boto3.session.Session().client(
+            "s3", endpoint_url=endpoint, aws_access_key_id=access_key,
+            aws_secret_access_key=secret_key, aws_session_token=session_token,
+            region_name=region, config=Config(**config_kwargs))
+        client.meta.events.register(
+            "before-sign.s3.PutObject", self._add_atomic_write_header)
+        self._s3_atomic_client = client
+        return client
+
+    def _add_atomic_write_header(self, request, **kwargs):
+        request.headers["If-None-Match"] = "*"
+
     def copy_file(self, source_path: str, target_path: str, overwrite: bool = False):
         if not overwrite and self.exists(target_path):
             raise FileExistsError(f"Target file {target_path} already exists and overwrite=False")

paimon-python/pypaimon/tests/s3_atomic_write_test.py

@@ -0,0 +1,163 @@
+# 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.
+
+import os
+import socketserver
+import threading
+import unittest
+from http.server import BaseHTTPRequestHandler, HTTPServer
+from unittest import mock
+from urllib.parse import urlsplit
+
+import pyarrow.fs as pafs
+
+try:
+    import boto3
+    from botocore.exceptions import ClientError
+except ImportError:
+    boto3 = None
+    ClientError = Exception
+
+from pypaimon.common.options import Options
+from pypaimon.common.options.config import S3Options
+from pypaimon.filesystem.pyarrow_file_io import PyArrowFileIO
+
+
+class _Server(socketserver.ThreadingMixIn, HTTPServer):
+    daemon_threads = True
+
+
+class _ConditionalPutHandler(BaseHTTPRequestHandler):
+    protocol_version = "HTTP/1.1"
+
+    def _respond(self, status, body=b""):
+        self.send_response(status)
+        self.send_header("Content-Type", "application/xml")
+        self.send_header("Content-Length", str(len(body)))
+        self.end_headers()
+        self.wfile.write(body)
+
+    def do_PUT(self):
+        body = self.rfile.read(int(self.headers["Content-Length"]))
+        if self.server.deny_put:
+            return self._respond(403, b"<Error><Code>AccessDenied</Code></Error>")
+        _, bucket, key = urlsplit(self.path).path.split("/", 2)
+        with self.server.lock:
+            old = self.server.objects.get((bucket, key))
+            if self.headers.get("If-None-Match") == "*" and old is not None:
+                return self._respond(412, b"<Error><Code>PreconditionFailed</Code></Error>")
+            self.server.objects[(bucket, key)] = body
+            self.server.headers_seen.append(dict(self.headers))
+        self._respond(200)
+
+    def log_message(self, *args):
+        pass
+
+
+@unittest.skipUnless(boto3 is not None, "requires boto3")
+class S3AtomicWriteTest(unittest.TestCase):
+    def setUp(self):
+        self.server = _Server(("127.0.0.1", 0), _ConditionalPutHandler)
+        self.server.objects = {}
+        self.server.headers_seen = []
+        self.server.deny_put = False
+        self.server.lock = threading.Lock()
+        self.thread = threading.Thread(target=self.server.serve_forever)
+        self.thread.start()
+        self.endpoint = "http://127.0.0.1:{}".format(self.server.server_port)
+        self.env = mock.patch.dict(os.environ, {
+            "NO_PROXY": "127.0.0.1,localhost",
+            "no_proxy": "127.0.0.1,localhost",
+        })
+        self.env.start()
+
+    def tearDown(self):
+        self.env.stop()
+        self.server.shutdown()
+        self.server.server_close()
+        self.thread.join()
+
+    def _s3_io(self):
+        options = Options({
+            S3Options.S3_ACCESS_KEY_ID.key(): "ak",
+            S3Options.S3_ACCESS_KEY_SECRET.key(): "sk",
+            S3Options.S3_ENDPOINT.key(): self.endpoint,
+            S3Options.S3_REGION.key(): "us-east-1",
+            "fs.s3.path.style.access": "true",
+        })
+        with mock.patch.object(PyArrowFileIO, "_initialize_s3_fs", return_value=mock.Mock()):
+            io = PyArrowFileIO("s3://test-bucket/", options)
+        io.filesystem = mock.Mock(spec=pafs.S3FileSystem)
+        io.filesystem.get_file_info.return_value = [
+            pafs.FileInfo("test-bucket/table", pafs.FileType.NotFound)]
+        return io
+
+    def test_concurrent_s3_writers_publish_only_one_value(self):
+        io = self._s3_io()
+        results = []
+        barrier = threading.Barrier(2)
+
+        def write(value):
+            barrier.wait()
+            results.append(io.try_to_write_atomic(
+                "s3://test-bucket/table/snapshot/snapshot-1", value))
+
+        threads = [threading.Thread(target=write, args=(value,))
+                   for value in ("first", "second")]
+        for thread in threads:
+            thread.start()
+        for thread in threads:
+            thread.join(10)
+
+        self.assertFalse(any(thread.is_alive() for thread in threads))
+        self.assertCountEqual([True, False], results)
+        self.assertIn(self.server.objects[("test-bucket", "table/snapshot/snapshot-1")],
+                      (b"first", b"second"))
+        self.assertEqual("*", self.server.headers_seen[0]["If-None-Match"])
+        self.assertIn("if-none-match", self.server.headers_seen[0]["Authorization"])
+
+    def test_cross_bucket_s3_write_uses_target_bucket(self):
+        io = self._s3_io()
+        io.try_to_write_atomic("s3://other-bucket/table/schema/schema-0", "data")
+
+        self.assertEqual({("other-bucket", "table/schema/schema-0"): b"data"},
+                         self.server.objects)
+
+    def test_wrong_uri_scheme_does_not_write(self):
+        io = self._s3_io()
+        with self.assertRaisesRegex(ValueError, "different filesystem"):
+            io.try_to_write_atomic("oss://test-bucket/table/schema/schema-0", "bad")
+        self.assertEqual({}, self.server.objects)
+
+    def test_write_permission_error_does_not_look_like_conflict(self):
+        io = self._s3_io()
+        self.server.deny_put = True
+
+        with self.assertRaises(ClientError):
+            io.try_to_write_atomic(
+                "s3://test-bucket/table/snapshot/snapshot-1", "data")
+
+        self.assertEqual({}, self.server.objects)
+
+    def test_atomic_write_rejects_directory(self):
+        io = self._s3_io()
+        io.filesystem.get_file_info.return_value = [
+            pafs.FileInfo("test-bucket/table", pafs.FileType.Directory)]
+
+        self.assertFalse(io.try_to_write_atomic("s3://test-bucket/table", "data"))
+
+        self.assertEqual({}, self.server.objects)