fix: raise error when creating a table during a transaction

Author: felixscherz

pyiceberg/catalog/s3tables.py

@@ -35,7 +35,7 @@
 from pyiceberg.partitioning import UNPARTITIONED_PARTITION_SPEC, PartitionSpec
 from pyiceberg.schema import Schema
 from pyiceberg.serializers import FromInputFile
-from pyiceberg.table import CommitTableResponse, Table
+from pyiceberg.table import CommitTableResponse, CreateTableTransaction, Table
 from pyiceberg.table.metadata import new_table_metadata
 from pyiceberg.table.sorting import UNSORTED_SORT_ORDER, SortOrder
 from pyiceberg.table.update import TableRequirement, TableUpdate
@@ -92,36 +92,57 @@ def commit_table(
         table_identifier = table.name()
         database_name, table_name = self.identifier_to_database_and_table(table_identifier, NoSuchTableError)
 
-        current_table, version_token = self._load_table_and_version(identifier=table_identifier)
-
-        updated_staged_table = self._update_and_stage_table(current_table, table_identifier, requirements, updates)
-        if current_table and updated_staged_table.metadata == current_table.metadata:
-            # no changes, do nothing
-            return CommitTableResponse(metadata=current_table.metadata, metadata_location=current_table.metadata_location)
-
-        self._write_metadata(
-            metadata=updated_staged_table.metadata,
-            io=updated_staged_table.io,
-            metadata_path=updated_staged_table.metadata_location,
-            overwrite=True,
-        )
-
-        # try to update metadata location which will fail if the versionToken changed meanwhile
+        current_table: Optional[Table]
+        version_token: Optional[str]
         try:
-            self.s3tables.update_table_metadata_location(
-                tableBucketARN=self.table_bucket_arn,
-                namespace=database_name,
-                name=table_name,
-                versionToken=version_token,
-                metadataLocation=updated_staged_table.metadata_location,
+            current_table, version_token = self._load_table_and_version(identifier=table_identifier)
+        except NoSuchTableError:
+            current_table = None
+            version_token = None
+
+        if current_table:
+            updated_staged_table = self._update_and_stage_table(current_table, table_identifier, requirements, updates)
+            if updated_staged_table.metadata == current_table.metadata:
+                # no changes, do nothing
+                return CommitTableResponse(metadata=current_table.metadata, metadata_location=current_table.metadata_location)
+
+            self._write_metadata(
+                metadata=updated_staged_table.metadata,
+                io=updated_staged_table.io,
+                metadata_path=updated_staged_table.metadata_location,
+                overwrite=True,
             )
-        except self.s3tables.exceptions.ConflictException as e:
-            raise CommitFailedException(
-                f"Cannot commit {database_name}.{table_name} because of a concurrent update to the table version {version_token}."
-            ) from e
-        return CommitTableResponse(
-            metadata=updated_staged_table.metadata, metadata_location=updated_staged_table.metadata_location
-        )
+
+            # try to update metadata location which will fail if the versionToken changed meanwhile
+            try:
+                self.s3tables.update_table_metadata_location(
+                    tableBucketARN=self.table_bucket_arn,
+                    namespace=database_name,
+                    name=table_name,
+                    versionToken=version_token,
+                    metadataLocation=updated_staged_table.metadata_location,
+                )
+            except self.s3tables.exceptions.ConflictException as e:
+                raise CommitFailedException(
+                    f"Cannot commit {database_name}.{table_name} because of a concurrent update to the table version {version_token}."
+                ) from e
+            return CommitTableResponse(
+                metadata=updated_staged_table.metadata, metadata_location=updated_staged_table.metadata_location
+            )
+        else:
+            # table does not exist, create it
+            raise NotImplementedError("Creating a table on commit is currently not supported.")
+
+    def create_table_transaction(
+        self,
+        identifier: Union[str, Identifier],
+        schema: Union[Schema, "pa.Schema"],
+        location: Optional[str] = None,
+        partition_spec: PartitionSpec = UNPARTITIONED_PARTITION_SPEC,
+        sort_order: SortOrder = UNSORTED_SORT_ORDER,
+        properties: Properties = EMPTY_DICT,
+    ) -> CreateTableTransaction:
+        raise NotImplementedError("create_table_transaction currently not supported.")
 
     def create_namespace(self, namespace: Union[str, Identifier], properties: Properties = EMPTY_DICT) -> None:
         if properties:

tests/catalog/integration_test_s3tables.py

@@ -18,7 +18,9 @@
 
 from pyiceberg.catalog.s3tables import S3TablesCatalog
 from pyiceberg.exceptions import NoSuchNamespaceError, NoSuchTableError, TableBucketNotFound
+from pyiceberg.partitioning import PartitionField, PartitionSpec
 from pyiceberg.schema import Schema
+from pyiceberg.transforms import IdentityTransform
 from pyiceberg.types import IntegerType
 
 
@@ -247,3 +249,37 @@ def test_commit_new_data_to_table(catalog: S3TablesCatalog, database_name: str,
     assert updated_table_metadata.metadata_log[-1].metadata_file == original_table_metadata_location
     assert updated_table_metadata.metadata_log[-1].timestamp_ms == original_table_last_updated_ms
     assert table.scan().to_arrow().num_rows == 2 * row_count
+
+
+def test_create_table_transaction(
+    catalog: S3TablesCatalog, database_name: str, table_name: str, table_schema_nested: str
+) -> None:
+    identifier = (database_name, table_name)
+    catalog.create_namespace(namespace=database_name)
+
+    with catalog.create_table_transaction(
+        identifier,
+        table_schema_nested,
+        partition_spec=PartitionSpec(PartitionField(source_id=1, field_id=1000, transform=IdentityTransform(), name="foo")),
+    ) as txn:
+        last_updated_metadata = txn.table_metadata.last_updated_ms
+        with txn.update_schema() as update_schema:
+            update_schema.add_column(path="b", field_type=IntegerType())
+
+        with txn.update_spec() as update_spec:
+            update_spec.add_identity("bar")
+
+        txn.set_properties(test_a="test_aa", test_b="test_b", test_c="test_c")
+
+    table = catalog.load_table(identifier)
+
+    assert table.schema().find_field("b").field_type == IntegerType()
+    assert table.properties == {"test_a": "test_aa", "test_b": "test_b", "test_c": "test_c"}
+    assert table.spec().last_assigned_field_id == 1001
+    assert table.spec().fields_by_source_id(1)[0].name == "foo"
+    assert table.spec().fields_by_source_id(1)[0].field_id == 1000
+    assert table.spec().fields_by_source_id(1)[0].transform == IdentityTransform()
+    assert table.spec().fields_by_source_id(2)[0].name == "bar"
+    assert table.spec().fields_by_source_id(2)[0].field_id == 1001
+    assert table.spec().fields_by_source_id(2)[0].transform == IdentityTransform()
+    assert table.metadata.last_updated_ms > last_updated_metadata