feat: Add AWS profile support for GlueCatalog and internal S3FileIO

Author: stats-dev

pyiceberg/catalog/glue.py

@@ -48,7 +48,7 @@
     NoSuchTableError,
     TableAlreadyExistsError,
 )
-from pyiceberg.io import AWS_ACCESS_KEY_ID, AWS_REGION, AWS_SECRET_ACCESS_KEY, AWS_SESSION_TOKEN
+from pyiceberg.io import AWS_ACCESS_KEY_ID, AWS_PROFILE_NAME, AWS_REGION, AWS_SECRET_ACCESS_KEY, AWS_SESSION_TOKEN
 from pyiceberg.partitioning import UNPARTITIONED_PARTITION_SPEC, PartitionSpec
 from pyiceberg.schema import Schema, SchemaVisitor, visit
 from pyiceberg.serializers import FromInputFile
@@ -329,7 +329,7 @@ def __init__(self, name: str, client: Optional["GlueClient"] = None, **propertie
             retry_mode_prop_value = get_first_property_value(properties, GLUE_RETRY_MODE)
 
             session = boto3.Session(
-                profile_name=properties.get(GLUE_PROFILE_NAME),
+                profile_name=get_first_property_value(properties, GLUE_PROFILE_NAME, AWS_PROFILE_NAME),
                 region_name=get_first_property_value(properties, GLUE_REGION, AWS_REGION),
                 botocore_session=properties.get(BOTOCORE_SESSION),
                 aws_access_key_id=get_first_property_value(properties, GLUE_ACCESS_KEY_ID, AWS_ACCESS_KEY_ID),

pyiceberg/io/__init__.py

@@ -41,12 +41,14 @@
 
 logger = logging.getLogger(__name__)
 
+AWS_PROFILE_NAME = "client.profile-name"
 AWS_REGION = "client.region"
 AWS_ACCESS_KEY_ID = "client.access-key-id"
 AWS_SECRET_ACCESS_KEY = "client.secret-access-key"
 AWS_SESSION_TOKEN = "client.session-token"
 AWS_ROLE_ARN = "client.role-arn"
 AWS_ROLE_SESSION_NAME = "client.role-session-name"
+S3_PROFILE_NAME = "s3.profile-name"
 S3_ANONYMOUS = "s3.anonymous"
 S3_ENDPOINT = "s3.endpoint"
 S3_ACCESS_KEY_ID = "s3.access-key-id"

pyiceberg/io/fsspec.py

@@ -51,6 +51,7 @@
     ADLS_TENANT_ID,
     ADLS_TOKEN,
     AWS_ACCESS_KEY_ID,
+    AWS_PROFILE_NAME,
     AWS_REGION,
     AWS_SECRET_ACCESS_KEY,
     AWS_SESSION_TOKEN,
@@ -71,6 +72,7 @@
     S3_CONNECT_TIMEOUT,
     S3_ENDPOINT,
     S3_FORCE_VIRTUAL_ADDRESSING,
+    S3_PROFILE_NAME,
     S3_PROXY_URI,
     S3_REGION,
     S3_REQUEST_TIMEOUT,
@@ -205,7 +207,13 @@ def _s3(properties: Properties) -> AbstractFileSystem:
     else:
         anon = False
 
-    fs = S3FileSystem(anon=anon, client_kwargs=client_kwargs, config_kwargs=config_kwargs)
+    session = None
+    if profile_name := get_first_property_value(properties, S3_PROFILE_NAME, AWS_PROFILE_NAME):
+        from aiobotocore.session import AioSession
+
+        session = AioSession(profile=profile_name)
+
+    fs = S3FileSystem(anon=anon, client_kwargs=client_kwargs, config_kwargs=config_kwargs, session=session)
 
     for event_name, event_function in register_events.items():
         fs.s3.meta.events.unregister(event_name, unique_id=1925)

tests/catalog/test_glue_profile.py

@@ -0,0 +1,51 @@
+# 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.
+
+from unittest import mock
+
+from moto import mock_aws
+
+from pyiceberg.catalog.glue import GlueCatalog
+from pyiceberg.typedef import Properties
+
+UNIFIED_AWS_SESSION_PROPERTIES = {
+    "client.access-key-id": "client.access-key-id",
+    "client.secret-access-key": "client.secret-access-key",
+    "client.region": "client.region",
+    "client.session-token": "client.session-token",
+}
+
+
+@mock_aws
+def test_passing_client_profile_name_properties_to_glue() -> None:
+    session_properties: Properties = {
+        "client.profile-name": "profile_name",
+        **UNIFIED_AWS_SESSION_PROPERTIES,
+    }
+
+    with mock.patch("boto3.Session") as mock_session:
+        test_catalog = GlueCatalog("glue", **session_properties)
+
+    mock_session.assert_called_with(
+        aws_access_key_id="client.access-key-id",
+        aws_secret_access_key="client.secret-access-key",
+        aws_session_token="client.session-token",
+        region_name="client.region",
+        profile_name="profile_name",
+        botocore_session=None,
+    )
+    assert test_catalog.glue is mock_session().client()

tests/io/test_fsspec.py

@@ -34,7 +34,13 @@
 from pyiceberg.io.fsspec import FsspecFileIO, S3V4RestSigner
 from pyiceberg.io.pyarrow import PyArrowFileIO
 from pyiceberg.typedef import Properties
-from tests.conftest import UNIFIED_AWS_SESSION_PROPERTIES
+
+UNIFIED_AWS_SESSION_PROPERTIES = {
+    "client.access-key-id": "client.access-key-id",
+    "client.secret-access-key": "client.secret-access-key",
+    "client.region": "client.region",
+    "client.session-token": "client.session-token",
+}
 
 
 def test_fsspec_infer_local_fs_from_path(fsspec_fileio: FsspecFileIO) -> None:
@@ -303,6 +309,7 @@ def test_fsspec_s3_session_properties() -> None:
                 "aws_session_token": "s3.session-token",
             },
             config_kwargs={},
+            session=None,
         )
 
 
