refactor

Author: kevinjqliu

pyiceberg/io/pyarrow.py

@@ -130,8 +130,6 @@
     S3_ROLE_SESSION_NAME,
     S3_SECRET_ACCESS_KEY,
     S3_SESSION_TOKEN,
-    S3_SIGNER_ENDPOINT,
-    S3_SIGNER_URI,
     FileIO,
     InputFile,
     InputStream,
@@ -197,7 +195,11 @@
 from pyiceberg.utils.datetime import millis_to_datetime
 from pyiceberg.utils.decimal import unscaled_to_decimal
 from pyiceberg.utils.deprecated import deprecation_message
-from pyiceberg.utils.properties import get_first_property_value, properties_with_prefix, property_as_bool, property_as_int
+from pyiceberg.utils.properties import (
+    filter_properties,
+    property_as_bool,
+    property_as_int,
+)
 from pyiceberg.utils.singleton import Singleton
 from pyiceberg.utils.truncate import truncate_upper_bound_binary_string, truncate_upper_bound_text_string
 
@@ -461,51 +463,76 @@ def _process_basic_properties(
 
         return client_kwargs
 
+    def _get_first_property_value_with_tracking(self, props: Properties, used_keys: set[str], *keys: str) -> Optional[Any]:
+        """Tracks all candidate keys and returns the first value found."""
+        used_keys.update(keys)
+        for key in keys:
+            if key in props:
+                return props[key]
+        return None
+
     def _initialize_oss_fs(self) -> FileSystem:
         from pyarrow.fs import S3FileSystem
 
+        properties = filter_properties(self.properties, key_predicate=lambda k: k.startswith(("oss.", "client.")))
+        used_keys: set[str] = set()
         client_kwargs = {}
-        if endpoint := get_first_property_value(self.properties, S3_ENDPOINT, "oss.endpoint_override"):
+
+        get = lambda *keys: self._get_first_property_value_with_tracking(properties, used_keys, *keys)  # noqa: E731
+
+        if endpoint := get(S3_ENDPOINT, "oss.endpoint_override"):
             client_kwargs["endpoint_override"] = endpoint
-        if access_key := get_first_property_value(self.properties, S3_ACCESS_KEY_ID, AWS_ACCESS_KEY_ID, "oss.access_key"):
+        if access_key := get(S3_ACCESS_KEY_ID, AWS_ACCESS_KEY_ID, "oss.access_key"):
             client_kwargs["access_key"] = access_key
-        if secret_key := get_first_property_value(self.properties, S3_SECRET_ACCESS_KEY, AWS_SECRET_ACCESS_KEY, "oss.secret_key"):
+        if secret_key := get(S3_SECRET_ACCESS_KEY, AWS_SECRET_ACCESS_KEY, "oss.secret_key"):
             client_kwargs["secret_key"] = secret_key
-        if session_token := get_first_property_value(self.properties, S3_SESSION_TOKEN, AWS_SESSION_TOKEN, "oss.session_token"):
+        if session_token := get(S3_SESSION_TOKEN, AWS_SESSION_TOKEN, "oss.session_token"):
             client_kwargs["session_token"] = session_token
-        if region := get_first_property_value(self.properties, S3_REGION, AWS_REGION, "oss.region"):
+        if region := get(S3_REGION, AWS_REGION, "oss.region"):
             client_kwargs["region"] = region
-        # Check for force_virtual_addressing in order of preference, defaulting to True if not found
-        if force_virtual_addressing := get_first_property_value(
-            self.properties, S3_FORCE_VIRTUAL_ADDRESSING, "oss.force_virtual_addressing"
-        ):
+        # Check for force_virtual_addressing in order of preference. For oss FS, defaulting to True if not found
+        if force_virtual_addressing := get(S3_FORCE_VIRTUAL_ADDRESSING, "oss.force_virtual_addressing"):
             if isinstance(force_virtual_addressing, str):  # S3_FORCE_VIRTUAL_ADDRESSING's value can be a string
                 force_virtual_addressing = strtobool(force_virtual_addressing)
             client_kwargs["force_virtual_addressing"] = force_virtual_addressing
         else:
             client_kwargs["force_virtual_addressing"] = True
-        if proxy_uri := get_first_property_value(self.properties, S3_PROXY_URI, "oss.proxy_options"):
+        if proxy_uri := get(S3_PROXY_URI, "oss.proxy_options"):
             client_kwargs["proxy_options"] = proxy_uri
-        if connect_timeout := get_first_property_value(self.properties, S3_CONNECT_TIMEOUT, "oss.connect_timeout"):
+        if connect_timeout := get(S3_CONNECT_TIMEOUT, "oss.connect_timeout"):
             client_kwargs["connect_timeout"] = float(connect_timeout)
-        if request_timeout := get_first_property_value(self.properties, S3_REQUEST_TIMEOUT, "oss.request_timeout"):
+        if request_timeout := get(S3_REQUEST_TIMEOUT, "oss.request_timeout"):
             client_kwargs["request_timeout"] = float(request_timeout)
-        if role_arn := get_first_property_value(self.properties, S3_ROLE_ARN, AWS_ROLE_ARN, "oss.role_arn"):
+        if role_arn := get(S3_ROLE_ARN, AWS_ROLE_ARN, "oss.role_arn"):
             client_kwargs["role_arn"] = role_arn
-        if session_name := get_first_property_value(
-            self.properties, S3_ROLE_SESSION_NAME, AWS_ROLE_SESSION_NAME, "oss.session_name"
-        ):
+        if session_name := get(S3_ROLE_SESSION_NAME, AWS_ROLE_SESSION_NAME, "oss.session_name"):
             client_kwargs["session_name"] = session_name
 
-        oss_properties = properties_with_prefix(self.properties, prefix="oss.")
-        client_kwargs = {**oss_properties, **client_kwargs}
+        remaining_oss_props = {
+            k.removeprefix("oss."): v for k, v in self.properties.items() if k.startswith("oss.") and k not in used_keys
+        }
+        client_kwargs = {**remaining_oss_props, **client_kwargs}
         return S3FileSystem(**client_kwargs)
 
     def _initialize_s3_fs(self, netloc: Optional[str]) -> FileSystem:
         from pyarrow.fs import S3FileSystem
 
-        provided_region = get_first_property_value(self.properties, S3_REGION, AWS_REGION)
+        properties = filter_properties(self.properties, key_predicate=lambda k: k.startswith(("s3.", "client.")))
+        used_keys: set[str] = set()
+        client_kwargs = {}
+
+        get = lambda *keys: self._get_first_property_value_with_tracking(properties, used_keys, *keys)  # noqa: E731
 
+        if endpoint := get(S3_ENDPOINT, "s3.endpoint_override"):
+            client_kwargs["endpoint_override"] = endpoint
+        if access_key := get(S3_ACCESS_KEY_ID, AWS_ACCESS_KEY_ID, "s3.access_key"):
+            client_kwargs["access_key"] = access_key
+        if secret_key := get(S3_SECRET_ACCESS_KEY, AWS_SECRET_ACCESS_KEY, "s3.secret_key"):
+            client_kwargs["secret_key"] = secret_key
+        if session_token := get(S3_SESSION_TOKEN, AWS_SESSION_TOKEN, "s3.session_token"):
+            client_kwargs["session_token"] = session_token
+
+        provided_region = get(S3_REGION, AWS_REGION)
         # Do this when we don't provide the region at all, or when we explicitly enable it
         if provided_region is None or property_as_bool(self.properties, S3_RESOLVE_REGION, False) is True:
             # Resolve region from netloc(bucket), fallback to user-provided region
@@ -518,66 +545,34 @@ def _initialize_s3_fs(self, netloc: Optional[str]) -> FileSystem:
                 )
         else:
             bucket_region = provided_region
