fix create tables

Author: flaneur2020

crates/catalog/s3tables/src/catalog.rs

@@ -6,6 +6,7 @@ use aws_sdk_s3tables::operation::create_table::CreateTableOutput;
 use aws_sdk_s3tables::operation::get_namespace::GetNamespaceOutput;
 use aws_sdk_s3tables::operation::get_table::GetTableOutput;
 use aws_sdk_s3tables::operation::list_tables::ListTablesOutput;
+use aws_sdk_s3tables::types::OpenTableFormat;
 use iceberg::io::FileIO;
 use iceberg::spec::{TableMetadata, TableMetadataBuilder};
 use iceberg::table::Table;
@@ -263,12 +264,14 @@ impl Catalog for S3TablesCatalog {
     /// We have to get this random warehouse location after the table is created.
     ///
     /// This function can return an error in the following situations:
+    /// - If the location of the table is set by user, identified by a specific
+    /// `DataInvalid` variant.
     /// - Errors from the underlying database creation process, converted using
     /// `from_aws_sdk_error`.
     async fn create_table(
         &self,
         namespace: &NamespaceIdent,
-        creation: TableCreation,
+        mut creation: TableCreation,
     ) -> Result<Table> {
         let table_ident = TableIdent::new(namespace.clone(), creation.name.clone());
 
@@ -278,27 +281,41 @@ impl Catalog for S3TablesCatalog {
             .create_table()
             .table_bucket_arn(self.config.table_bucket_arn.clone())
             .namespace(namespace.to_url_string())
+            .format(OpenTableFormat::Iceberg)
             .name(table_ident.name())
             .send()
             .await
             .map_err(from_aws_sdk_error)?;
 
-        // get warehouse location
-        let get_resp: GetTableOutput = self
-            .s3tables_client
-            .get_table()
-            .table_bucket_arn(self.config.table_bucket_arn.clone())
-            .namespace(namespace.to_url_string())
-            .name(table_ident.name())
-            .send()
-            .await
-            .map_err(from_aws_sdk_error)?;
+        // prepare metadata location. the warehouse location is generated by s3tables catalog,
+        // which looks like: s3://e6c9bf20-991a-46fb-kni5xs1q2yxi3xxdyxzjzigdeop1quse2b--table-s3
+        let metadata_location = match &creation.location {
+            Some(_) => {
+                return Err(Error::new(
+                    ErrorKind::DataInvalid,
+                    "The location of the table is generated by s3tables catalog, can't be set by user.",
+                ));
+            }
+            None => {
+                let get_resp: GetTableOutput = self
+                    .s3tables_client
+                    .get_table()
+                    .table_bucket_arn(self.config.table_bucket_arn.clone())
+                    .namespace(namespace.to_url_string())
+                    .name(table_ident.name())
+                    .send()
+                    .await
+                    .map_err(from_aws_sdk_error)?;
+                let warehouse_location = get_resp.warehouse_location().to_string();
+                create_metadata_location(warehouse_location, 0)?
+            }
+        };
 
         // write metadata to file
+        creation.location = Some(metadata_location.clone());
         let metadata = TableMetadataBuilder::from_table_creation(creation)?
             .build()?
             .metadata;
-        let metadata_location = create_metadata_location(get_resp.warehouse_location(), 0)?;
         self.file_io
             .new_output(&metadata_location)?
             .write(serde_json::to_vec(&metadata)?.into())
@@ -460,6 +477,8 @@ where T: std::fmt::Debug {
 
 #[cfg(test)]
 mod tests {
+    use iceberg::spec::{NestedField, PrimitiveType, Schema, Type};
+
     use super::*;
 
     async fn load_s3tables_catalog_from_env() -> Result<Option<S3TablesCatalog>> {
@@ -541,4 +560,47 @@ mod tests {
         catalog.drop_namespace(&namespace).await.unwrap();
         assert!(!catalog.namespace_exists(&namespace).await.unwrap());
     }
+
+    #[tokio::test]
+    async fn test_s3tables_create_delete_table() {
+        let catalog = match load_s3tables_catalog_from_env().await {
+            Ok(Some(catalog)) => catalog,
+            Ok(None) => return,
+            Err(e) => panic!("Error loading catalog: {}", e),
+        };
+
+        let creation = {
+            let schema = Schema::builder()
+                .with_schema_id(0)
+                .with_fields(vec![
+                    NestedField::required(1, "foo", Type::Primitive(PrimitiveType::Int)).into(),
+                    NestedField::required(2, "bar", Type::Primitive(PrimitiveType::String)).into(),
+                ])
+                .build()
+                .unwrap();
+            TableCreation::builder()
+                .name("test_s3tables_create_delete_table".to_string())
+                .properties(HashMap::new())
+                .schema(schema)
+                .build()
+        };
+
+        let namespace = NamespaceIdent::new("test_s3tables_create_delete_table".to_string());
+        let table_ident = TableIdent::new(
+            namespace.clone(),
+            "test_s3tables_create_delete_table".to_string(),
+        );
+        catalog.drop_namespace(&namespace).await.ok();
+        catalog.drop_table(&table_ident).await.ok();
+
+        catalog
+            .create_namespace(&namespace, HashMap::new())
+            .await
+            .unwrap();
+        catalog.create_table(&namespace, creation).await.unwrap();
+        assert!(catalog.table_exists(&table_ident).await.unwrap());
+        catalog.drop_table(&table_ident).await.unwrap();
+        assert!(!catalog.table_exists(&table_ident).await.unwrap());
+        catalog.drop_namespace(&namespace).await.unwrap();
+    }
 }