@@ -333,6 +340,7 @@ def test_fsspec_s3_session_properties_force_virtual_addressing() -> None:
                 "aws_session_token": "s3.session-token",
             },
             config_kwargs={"s3": {"addressing_style": "virtual"}},
+            session=None,
         )
 
 
@@ -363,6 +371,7 @@ def test_fsspec_s3_session_properties_with_anonymous() -> None:
                 "aws_session_token": "s3.session-token",
             },
             config_kwargs={},
+            session=None,
         )
 
 
@@ -388,6 +397,7 @@ def test_fsspec_unified_session_properties() -> None:
                 "aws_session_token": "client.session-token",
             },
             config_kwargs={},
+            session=None,
         )
 
 

tests/io/test_fsspec_profile.py

@@ -0,0 +1,86 @@
+# 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 uuid
+from unittest import mock
+
+from pyiceberg.io.fsspec import FsspecFileIO
+from pyiceberg.typedef import Properties
+
+UNIFIED_AWS_SESSION_PROPERTIES = {
+    "client.access-key-id": "client.access-key-id",
+    "client.secret-access-key": "client.secret-access-key",
+    "client.region": "client.region",
+    "client.session-token": "client.session-token",
+}
+
+
+def test_fsspec_s3_session_properties_with_profile() -> None:
+    session_properties: Properties = {
+        "s3.profile-name": "test-profile",
+        "s3.endpoint": "http://localhost:9000",
+        **UNIFIED_AWS_SESSION_PROPERTIES,
+    }
+
+    with mock.patch("s3fs.S3FileSystem") as mock_s3fs, mock.patch("aiobotocore.session.AioSession") as mock_aio_session:
+        s3_fileio = FsspecFileIO(properties=session_properties)
+        filename = str(uuid.uuid4())
+
+        s3_fileio.new_input(location=f"s3://warehouse/{filename}")
+
+        mock_aio_session.assert_called_with(profile="test-profile")
+        mock_s3fs.assert_called_with(
+            anon=False,
+            client_kwargs={
+                "endpoint_url": "http://localhost:9000",
+                "aws_access_key_id": "client.access-key-id",
+                "aws_secret_access_key": "client.secret-access-key",
+                "region_name": "client.region",
+                "aws_session_token": "client.session-token",
+            },
+            config_kwargs={},
+            session=mock_aio_session(),
+        )
+
+
+def test_fsspec_s3_session_properties_with_client_profile() -> None:
+    session_properties: Properties = {
+        "client.profile-name": "test-profile",
+        "s3.endpoint": "http://localhost:9000",
+        **UNIFIED_AWS_SESSION_PROPERTIES,
+    }
+
+    with mock.patch("s3fs.S3FileSystem") as mock_s3fs, mock.patch("aiobotocore.session.AioSession") as mock_aio_session:
+        s3_fileio = FsspecFileIO(properties=session_properties)
+        filename = str(uuid.uuid4())
+
+        s3_fileio.new_input(location=f"s3://warehouse/{filename}")
+
+        mock_aio_session.assert_called_with(profile="test-profile")
+        mock_s3fs.assert_called_with(
+            anon=False,
+            client_kwargs={
+                "endpoint_url": "http://localhost:9000",
+                "aws_access_key_id": "client.access-key-id",
+                "aws_secret_access_key": "client.secret-access-key",
+                "region_name": "client.region",
+                "aws_session_token": "client.session-token",
+            },
+            config_kwargs={},
+            session=mock_aio_session(),
+        )