-
-        # Mapping from PyIceberg properties to S3FileSystem parameter names
-        property_mapping = {
-            S3_ENDPOINT: "endpoint_override",
-            S3_PROXY_URI: "proxy_options",
-            S3_CONNECT_TIMEOUT: "connect_timeout",
-            S3_REQUEST_TIMEOUT: "request_timeout",
-            S3_RETRY_STRATEGY_IMPL: "retry_strategy",
-        }
-
-        # Properties that need special handling
-        special_properties = {
-            S3_ACCESS_KEY_ID,
-            S3_SECRET_ACCESS_KEY,
-            S3_SESSION_TOKEN,
-            S3_ROLE_ARN,
-            S3_ROLE_SESSION_NAME,
-            S3_RESOLVE_REGION,
-            S3_REGION,
-            S3_RETRY_STRATEGY_IMPL,
-            S3_CONNECT_TIMEOUT,
-            S3_REQUEST_TIMEOUT,
-            S3_SIGNER_ENDPOINT,
-            S3_SIGNER_URI,
-            S3_FORCE_VIRTUAL_ADDRESSING,
-        }
-
-        client_kwargs = self._process_basic_properties(property_mapping, special_properties, "s3")
         client_kwargs["region"] = bucket_region
+        used_keys.add(S3_RESOLVE_REGION)
 
