Robustify boto3 session handling (DynamoDB, RestCatalog) (apache#2071)

3 changes here...

1. Basically a mirror of
apache#1920 but for dynamo
2. Allow users to pass in an existing client/pass initialize client with
properties in `RestCatalog` with sigv4 enabled. (closes apache#2070, closes
apache#2008)
3. Re-use of the client in `add_headers` in `RestCatalog` with sigv4.
(closes apache#2069 )

---------

Co-authored-by: Marko Grujic <[email protected]>

Author: jayceslesar

poetry.lock

@@ -66,6 +66,8 @@
 
 if TYPE_CHECKING:
     import pyarrow as pa
+    from mypy_boto3_dynamodb.client import DynamoDBClient
+
 
 DYNAMODB_CLIENT = "dynamodb"
 
@@ -94,18 +96,28 @@
 
 
 class DynamoDbCatalog(MetastoreCatalog):
-    def __init__(self, name: str, **properties: str):
+    def __init__(self, name: str, client: Optional["DynamoDBClient"] = None, **properties: str):
+        """Dynamodb catalog.
+
+        Args:
+            name: Name to identify the catalog.
+            client: An optional boto3 dynamodb client.
+            properties: Properties for dynamodb client construction and configuration.
+        """
         super().__init__(name, **properties)
+        if client is not None:
+            self.dynamodb = client
+        else:
+            session = boto3.Session(
+                profile_name=properties.get(DYNAMODB_PROFILE_NAME),
+                region_name=get_first_property_value(properties, DYNAMODB_REGION, AWS_REGION),
+                botocore_session=properties.get(BOTOCORE_SESSION),
+                aws_access_key_id=get_first_property_value(properties, DYNAMODB_ACCESS_KEY_ID, AWS_ACCESS_KEY_ID),
+                aws_secret_access_key=get_first_property_value(properties, DYNAMODB_SECRET_ACCESS_KEY, AWS_SECRET_ACCESS_KEY),
+                aws_session_token=get_first_property_value(properties, DYNAMODB_SESSION_TOKEN, AWS_SESSION_TOKEN),
+            )
+            self.dynamodb = session.client(DYNAMODB_CLIENT)
 
-        session = boto3.Session(
-            profile_name=properties.get(DYNAMODB_PROFILE_NAME),
-            region_name=get_first_property_value(properties, DYNAMODB_REGION, AWS_REGION),
-            botocore_session=properties.get(BOTOCORE_SESSION),
-            aws_access_key_id=get_first_property_value(properties, DYNAMODB_ACCESS_KEY_ID, AWS_ACCESS_KEY_ID),
-            aws_secret_access_key=get_first_property_value(properties, DYNAMODB_SECRET_ACCESS_KEY, AWS_SECRET_ACCESS_KEY),
-            aws_session_token=get_first_property_value(properties, DYNAMODB_SESSION_TOKEN, AWS_SESSION_TOKEN),
-        )
-        self.dynamodb = session.client(DYNAMODB_CLIENT)
         self.dynamodb_table_name = self.properties.get(DYNAMODB_TABLE_NAME, DYNAMODB_TABLE_NAME_DEFAULT)
         self._ensure_catalog_table_exists_or_create()
 

pyiceberg/catalog/dynamodb.py

@@ -30,15 +30,6 @@
 
 import boto3
 from botocore.config import Config
-from mypy_boto3_glue.client import GlueClient
-from mypy_boto3_glue.type_defs import (
-    ColumnTypeDef,
-    DatabaseInputTypeDef,
-    DatabaseTypeDef,
-    StorageDescriptorTypeDef,
-    TableInputTypeDef,
-    TableTypeDef,
-)
 
 from pyiceberg.catalog import (
     BOTOCORE_SESSION,
@@ -101,6 +92,15 @@
 
 if TYPE_CHECKING:
     import pyarrow as pa
+    from mypy_boto3_glue.client import GlueClient
+    from mypy_boto3_glue.type_defs import (
+        ColumnTypeDef,
+        DatabaseInputTypeDef,
+        DatabaseTypeDef,
+        StorageDescriptorTypeDef,
+        TableInputTypeDef,
+        TableTypeDef,
+    )
 
 
 # There is a unique Glue metastore in each AWS account and each AWS region. By default, GlueCatalog chooses the Glue
@@ -140,7 +140,7 @@
 
 
 def _construct_parameters(
-    metadata_location: str, glue_table: Optional[TableTypeDef] = None, prev_metadata_location: Optional[str] = None
+    metadata_location: str, glue_table: Optional["TableTypeDef"] = None, prev_metadata_location: Optional[str] = None
 ) -> Properties:
     new_parameters = glue_table.get("Parameters", {}) if glue_table else {}
     new_parameters.update({TABLE_TYPE: ICEBERG.upper(), METADATA_LOCATION: metadata_location})
@@ -190,15 +190,15 @@ def primitive(self, primitive: PrimitiveType) -> str:
         return GLUE_PRIMITIVE_TYPES[primitive_type]
 
 
-def _to_columns(metadata: TableMetadata) -> List[ColumnTypeDef]:
-    results: Dict[str, ColumnTypeDef] = {}
+def _to_columns(metadata: TableMetadata) -> List["ColumnTypeDef"]:
+    results: Dict[str, "ColumnTypeDef"] = {}
 
     def _append_to_results(field: NestedField, is_current: bool) -> None:
         if field.name in results:
             return
 
         results[field.name] = cast(
-            ColumnTypeDef,
+            "ColumnTypeDef",
             {
                 "Name": field.name,
                 "Type": visit(field.field_type, _IcebergSchemaToGlueType()),
@@ -230,10 +230,10 @@ def _construct_table_input(
     metadata_location: str,
     properties: Properties,
     metadata: TableMetadata,
-    glue_table: Optional[TableTypeDef] = None,
+    glue_table: Optional["TableTypeDef"] = None,
     prev_metadata_location: Optional[str] = None,
-) -> TableInputTypeDef:
-    table_input: TableInputTypeDef = {
+) -> "TableInputTypeDef":
+    table_input: "TableInputTypeDef" = {
         "Name": table_name,
         "TableType": EXTERNAL_TABLE,
         "Parameters": _construct_parameters(metadata_location, glue_table, prev_metadata_location),
@@ -249,8 +249,8 @@ def _construct_table_input(
     return table_input
 
 
-def _construct_rename_table_input(to_table_name: str, glue_table: TableTypeDef) -> TableInputTypeDef:
-    rename_table_input: TableInputTypeDef = {"Name": to_table_name}
+def _construct_rename_table_input(to_table_name: str, glue_table: "TableTypeDef") -> "TableInputTypeDef":
+    rename_table_input: "TableInputTypeDef" = {"Name": to_table_name}
     # use the same Glue info to create the new table, pointing to the old metadata
     assert glue_table["TableType"]
     rename_table_input["TableType"] = glue_table["TableType"]
@@ -264,16 +264,16 @@ def _construct_rename_table_input(to_table_name: str, glue_table: TableTypeDef)
         # It turns out the output of StorageDescriptor is not the same as the input type
         # because the Column can have a different type, but for now it seems to work, so
         # silence the type error.
-        rename_table_input["StorageDescriptor"] = cast(StorageDescriptorTypeDef, glue_table["StorageDescriptor"])
+        rename_table_input["StorageDescriptor"] = cast("StorageDescriptorTypeDef", glue_table["StorageDescriptor"])
 
     if "Description" in glue_table:
         rename_table_input["Description"] = glue_table["Description"]
 
     return rename_table_input
 
 
-def _construct_database_input(database_name: str, properties: Properties) -> DatabaseInputTypeDef:
-    database_input: DatabaseInputTypeDef = {"Name": database_name}
+def _construct_database_input(database_name: str, properties: Properties) -> "DatabaseInputTypeDef":
+    database_input: "DatabaseInputTypeDef" = {"Name": database_name}
     parameters = {}
     for k, v in properties.items():
         if k == "Description":
@@ -286,7 +286,7 @@ def _construct_database_input(database_name: str, properties: Properties) -> Dat
     return database_input
 
 
-def _register_glue_catalog_id_with_glue_client(glue: GlueClient, glue_catalog_id: str) -> None:
+def _register_glue_catalog_id_with_glue_client(glue: "GlueClient", glue_catalog_id: str) -> None:
     """
     Register the Glue Catalog ID (AWS Account ID) as a parameter on all Glue client methods.
 
@@ -303,9 +303,9 @@ def add_glue_catalog_id(params: Dict[str, str], **kwargs: Any) -> None:
 
 
 class GlueCatalog(MetastoreCatalog):
-    glue: GlueClient
+    glue: "GlueClient"
 
-    def __init__(self, name: str, client: Optional[GlueClient] = None, **properties: Any):
+    def __init__(self, name: str, client: Optional["GlueClient"] = None, **properties: Any):
         """Glue Catalog.
 
         You either need to provide a boto3 glue client, or one will be constructed from the properties.
@@ -317,7 +317,7 @@ def __init__(self, name: str, client: Optional[GlueClient] = None, **properties:
         """
         super().__init__(name, **properties)
 
-        if client:
+        if client is not None:
             self.glue = client
         else:
             retry_mode_prop_value = get_first_property_value(properties, GLUE_RETRY_MODE)
@@ -344,7 +344,7 @@ def __init__(self, name: str, client: Optional[GlueClient] = None, **properties:
             if glue_catalog_id := properties.get(GLUE_ID):
                 _register_glue_catalog_id_with_glue_client(self.glue, glue_catalog_id)
 
-    def _convert_glue_to_iceberg(self, glue_table: TableTypeDef) -> Table:
+    def _convert_glue_to_iceberg(self, glue_table: "TableTypeDef") -> Table:
         properties: Properties = glue_table["Parameters"]
 
         assert glue_table["DatabaseName"]
@@ -380,15 +380,15 @@ def _convert_glue_to_iceberg(self, glue_table: TableTypeDef) -> Table:
             catalog=self,
         )
 
-    def _create_glue_table(self, database_name: str, table_name: str, table_input: TableInputTypeDef) -> None:
+    def _create_glue_table(self, database_name: str, table_name: str, table_input: "TableInputTypeDef") -> None:
         try:
             self.glue.create_table(DatabaseName=database_name, TableInput=table_input)
         except self.glue.exceptions.AlreadyExistsException as e:
             raise TableAlreadyExistsError(f"Table {database_name}.{table_name} already exists") from e
         except self.glue.exceptions.EntityNotFoundException as e:
             raise NoSuchNamespaceError(f"Database {database_name} does not exist") from e
 
-    def _update_glue_table(self, database_name: str, table_name: str, table_input: TableInputTypeDef, version_id: str) -> None:
+    def _update_glue_table(self, database_name: str, table_name: str, table_input: "TableInputTypeDef", version_id: str) -> None:
         try:
             self.glue.update_table(
                 DatabaseName=database_name,
@@ -403,7 +403,7 @@ def _update_glue_table(self, database_name: str, table_name: str, table_input: T
                 f"Cannot commit {database_name}.{table_name} because Glue detected concurrent update to table version {version_id}"
             ) from e
 
-    def _get_glue_table(self, database_name: str, table_name: str) -> TableTypeDef:
+    def _get_glue_table(self, database_name: str, table_name: str) -> "TableTypeDef":
         try:
             load_table_response = self.glue.get_table(DatabaseName=database_name, Name=table_name)
             return load_table_response["Table"]
@@ -496,7 +496,7 @@ def commit_table(
         table_identifier = table.name()
         database_name, table_name = self.identifier_to_database_and_table(table_identifier, NoSuchTableError)
 
-        current_glue_table: Optional[TableTypeDef]
+        current_glue_table: Optional["TableTypeDef"]
         glue_table_version_id: Optional[str]
         current_table: Optional[Table]
         try:
@@ -702,7 +702,7 @@ def list_tables(self, namespace: Union[str, Identifier]) -> List[Identifier]:
             NoSuchNamespaceError: If a namespace with the given name does not exist, or the identifier is invalid.
         """
         database_name = self.identifier_to_database(namespace, NoSuchNamespaceError)
-        table_list: List[TableTypeDef] = []
+        table_list: List["TableTypeDef"] = []
         next_token: Optional[str] = None
         try:
             while True:
@@ -730,7 +730,7 @@ def list_namespaces(self, namespace: Union[str, Identifier] = ()) -> List[Identi
         if namespace:
             return []
 
-        database_list: List[DatabaseTypeDef] = []
+        database_list: List["DatabaseTypeDef"] = []
         next_token: Optional[str] = None
 
         while True:
@@ -806,5 +806,5 @@ def view_exists(self, identifier: Union[str, Identifier]) -> bool:
         raise NotImplementedError
 
     @staticmethod
-    def __is_iceberg_table(table: TableTypeDef) -> bool:
+    def __is_iceberg_table(table: "TableTypeDef") -> bool:
         return table.get("Parameters", {}).get(TABLE_TYPE, "").lower() == ICEBERG

pyiceberg/catalog/glue.py

@@ -32,6 +32,7 @@
 
 from pyiceberg import __version__
 from pyiceberg.catalog import (
+    BOTOCORE_SESSION,
     TOKEN,
     URI,
     WAREHOUSE_LOCATION,
@@ -53,6 +54,7 @@
     TableAlreadyExistsError,
     UnauthorizedError,
 )
+from pyiceberg.io import AWS_ACCESS_KEY_ID, AWS_REGION, AWS_SECRET_ACCESS_KEY, AWS_SESSION_TOKEN
 from pyiceberg.partitioning import UNPARTITIONED_PARTITION_SPEC, PartitionSpec, assign_fresh_partition_spec_ids
 from pyiceberg.schema import Schema, assign_fresh_schema_ids
 from pyiceberg.table import (
@@ -72,7 +74,7 @@
 from pyiceberg.typedef import EMPTY_DICT, UTF8, IcebergBaseModel, Identifier, Properties
 from pyiceberg.types import transform_dict_value_to_str
 from pyiceberg.utils.deprecated import deprecation_message
-from pyiceberg.utils.properties import get_header_properties, property_as_bool
+from pyiceberg.utils.properties import get_first_property_value, get_header_properties, property_as_bool
 
 if TYPE_CHECKING:
     import pyarrow as pa
@@ -390,11 +392,17 @@ class SigV4Adapter(HTTPAdapter):
             def __init__(self, **properties: str):
                 super().__init__()
                 self._properties = properties
+                self._boto_session = boto3.Session(
+                    region_name=get_first_property_value(self._properties, AWS_REGION),
+                    botocore_session=self._properties.get(BOTOCORE_SESSION),
+                    aws_access_key_id=get_first_property_value(self._properties, AWS_ACCESS_KEY_ID),
+                    aws_secret_access_key=get_first_property_value(self._properties, AWS_SECRET_ACCESS_KEY),
+                    aws_session_token=get_first_property_value(self._properties, AWS_SESSION_TOKEN),
+                )
 
             def add_headers(self, request: PreparedRequest, **kwargs: Any) -> None:  # pylint: disable=W0613
-                boto_session = boto3.Session()
-                credentials = boto_session.get_credentials().get_frozen_credentials()
-                region = self._properties.get(SIGV4_REGION, boto_session.region_name)
+                credentials = self._boto_session.get_credentials().get_frozen_credentials()
+                region = self._properties.get(SIGV4_REGION, self._boto_session.region_name)
                 service = self._properties.get(SIGV4_SERVICE, "execute-api")
 
                 url = str(request.url).split("?")[0]

pyiceberg/catalog/rest/__init__.py

@@ -67,11 +67,10 @@ pandas = { version = ">=1.0.0,<3.0.0", optional = true }
 duckdb = { version = ">=0.5.0,<2.0.0", optional = true }
 ray = [
   { version = "==2.10.0", python = "<3.9", optional = true },
-  { version = ">=2.10.0,<3.0.0", python = ">=3.9", optional = true },
+  { version = ">=2.10.0,<=2.44.0", python = ">=3.9", optional = true },
 ]
 python-snappy = { version = ">=0.6.0,<1.0.0", optional = true }
 thrift = { version = ">=0.13.0,<1.0.0", optional = true }
-mypy-boto3-glue = { version = ">=1.28.18", optional = true }
 boto3 = { version = ">=1.24.59", optional = true }
 s3fs = { version = ">=2023.1.0", optional = true }
 adlfs = { version = ">=2023.1.0", optional = true }
@@ -102,6 +101,8 @@ cython = "3.1.2"
 deptry = ">=0.14,<0.24"
 datafusion = ">=44,<48"
 docutils = "!=0.21.post1"   # https://github.com/python-poetry/poetry/issues/9248#issuecomment-2026240520
+mypy-boto3-glue = ">=1.28.18"
+mypy-boto3-dynamodb = ">=1.28.18"
 
 [tool.poetry.group.docs.dependencies]
 # for mkdocs
@@ -217,6 +218,10 @@ ignore_missing_imports = true
 module = "mypy_boto3_glue.*"
 ignore_missing_imports = true
 
+[[tool.mypy.overrides]]
+module = "mypy_boto3_dynamodb.*"
+ignore_missing_imports = true
+
 [[tool.mypy.overrides]]
 module = "moto"
 ignore_missing_imports = true
@@ -299,7 +304,7 @@ snappy = ["python-snappy"]
 hive = ["thrift"]
 hive-kerberos = ["thrift", "thrift_sasl", "kerberos"]
 s3fs = ["s3fs"]
-glue = ["boto3", "mypy-boto3-glue"]
+glue = ["boto3"]
 adlfs = ["adlfs"]
 dynamodb = ["boto3"]
 zstandard = ["zstandard"]

pyproject.toml

@@ -626,3 +626,11 @@ def test_table_exists(
     assert test_catalog.table_exists(identifier) is True
     # Act and Assert for an non-existing  table
     assert test_catalog.table_exists(("non", "exist")) is False
+
+
+@mock_aws
+def test_dynamodb_client_override() -> None:
+    catalog_name = "glue"
+    test_client = boto3.client("dynamodb", region_name="us-west-2")
+    test_catalog = DynamoDbCatalog(catalog_name, test_client)
+    assert test_catalog.dynamodb is test_client