feat: Add S3 SSE configs (#3173)

Author: achasnovskiy

mkdocs/docs/configuration.md

@@ -172,6 +172,8 @@ For the FileIO there are several configuration options available:
 | s3.force-virtual-addressing | False                      | Whether to use virtual addressing of buckets. If true, then virtual addressing is always enabled. If false, then virtual addressing is only enabled if endpoint_override is empty. This can be used for non-AWS backends that only support virtual hosted-style access. |
 | s3.retry-strategy-impl      | None                       | Ability to set a custom S3 retry strategy. A full path to a class needs to be given that extends the [S3RetryStrategy](https://github.com/apache/arrow/blob/639201bfa412db26ce45e73851432018af6c945e/python/pyarrow/_s3fs.pyx#L110) base class.                         |
 | s3.anonymous                | True                       | Configure whether to use anonymous connection. If False (default), uses key/secret if configured or boto's credential resolver.                                                                                                                                         |
+| s3.server-side-encryption   | aws:kms                    | Configure server-side encryption (e.g. `AES256` or `aws:kms`). Only supported by `FsspecFileIO`.                                                                                                                                     |
+| s3.sse-kms-key-id           | alias/my-key               | Configure the SSE-KMS key id (or ARN) for multipart uploads. Only supported by `FsspecFileIO`.                                                                                                                                                               |
 
 <!-- markdown-link-check-enable-->
 

pyiceberg/io/__init__.py

@@ -80,6 +80,8 @@ def _is_local_path(path: str) -> bool:
 S3_ROLE_SESSION_NAME = "s3.role-session-name"
 S3_FORCE_VIRTUAL_ADDRESSING = "s3.force-virtual-addressing"
 S3_RETRY_STRATEGY_IMPL = "s3.retry-strategy-impl"
+S3_SERVER_SIDE_ENCRYPTION = "s3.server-side-encryption"
+S3_SSE_KMS_KEY_ID = "s3.sse-kms-key-id"
 HDFS_HOST = "hdfs.host"
 HDFS_PORT = "hdfs.port"
 HDFS_USER = "hdfs.user"

pyiceberg/io/fsspec.py

@@ -79,11 +79,13 @@
     S3_REGION,
     S3_REQUEST_TIMEOUT,
     S3_SECRET_ACCESS_KEY,
+    S3_SERVER_SIDE_ENCRYPTION,
     S3_SESSION_TOKEN,
     S3_SIGNER,
     S3_SIGNER_ENDPOINT,
     S3_SIGNER_ENDPOINT_DEFAULT,
     S3_SIGNER_URI,
+    S3_SSE_KMS_KEY_ID,
     FileIO,
     InputFile,
     InputStream,
@@ -178,6 +180,7 @@ def _s3(properties: Properties) -> AbstractFileSystem:
         "region_name": get_first_property_value(properties, S3_REGION, AWS_REGION),
     }
     config_kwargs = {}
+    s3_additional_kwargs = {}
     register_events: dict[str, Callable[[AWSRequest], None]] = {}
 
     if signer := properties.get(S3_SIGNER):
@@ -210,6 +213,12 @@ def _s3(properties: Properties) -> AbstractFileSystem:
     else:
         anon = False
 
+    if server_side_encryption := properties.get(S3_SERVER_SIDE_ENCRYPTION):
+        s3_additional_kwargs["ServerSideEncryption"] = server_side_encryption
+
+    if sse_kms_key_id := properties.get(S3_SSE_KMS_KEY_ID):
+        s3_additional_kwargs["SSEKMSKeyId"] = sse_kms_key_id
+
     s3_fs_kwargs = {
         "anon": anon,
         "client_kwargs": client_kwargs,
@@ -219,6 +228,9 @@ def _s3(properties: Properties) -> AbstractFileSystem:
     if profile_name := get_first_property_value(properties, S3_PROFILE_NAME, AWS_PROFILE_NAME):
         s3_fs_kwargs["profile"] = profile_name
 
+    if s3_additional_kwargs:
+        s3_fs_kwargs["s3_additional_kwargs"] = s3_additional_kwargs
+
     fs = S3FileSystem(**s3_fs_kwargs)
 
     for event_name, event_function in register_events.items():

tests/io/test_fsspec.py

@@ -405,6 +405,42 @@ def test_fsspec_unified_session_properties() -> None:
         )
 
 
+def test_fsspec_s3_encryption_additional_kwargs() -> None:
+    session_properties: Properties = {
+        "s3.server-side-encryption": "aws:kms",
+        "s3.sse-kms-key-id": "arn:aws:kms:us-east-1:123456789012:key/test-key",
+        **UNIFIED_AWS_SESSION_PROPERTIES,
+    }
+
+    with mock.patch("s3fs.S3FileSystem") as mock_s3fs:
+        s3_fileio = FsspecFileIO(properties=session_properties)
+        filename = str(uuid.uuid4())
+
+        s3_fileio.new_input(location=f"s3://warehouse/{filename}")
+
+        call_kwargs = mock_s3fs.call_args.kwargs
+        assert call_kwargs["s3_additional_kwargs"] == {
+            "ServerSideEncryption": "aws:kms",
+            "SSEKMSKeyId": "arn:aws:kms:us-east-1:123456789012:key/test-key",
+        }
+
+
+def test_fsspec_s3_encryption_additional_kwargs_partial() -> None:
+    session_properties: Properties = {
+        "s3.server-side-encryption": "AES256",
+        **UNIFIED_AWS_SESSION_PROPERTIES,
+    }
+
+    with mock.patch("s3fs.S3FileSystem") as mock_s3fs:
+        s3_fileio = FsspecFileIO(properties=session_properties)
+        filename = str(uuid.uuid4())
+
+        s3_fileio.new_input(location=f"s3://warehouse/{filename}")
+
+        call_kwargs = mock_s3fs.call_args.kwargs
+        assert call_kwargs["s3_additional_kwargs"] == {"ServerSideEncryption": "AES256"}
+
+
 @pytest.mark.adls
 def test_fsspec_new_input_file_adls(adls_fsspec_fileio: FsspecFileIO) -> None:
     """Test creating a new input file from an fsspec file-io"""