-        if S3_ACCESS_KEY_ID in self.properties or AWS_ACCESS_KEY_ID in self.properties:
-            client_kwargs["access_key"] = get_first_property_value(self.properties, S3_ACCESS_KEY_ID, AWS_ACCESS_KEY_ID)
-
-        if S3_SECRET_ACCESS_KEY in self.properties or AWS_SECRET_ACCESS_KEY in self.properties:
-            client_kwargs["secret_key"] = get_first_property_value(self.properties, S3_SECRET_ACCESS_KEY, AWS_SECRET_ACCESS_KEY)
-
-        if S3_SESSION_TOKEN in self.properties or AWS_SESSION_TOKEN in self.properties:
-            client_kwargs["session_token"] = get_first_property_value(self.properties, S3_SESSION_TOKEN, AWS_SESSION_TOKEN)
-
-        if S3_ROLE_ARN in self.properties or AWS_ROLE_ARN in self.properties:
-            client_kwargs["role_arn"] = get_first_property_value(self.properties, S3_ROLE_ARN, AWS_ROLE_ARN)
-
-        if S3_ROLE_SESSION_NAME in self.properties or AWS_ROLE_SESSION_NAME in self.properties:
-            client_kwargs["session_name"] = get_first_property_value(self.properties, S3_ROLE_SESSION_NAME, AWS_ROLE_SESSION_NAME)
+        if force_virtual_addressing := get(S3_FORCE_VIRTUAL_ADDRESSING, "s3.force_virtual_addressing"):
+            if isinstance(force_virtual_addressing, str):  # S3_FORCE_VIRTUAL_ADDRESSING's value can be a string
+                force_virtual_addressing = strtobool(force_virtual_addressing)
+            client_kwargs["force_virtual_addressing"] = force_virtual_addressing
 
-        if connect_timeout := self.properties.get(S3_CONNECT_TIMEOUT):
+        if proxy_uri := get(S3_PROXY_URI, "s3.proxy_options"):
+            client_kwargs["proxy_options"] = proxy_uri
+        if connect_timeout := get(S3_CONNECT_TIMEOUT, "s3.connect_timeout"):
             client_kwargs["connect_timeout"] = float(connect_timeout)
-
-        if request_timeout := self.properties.get(S3_REQUEST_TIMEOUT):
+        if request_timeout := get(S3_REQUEST_TIMEOUT, "s3.request_timeout"):
             client_kwargs["request_timeout"] = float(request_timeout)
-
-        if self.properties.get(S3_FORCE_VIRTUAL_ADDRESSING) is not None:
-            client_kwargs["force_virtual_addressing"] = property_as_bool(self.properties, S3_FORCE_VIRTUAL_ADDRESSING, False)
+        if role_arn := get(S3_ROLE_ARN, AWS_ROLE_ARN, "s3.role_arn"):
+            client_kwargs["role_arn"] = role_arn
+        if session_name := get(S3_ROLE_SESSION_NAME, AWS_ROLE_SESSION_NAME, "s3.session_name"):
+            client_kwargs["session_name"] = session_name
 
         # Handle retry strategy special case
-        if (retry_strategy_impl := self.properties.get(S3_RETRY_STRATEGY_IMPL)) and (
-            retry_instance := _import_retry_strategy(retry_strategy_impl)
-        ):
-            client_kwargs["retry_strategy"] = retry_instance
+        if retry_strategy_impl := get(S3_RETRY_STRATEGY_IMPL, "s3.retry_strategy"):
+            if retry_instance := _import_retry_strategy(retry_strategy_impl):
+                client_kwargs["retry_strategy"] = retry_instance
 
+        remaining_s3_props = {
+            k.removeprefix("s3."): v for k, v in self.properties.items() if k.startswith("s3.") and k not in used_keys
+        }
+        client_kwargs = {**remaining_s3_props, **client_kwargs}
         return S3FileSystem(**client_kwargs)
 
     def _initialize_azure_fs(self) -> FileSystem: