Support CreateTableTransaction in Glue and Rest

mkdocs/docs/api.md

@@ -165,6 +165,25 @@ catalog.create_table(
 )
 ```
 
+To create a table with some subsequent changes in a transaction:
+
+```python
+with catalog.create_table_transaction(
+    identifier="docs_example.bids",
+    schema=schema,
+    location="s3://pyiceberg",
+    partition_spec=partition_spec,
+    sort_order=sort_order,
+) as txn:
+    with txn.update_schema() as update_schema:
+        update_schema.add_column(path="new_column", field_type=StringType())
+
+    with txn.update_spec() as update_spec:
+        update_spec.add_identity("symbol")
+
+    txn.set_properties(test_a="test_aa", test_b="test_b", test_c="test_c")
+```
+
 ## Load a table
 
 ### Catalog table

pyiceberg/catalog/__init__.py

@@ -45,9 +45,11 @@
 from pyiceberg.table import (
     CommitTableRequest,
     CommitTableResponse,
+    CreateTableTransaction,
+    StagedTable,
     Table,
 )
-from pyiceberg.table.metadata import TableMetadata
+from pyiceberg.table.metadata import TableMetadata, TableMetadataV1, new_table_metadata
 from pyiceberg.table.sorting import UNSORTED_SORT_ORDER, SortOrder
 from pyiceberg.typedef import (
     EMPTY_DICT,
@@ -288,6 +290,78 @@ def __init__(self, name: str, **properties: str):
     def _load_file_io(self, properties: Properties = EMPTY_DICT, location: Optional[str] = None) -> FileIO:
         return load_file_io({**self.properties, **properties}, location)
 
+    def _create_staged_table(
+        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,
+    ) -> StagedTable:
+        """Create a table and return the table instance without committing the changes.
+
+        Args:
+            identifier (str | Identifier): Table identifier.
+            schema (Schema): Table's schema.
+            location (str | None): Location for the table. Optional Argument.
+            partition_spec (PartitionSpec): PartitionSpec for the table.
+            sort_order (SortOrder): SortOrder for the table.
+            properties (Properties): Table properties that can be a string based dictionary.
+
+        Returns:
+            Table: the created table instance.
+
+        Raises:
+            TableAlreadyExistsError: If a table with the name already exists.
+        """
+        schema: Schema = self._convert_schema_if_needed(schema)  # type: ignore
+
+        database_name, table_name = self.identifier_to_database_and_table(identifier)
+
+        location = self._resolve_table_location(location, database_name, table_name)
+        metadata_location = self._get_metadata_location(location=location)
+        metadata = new_table_metadata(
+            location=location, schema=schema, partition_spec=partition_spec, sort_order=sort_order, properties=properties
+        )
+        io = load_file_io(properties=self.properties, location=metadata_location)
+        return StagedTable(
+            identifier=(self.name, database_name, table_name),
+            metadata=metadata,
+            metadata_location=metadata_location,
+            io=io,
+            catalog=self,
+        )
+
+    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:
+        """Create a CreateTableTransaction.
+
+        Args:
+            identifier (str | Identifier): Table identifier.
+            schema (Schema): Table's schema.
+            location (str | None): Location for the table. Optional Argument.
+            partition_spec (PartitionSpec): PartitionSpec for the table.
+            sort_order (SortOrder): SortOrder for the table.
+            properties (Properties): Table properties that can be a string based dictionary.
+
+        Returns:
+            CreateTableTransaction: createTableTransaction instance.
+
+        Raises:
+            TableAlreadyExistsError: If a table with the name already exists.
+        """
+        return CreateTableTransaction(
+            self._create_staged_table(identifier, schema, location, partition_spec, sort_order, properties)
+        )
+
     @abstractmethod
     def create_table(
         self,
@@ -717,6 +791,10 @@ def _get_updated_props_and_update_summary(
 
         return properties_update_summary, updated_properties
 
+    @staticmethod
+    def empty_table_metadata() -> TableMetadata:
+        return TableMetadataV1(location="", last_column_id=-1, schema=Schema())
+
     def __repr__(self) -> str:
         """Return the string representation of the Catalog class."""
         return f"{self.name} ({self.__class__})"

pyiceberg/catalog/glue.py

@@ -62,8 +62,13 @@
 from pyiceberg.partitioning import UNPARTITIONED_PARTITION_SPEC, PartitionSpec
 from pyiceberg.schema import Schema, SchemaVisitor, visit
 from pyiceberg.serializers import FromInputFile
-from pyiceberg.table import CommitTableRequest, CommitTableResponse, Table, update_table_metadata
-from pyiceberg.table.metadata import TableMetadata, new_table_metadata
+from pyiceberg.table import (
+    CommitTableRequest,
+    CommitTableResponse,
+    Table,
+    update_table_metadata,
+)
+from pyiceberg.table.metadata import TableMetadata
 from pyiceberg.table.sorting import UNSORTED_SORT_ORDER, SortOrder
 from pyiceberg.typedef import EMPTY_DICT, Identifier, Properties
 from pyiceberg.types import (
@@ -384,20 +389,18 @@ def create_table(
             ValueError: If the identifier is invalid, or no path is given to store metadata.
 
         """
-        schema: Schema = self._convert_schema_if_needed(schema)  # type: ignore
-
-        database_name, table_name = self.identifier_to_database_and_table(identifier)
-
-        location = self._resolve_table_location(location, database_name, table_name)
-        metadata_location = self._get_metadata_location(location=location)
-        metadata = new_table_metadata(
-            location=location, schema=schema, partition_spec=partition_spec, sort_order=sort_order, properties=properties
+        staged_table = self._create_staged_table(
+            identifier=identifier,
+            schema=schema,
+            location=location,
+            partition_spec=partition_spec,
+            sort_order=sort_order,
+            properties=properties,
         )
-        io = load_file_io(properties=self.properties, location=metadata_location)
-        self._write_metadata(metadata, io, metadata_location)
-
-        table_input = _construct_table_input(table_name, metadata_location, properties, metadata)
         database_name, table_name = self.identifier_to_database_and_table(identifier)
+
+        self._write_metadata(staged_table.metadata, staged_table.io, staged_table.metadata_location)
+        table_input = _construct_table_input(table_name, staged_table.metadata_location, properties, staged_table.metadata)
         self._create_glue_table(database_name=database_name, table_name=table_name, table_input=table_input)
 
         return self.load_table(identifier=identifier)
@@ -435,46 +438,71 @@ def _commit_table(self, table_request: CommitTableRequest) -> CommitTableRespons
         )
         database_name, table_name = self.identifier_to_database_and_table(identifier_tuple)
 
-        current_glue_table = self._get_glue_table(database_name=database_name, table_name=table_name)
-        glue_table_version_id = current_glue_table.get("VersionId")
-        if not glue_table_version_id:
-            raise CommitFailedException(f"Cannot commit {database_name}.{table_name} because Glue table version id is missing")
-        current_table = self._convert_glue_to_iceberg(glue_table=current_glue_table)
-        base_metadata = current_table.metadata
-
-        # Validate the update requirements
-        for requirement in table_request.requirements:
-            requirement.validate(base_metadata)
-
-        updated_metadata = update_table_metadata(base_metadata, table_request.updates)
-        if updated_metadata == base_metadata:
-            # no changes, do nothing
-            return CommitTableResponse(metadata=base_metadata, metadata_location=current_table.metadata_location)
-
-        # write new metadata
-        new_metadata_version = self._parse_metadata_version(current_table.metadata_location) + 1
-        new_metadata_location = self._get_metadata_location(current_table.metadata.location, new_metadata_version)
-        self._write_metadata(updated_metadata, current_table.io, new_metadata_location)
-
-        update_table_input = _construct_table_input(
-            table_name=table_name,
-            metadata_location=new_metadata_location,
-            properties=current_table.properties,
-            metadata=updated_metadata,
-            glue_table=current_glue_table,
-            prev_metadata_location=current_table.metadata_location,
-        )
+        try:
+            current_glue_table = self._get_glue_table(database_name=database_name, table_name=table_name)
+            # Update the table
+            glue_table_version_id = current_glue_table.get("VersionId")
+            if not glue_table_version_id:
+                raise CommitFailedException(
+                    f"Cannot commit {database_name}.{table_name} because Glue table version id is missing"
+                )
+            current_table = self._convert_glue_to_iceberg(glue_table=current_glue_table)
+            base_metadata = current_table.metadata
+
+            # Validate the update requirements
+            for requirement in table_request.requirements:
+                requirement.validate(base_metadata)
+
+            updated_metadata = update_table_metadata(base_metadata=base_metadata, updates=table_request.updates)
+            if updated_metadata == base_metadata:
+                # no changes, do nothing
+                return CommitTableResponse(metadata=base_metadata, metadata_location=current_table.metadata_location)
+
+            # write new metadata
+            new_metadata_version = self._parse_metadata_version(current_table.metadata_location) + 1
+            new_metadata_location = self._get_metadata_location(current_table.metadata.location, new_metadata_version)
+            self._write_metadata(updated_metadata, current_table.io, new_metadata_location)
+
+            update_table_input = _construct_table_input(
+                table_name=table_name,
+                metadata_location=new_metadata_location,
+                properties=current_table.properties,
+                metadata=updated_metadata,
+                glue_table=current_glue_table,
+                prev_metadata_location=current_table.metadata_location,
+            )
 
-        # Pass `version_id` to implement optimistic locking: it ensures updates are rejected if concurrent
-        # modifications occur. See more details at https://iceberg.apache.org/docs/latest/aws/#optimistic-locking
-        self._update_glue_table(
-            database_name=database_name,
-            table_name=table_name,
-            table_input=update_table_input,
-            version_id=glue_table_version_id,
-        )
+            # Pass `version_id` to implement optimistic locking: it ensures updates are rejected if concurrent
+            # modifications occur. See more details at https://iceberg.apache.org/docs/latest/aws/#optimistic-locking
+            self._update_glue_table(
+                database_name=database_name,
+                table_name=table_name,
+                table_input=update_table_input,
+                version_id=glue_table_version_id,
+            )
+
+            return CommitTableResponse(metadata=updated_metadata, metadata_location=new_metadata_location)
+        except NoSuchTableError:
+            # Create the table
+            updated_metadata = update_table_metadata(
+                base_metadata=self.empty_table_metadata(), updates=table_request.updates, enforce_validation=True
+            )
+            new_metadata_version = 0
+            new_metadata_location = self._get_metadata_location(updated_metadata.location, new_metadata_version)
+            self._write_metadata(
+                updated_metadata, self._load_file_io(updated_metadata.properties, new_metadata_location), new_metadata_location
+            )
+
+            create_table_input = _construct_table_input(
+                table_name=table_name,
+                metadata_location=new_metadata_location,
+                properties=updated_metadata.properties,
+                metadata=updated_metadata,
+            )
+
+            self._create_glue_table(database_name=database_name, table_name=table_name, table_input=create_table_input)
 
-        return CommitTableResponse(metadata=updated_metadata, metadata_location=new_metadata_location)
+            return CommitTableResponse(metadata=updated_metadata, metadata_location=new_metadata_location)
 
     def load_table(self, identifier: Union[str, Identifier]) -> Table:
         """Load the table's metadata and returns the table instance.

pyiceberg/catalog/rest.py

@@ -61,6 +61,7 @@
 from pyiceberg.table import (
     CommitTableRequest,
     CommitTableResponse,
+    StagedTable,
     Table,
     TableIdentifier,
 )
@@ -135,7 +136,7 @@ def _retry_hook(retry_state: RetryCallState) -> None:
 
 
 class TableResponse(IcebergBaseModel):
-    metadata_location: str = Field(alias="metadata-location")
+    metadata_location: Optional[str] = Field(alias="metadata-location")
     metadata: TableMetadata
     config: Properties = Field(default_factory=dict)
 
@@ -460,7 +461,18 @@ def add_headers(self, request: PreparedRequest, **kwargs: Any) -> None:  # pylin
     def _response_to_table(self, identifier_tuple: Tuple[str, ...], table_response: TableResponse) -> Table:
         return Table(
             identifier=(self.name,) + identifier_tuple if self.name else identifier_tuple,
-            metadata_location=table_response.metadata_location,
+            metadata_location=table_response.metadata_location,  # type: ignore
+            metadata=table_response.metadata,
+            io=self._load_file_io(
+                {**table_response.metadata.properties, **table_response.config}, table_response.metadata_location
+            ),
+            catalog=self,
+        )
+
+    def _response_to_staged_table(self, identifier_tuple: Tuple[str, ...], table_response: TableResponse) -> StagedTable:
+        return StagedTable(
+            identifier=(self.name,) + identifier_tuple if self.name else identifier_tuple,
+            metadata_location=table_response.metadata_location,  # type: ignore
             metadata=table_response.metadata,
             io=self._load_file_io(
                 {**table_response.metadata.properties, **table_response.config}, table_response.metadata_location
@@ -490,16 +502,16 @@ def _config_headers(self, session: Session) -> None:
     def _extract_headers_from_properties(self) -> Dict[str, str]:
         return {key[len(HEADER_PREFIX) :]: value for key, value in self.properties.items() if key.startswith(HEADER_PREFIX)}
 
-    @retry(**_RETRY_ARGS)
-    def create_table(
+    def _create_table(
         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,
-    ) -> Table:
+        stage_create: bool = False,
+    ) -> TableResponse:
         iceberg_schema = self._convert_schema_if_needed(schema)
         fresh_schema = assign_fresh_schema_ids(iceberg_schema)
         fresh_partition_spec = assign_fresh_partition_spec_ids(partition_spec, iceberg_schema, fresh_schema)
@@ -512,6 +524,7 @@ def create_table(
             table_schema=fresh_schema,
             partition_spec=fresh_partition_spec,
             write_order=fresh_sort_order,
+            stage_create=stage_create,
             properties=properties,
         )
         serialized_json = request.model_dump_json().encode(UTF8)
@@ -524,7 +537,32 @@ def create_table(
         except HTTPError as exc:
             self._handle_non_200_response(exc, {409: TableAlreadyExistsError})
 
-        table_response = TableResponse(**response.json())
+        return TableResponse(**response.json())
+
+    @retry(**_RETRY_ARGS)
+    def _create_staged_table(
+        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,
+    ) -> StagedTable:
+        table_response = self._create_table(identifier, schema, location, partition_spec, sort_order, properties, True)
+        return self._response_to_staged_table(self.identifier_to_tuple(identifier), table_response)
+
+    @retry(**_RETRY_ARGS)
+    def create_table(
+        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,
+    ) -> Table:
+        table_response = self._create_table(identifier, schema, location, partition_spec, sort_order, properties)
         return self._response_to_table(self.identifier_to_tuple(identifier), table_response)
 
     @retry(**_RETRY_ARGS)