1use std::collections::HashMap;
19use std::future::Future;
20use std::str::FromStr;
21use std::sync::Arc;
22
23use async_trait::async_trait;
24use aws_sdk_s3tables::operation::create_table::CreateTableOutput;
25use aws_sdk_s3tables::operation::get_namespace::GetNamespaceOutput;
26use aws_sdk_s3tables::operation::get_table::{GetTableError, GetTableOutput};
27use aws_sdk_s3tables::operation::list_tables::ListTablesOutput;
28use aws_sdk_s3tables::operation::update_table_metadata_location::UpdateTableMetadataLocationError;
29use aws_sdk_s3tables::types::OpenTableFormat;
30use iceberg::encryption::kms::{KeyManagementClient, KmsClientFactory};
31use iceberg::io::{FileIO, FileIOBuilder, StorageFactory};
32use iceberg::spec::{TableMetadata, TableMetadataBuilder};
33use iceberg::table::Table;
34use iceberg::{
35 Catalog, CatalogBuilder, Error, ErrorKind, MetadataLocation, Namespace, NamespaceIdent, Result,
36 Runtime, TableCommit, TableCreation, TableIdent,
37};
38use iceberg_storage_opendal::OpenDalStorageFactory;
39
40use crate::utils::create_sdk_config;
41
42pub const S3TABLES_CATALOG_PROP_TABLE_BUCKET_ARN: &str = "table_bucket_arn";
44pub const S3TABLES_CATALOG_PROP_ENDPOINT_URL: &str = "endpoint_url";
46
47#[derive(Debug)]
49struct S3TablesCatalogConfig {
50 name: Option<String>,
52 table_bucket_arn: String,
57 endpoint_url: Option<String>,
59 client: Option<aws_sdk_s3tables::Client>,
61 props: HashMap<String, String>,
68}
69
70#[derive(Debug)]
72pub struct S3TablesCatalogBuilder {
73 config: S3TablesCatalogConfig,
74 storage_factory: Option<Arc<dyn StorageFactory>>,
75 kms_client_factory: Option<Arc<dyn KmsClientFactory>>,
76 runtime: Option<Runtime>,
77}
78
79impl Default for S3TablesCatalogBuilder {
81 fn default() -> Self {
82 Self {
83 config: S3TablesCatalogConfig {
84 name: None,
85 table_bucket_arn: "".to_string(),
86 endpoint_url: None,
87 client: None,
88 props: HashMap::new(),
89 },
90 storage_factory: None,
91 kms_client_factory: None,
92 runtime: None,
93 }
94 }
95}
96
97impl S3TablesCatalogBuilder {
99 pub fn with_endpoint_url(mut self, endpoint_url: impl Into<String>) -> Self {
108 self.config.endpoint_url = Some(endpoint_url.into());
109 self
110 }
111
112 pub fn with_client(mut self, client: aws_sdk_s3tables::Client) -> Self {
114 self.config.client = Some(client);
115 self
116 }
117
118 pub fn with_table_bucket_arn(mut self, table_bucket_arn: impl Into<String>) -> Self {
127 self.config.table_bucket_arn = table_bucket_arn.into();
128 self
129 }
130}
131
132impl CatalogBuilder for S3TablesCatalogBuilder {
133 type C = S3TablesCatalog;
134
135 fn with_storage_factory(mut self, storage_factory: Arc<dyn StorageFactory>) -> Self {
136 self.storage_factory = Some(storage_factory);
137 self
138 }
139
140 fn with_kms_client_factory(mut self, kms_client_factory: Arc<dyn KmsClientFactory>) -> Self {
141 self.kms_client_factory = Some(kms_client_factory);
142 self
143 }
144
145 fn with_runtime(mut self, runtime: Runtime) -> Self {
146 self.runtime = Some(runtime);
147 self
148 }
149
150 fn load(
151 mut self,
152 name: impl Into<String>,
153 props: HashMap<String, String>,
154 ) -> impl Future<Output = Result<Self::C>> + Send {
155 let catalog_name = name.into();
156 self.config.name = Some(catalog_name.clone());
157
158 if props.contains_key(S3TABLES_CATALOG_PROP_TABLE_BUCKET_ARN) {
159 self.config.table_bucket_arn = props
160 .get(S3TABLES_CATALOG_PROP_TABLE_BUCKET_ARN)
161 .cloned()
162 .unwrap_or_default();
163 }
164
165 if props.contains_key(S3TABLES_CATALOG_PROP_ENDPOINT_URL) {
166 self.config.endpoint_url = props.get(S3TABLES_CATALOG_PROP_ENDPOINT_URL).cloned();
167 }
168
169 self.config.props = props
171 .into_iter()
172 .filter(|(k, _)| {
173 k != S3TABLES_CATALOG_PROP_TABLE_BUCKET_ARN
174 && k != S3TABLES_CATALOG_PROP_ENDPOINT_URL
175 })
176 .collect();
177
178 async move {
179 if catalog_name.trim().is_empty() {
180 Err(Error::new(
181 ErrorKind::DataInvalid,
182 "Catalog name cannot be empty",
183 ))
184 } else if self.config.table_bucket_arn.is_empty() {
185 Err(Error::new(
186 ErrorKind::DataInvalid,
187 "Table bucket ARN is required",
188 ))
189 } else {
190 let runtime = match self.runtime {
191 Some(rt) => rt,
192 None => Runtime::try_current()?,
193 };
194 let kms_client = match self.kms_client_factory {
195 Some(factory) => Some(factory.create_kms_client(&self.config.props).await?),
196 None => None,
197 };
198 S3TablesCatalog::new(self.config, self.storage_factory, runtime, kms_client).await
199 }
200 }
201 }
202}
203
204#[derive(Debug)]
206pub struct S3TablesCatalog {
207 config: S3TablesCatalogConfig,
208 s3tables_client: aws_sdk_s3tables::Client,
209 file_io: FileIO,
210 runtime: Runtime,
211 kms_client: Option<Arc<dyn KeyManagementClient>>,
212}
213
214impl S3TablesCatalog {
215 async fn new(
217 config: S3TablesCatalogConfig,
218 storage_factory: Option<Arc<dyn StorageFactory>>,
219 runtime: Runtime,
220 kms_client: Option<Arc<dyn KeyManagementClient>>,
221 ) -> Result<Self> {
222 let s3tables_client = if let Some(client) = config.client.clone() {
223 client
224 } else {
225 let aws_config = create_sdk_config(&config.props, config.endpoint_url.clone()).await;
226 aws_sdk_s3tables::Client::new(&aws_config)
227 };
228
229 let factory = storage_factory.unwrap_or_else(|| {
231 Arc::new(OpenDalStorageFactory::S3 {
232 customized_credential_load: None,
233 })
234 });
235 let file_io = FileIOBuilder::new(factory)
236 .with_props(&config.props)
237 .build();
238
239 Ok(Self {
240 config,
241 s3tables_client,
242 file_io,
243 runtime,
244 kms_client,
245 })
246 }
247
248 async fn load_table_with_version_token(
249 &self,
250 table_ident: &TableIdent,
251 ) -> Result<(Table, String)> {
252 let req = self
253 .s3tables_client
254 .get_table()
255 .table_bucket_arn(self.config.table_bucket_arn.clone())
256 .namespace(table_ident.namespace().to_url_string())
257 .name(table_ident.name());
258
259 let resp = req.send().await.map_err(|err| {
260 if err
261 .as_service_error()
262 .is_some_and(GetTableError::is_not_found_exception)
263 {
264 Error::new(
267 ErrorKind::TableNotFound,
268 format!("Table {table_ident} is not found, either because the namespace or table did not exist"),
269 )
270 .with_source(err)
271 } else {
272 from_aws_sdk_error(err)
273 }
274 })?;
275
276 let metadata_location = resp.metadata_location().ok_or_else(|| {
278 Error::new(
279 ErrorKind::Unexpected,
280 format!(
281 "Table {} does not have metadata location",
282 table_ident.name()
283 ),
284 )
285 })?;
286 let metadata = TableMetadata::read_from(&self.file_io, metadata_location).await?;
287
288 let mut builder = Table::builder()
289 .identifier(table_ident.clone())
290 .metadata(metadata)
291 .metadata_location(metadata_location)
292 .file_io(self.file_io.clone())
293 .runtime(self.runtime.clone());
294 if let Some(kms_client) = self.kms_client.clone() {
295 builder = builder.kms_client(kms_client);
296 }
297 let table = builder.build()?;
298 Ok((table, resp.version_token))
299 }
300}
301
302#[async_trait]
303impl Catalog for S3TablesCatalog {
304 async fn list_namespaces(
309 &self,
310 parent: Option<&NamespaceIdent>,
311 ) -> Result<Vec<NamespaceIdent>> {
312 if parent.is_some() {
313 return Ok(vec![]);
314 }
315
316 let mut result = Vec::new();
317 let mut continuation_token = None;
318 loop {
319 let mut req = self
320 .s3tables_client
321 .list_namespaces()
322 .table_bucket_arn(self.config.table_bucket_arn.clone());
323 if let Some(token) = continuation_token {
324 req = req.continuation_token(token);
325 }
326 let resp = req.send().await.map_err(from_aws_sdk_error)?;
327 for ns in resp.namespaces() {
328 result.push(NamespaceIdent::from_vec(ns.namespace().to_vec())?);
329 }
330 continuation_token = resp.continuation_token().map(|s| s.to_string());
331 if continuation_token.is_none() {
332 break;
333 }
334 }
335 Ok(result)
336 }
337
338 async fn create_namespace(
355 &self,
356 namespace: &NamespaceIdent,
357 _properties: HashMap<String, String>,
358 ) -> Result<Namespace> {
359 if self.namespace_exists(namespace).await? {
360 return Err(Error::new(
361 ErrorKind::NamespaceAlreadyExists,
362 format!("Namespace {namespace:?} already exists"),
363 ));
364 }
365
366 let req = self
367 .s3tables_client
368 .create_namespace()
369 .table_bucket_arn(self.config.table_bucket_arn.clone())
370 .namespace(namespace.to_url_string());
371 req.send().await.map_err(from_aws_sdk_error)?;
372 Ok(Namespace::with_properties(
373 namespace.clone(),
374 HashMap::new(),
375 ))
376 }
377
378 async fn get_namespace(&self, namespace: &NamespaceIdent) -> Result<Namespace> {
388 if !self.namespace_exists(namespace).await? {
389 return Err(Error::new(
390 ErrorKind::NamespaceNotFound,
391 format!("Namespace {namespace:?} does not exist"),
392 ));
393 }
394
395 let req = self
396 .s3tables_client
397 .get_namespace()
398 .table_bucket_arn(self.config.table_bucket_arn.clone())
399 .namespace(namespace.to_url_string());
400 let resp: GetNamespaceOutput = req.send().await.map_err(from_aws_sdk_error)?;
401 let properties = HashMap::new();
402 Ok(Namespace::with_properties(
403 NamespaceIdent::from_vec(resp.namespace().to_vec())?,
404 properties,
405 ))
406 }
407
408 async fn namespace_exists(&self, namespace: &NamespaceIdent) -> Result<bool> {
421 let req = self
422 .s3tables_client
423 .get_namespace()
424 .table_bucket_arn(self.config.table_bucket_arn.clone())
425 .namespace(namespace.to_url_string());
426 match req.send().await {
427 Ok(_) => Ok(true),
428 Err(err) => {
429 if err.as_service_error().map(|e| e.is_not_found_exception()) == Some(true) {
430 Ok(false)
431 } else {
432 Err(from_aws_sdk_error(err))
433 }
434 }
435 }
436 }
437
438 async fn update_namespace(
443 &self,
444 _namespace: &NamespaceIdent,
445 _properties: HashMap<String, String>,
446 ) -> Result<()> {
447 Err(Error::new(
448 ErrorKind::FeatureUnsupported,
449 "Update namespace is not supported for s3tables catalog",
450 ))
451 }
452
453 async fn drop_namespace(&self, namespace: &NamespaceIdent) -> Result<()> {
462 if !self.namespace_exists(namespace).await? {
463 return Err(Error::new(
464 ErrorKind::NamespaceNotFound,
465 format!("Namespace {namespace:?} does not exist"),
466 ));
467 }
468
469 let req = self
470 .s3tables_client
471 .delete_namespace()
472 .table_bucket_arn(self.config.table_bucket_arn.clone())
473 .namespace(namespace.to_url_string());
474 req.send().await.map_err(from_aws_sdk_error)?;
475 Ok(())
476 }
477
478 async fn list_tables(&self, namespace: &NamespaceIdent) -> Result<Vec<TableIdent>> {
487 let mut result = Vec::new();
488 let mut continuation_token = None;
489 loop {
490 let mut req = self
491 .s3tables_client
492 .list_tables()
493 .table_bucket_arn(self.config.table_bucket_arn.clone())
494 .namespace(namespace.to_url_string());
495 if let Some(token) = continuation_token {
496 req = req.continuation_token(token);
497 }
498 let resp: ListTablesOutput = req.send().await.map_err(from_aws_sdk_error)?;
499 for table in resp.tables() {
500 result.push(TableIdent::new(
501 NamespaceIdent::from_vec(table.namespace().to_vec())?,
502 table.name().to_string(),
503 ));
504 }
505 continuation_token = resp.continuation_token().map(|s| s.to_string());
506 if continuation_token.is_none() {
507 break;
508 }
509 }
510 Ok(result)
511 }
512
513 async fn create_table(
528 &self,
529 namespace: &NamespaceIdent,
530 mut creation: TableCreation,
531 ) -> Result<Table> {
532 let table_ident = TableIdent::new(namespace.clone(), creation.name.clone());
533
534 let create_resp: CreateTableOutput = self
536 .s3tables_client
537 .create_table()
538 .table_bucket_arn(self.config.table_bucket_arn.clone())
539 .namespace(namespace.to_url_string())
540 .format(OpenTableFormat::Iceberg)
541 .name(table_ident.name())
542 .send()
543 .await
544 .map_err(from_aws_sdk_error)?;
545
546 let table_location = match &creation.location {
549 Some(_) => {
550 return Err(Error::new(
551 ErrorKind::DataInvalid,
552 "The location of the table is generated by s3tables catalog, can't be set by user.",
553 ));
554 }
555 None => {
556 let get_resp: GetTableOutput = self
557 .s3tables_client
558 .get_table()
559 .table_bucket_arn(self.config.table_bucket_arn.clone())
560 .namespace(namespace.to_url_string())
561 .name(table_ident.name())
562 .send()
563 .await
564 .map_err(from_aws_sdk_error)?;
565 get_resp.warehouse_location().to_string()
566 }
567 };
568
569 creation.location = Some(table_location.clone());
571 let metadata = TableMetadataBuilder::from_table_creation(creation)?
572 .build()?
573 .metadata;
574 let metadata_location = MetadataLocation::new_with_metadata(table_location, &metadata);
575 metadata.write_to(&self.file_io, &metadata_location).await?;
576
577 let metadata_location_str = metadata_location.to_string();
579 self.s3tables_client
580 .update_table_metadata_location()
581 .table_bucket_arn(self.config.table_bucket_arn.clone())
582 .namespace(namespace.to_url_string())
583 .name(table_ident.name())
584 .metadata_location(metadata_location_str.clone())
585 .version_token(create_resp.version_token())
586 .send()
587 .await
588 .map_err(from_aws_sdk_error)?;
589
590 let mut builder = Table::builder()
591 .identifier(table_ident)
592 .metadata_location(metadata_location_str)
593 .metadata(metadata)
594 .file_io(self.file_io.clone())
595 .runtime(self.runtime.clone());
596 if let Some(kms_client) = self.kms_client.clone() {
597 builder = builder.kms_client(kms_client);
598 }
599 let table = builder.build()?;
600 Ok(table)
601 }
602
603 async fn load_table(&self, table_ident: &TableIdent) -> Result<Table> {
614 Ok(self.load_table_with_version_token(table_ident).await?.0)
615 }
616
617 async fn drop_table(&self, _table: &TableIdent) -> Result<()> {
621 Err(Error::new(
622 ErrorKind::FeatureUnsupported,
623 "drop_table is not supported for S3Tables; use purge_table instead",
624 ))
625 }
626
627 async fn purge_table(&self, table: &TableIdent) -> Result<()> {
629 let req = self
630 .s3tables_client
631 .delete_table()
632 .table_bucket_arn(self.config.table_bucket_arn.clone())
633 .namespace(table.namespace().to_url_string())
634 .name(table.name());
635 req.send().await.map_err(from_aws_sdk_error)?;
636 Ok(())
637 }
638
639 async fn table_exists(&self, table_ident: &TableIdent) -> Result<bool> {
652 let req = self
653 .s3tables_client
654 .get_table()
655 .table_bucket_arn(self.config.table_bucket_arn.clone())
656 .namespace(table_ident.namespace().to_url_string())
657 .name(table_ident.name());
658 match req.send().await {
659 Ok(_) => Ok(true),
660 Err(err) => {
661 if err.as_service_error().map(|e| e.is_not_found_exception()) == Some(true) {
662 Ok(false)
663 } else {
664 Err(from_aws_sdk_error(err))
665 }
666 }
667 }
668 }
669
670 async fn rename_table(&self, src: &TableIdent, dest: &TableIdent) -> Result<()> {
679 let req = self
680 .s3tables_client
681 .rename_table()
682 .table_bucket_arn(self.config.table_bucket_arn.clone())
683 .namespace(src.namespace().to_url_string())
684 .name(src.name())
685 .new_namespace_name(dest.namespace().to_url_string())
686 .new_name(dest.name());
687 req.send().await.map_err(from_aws_sdk_error)?;
688 Ok(())
689 }
690
691 async fn register_table(
692 &self,
693 _table_ident: &TableIdent,
694 _metadata_location: String,
695 ) -> Result<Table> {
696 Err(Error::new(
697 ErrorKind::FeatureUnsupported,
698 "Registering a table is not supported yet",
699 ))
700 }
701
702 async fn update_table(&self, commit: TableCommit) -> Result<Table> {
704 let table_ident = commit.identifier().clone();
705 let table_namespace = table_ident.namespace();
706 let (current_table, version_token) =
707 self.load_table_with_version_token(&table_ident).await?;
708
709 let staged_table = commit.apply(current_table)?;
710 let staged_metadata_location_str = staged_table.metadata_location_result()?;
711 let staged_metadata_location = MetadataLocation::from_str(staged_metadata_location_str)?;
712
713 staged_table
714 .metadata()
715 .write_to(staged_table.file_io(), &staged_metadata_location)
716 .await?;
717
718 let builder = self
719 .s3tables_client
720 .update_table_metadata_location()
721 .table_bucket_arn(&self.config.table_bucket_arn)
722 .namespace(table_namespace.to_url_string())
723 .name(table_ident.name())
724 .version_token(version_token)
725 .metadata_location(staged_metadata_location_str);
726
727 let _ = builder.send().await.map_err(|e| {
728 let error = e.into_service_error();
729 match error {
730 UpdateTableMetadataLocationError::ConflictException(_) => Error::new(
731 ErrorKind::CatalogCommitConflicts,
732 format!("Commit conflicted for table: {table_ident}"),
733 )
734 .with_retryable(true),
735 UpdateTableMetadataLocationError::NotFoundException(_) => Error::new(
736 ErrorKind::TableNotFound,
737 format!("Table {table_ident} is not found"),
738 ),
739 _ => Error::new(
740 ErrorKind::Unexpected,
741 "Operation failed for hitting aws sdk error",
742 ),
743 }
744 .with_source(anyhow::Error::msg(format!("aws sdk error: {error:?}")))
745 })?;
746
747 Ok(staged_table)
748 }
749}
750
751pub(crate) fn from_aws_sdk_error<T>(error: aws_sdk_s3tables::error::SdkError<T>) -> Error
753where T: std::fmt::Debug {
754 Error::new(
755 ErrorKind::Unexpected,
756 format!("Operation failed for hitting aws sdk error: {error:?}"),
757 )
758}
759
760#[cfg(test)]
761mod tests {
762 use futures::TryStreamExt;
763 use iceberg::spec::{NestedField, PrimitiveType, Schema, Type};
764 use iceberg::transaction::{ApplyTransactionAction, Transaction};
765
766 use super::*;
767
768 async fn load_s3tables_catalog_from_env() -> Result<Option<S3TablesCatalog>> {
769 let table_bucket_arn = match std::env::var("TABLE_BUCKET_ARN").ok() {
770 Some(table_bucket_arn) => table_bucket_arn,
771 None => return Ok(None),
772 };
773
774 let config = S3TablesCatalogConfig {
775 name: None,
776 table_bucket_arn,
777 endpoint_url: None,
778 client: None,
779 props: HashMap::new(),
780 };
781
782 Ok(Some(
783 S3TablesCatalog::new(config, None, Runtime::current(), None).await?,
784 ))
785 }
786
787 #[tokio::test]
788 async fn test_s3tables_list_namespace() {
789 let catalog = match load_s3tables_catalog_from_env().await {
790 Ok(Some(catalog)) => catalog,
791 Ok(None) => return,
792 Err(e) => panic!("Error loading catalog: {e}"),
793 };
794
795 let namespaces = catalog.list_namespaces(None).await.unwrap();
796 assert!(!namespaces.is_empty());
797 }
798
799 #[tokio::test]
800 async fn test_s3tables_list_tables() {
801 let catalog = match load_s3tables_catalog_from_env().await {
802 Ok(Some(catalog)) => catalog,
803 Ok(None) => return,
804 Err(e) => panic!("Error loading catalog: {e}"),
805 };
806
807 let tables = catalog
808 .list_tables(&NamespaceIdent::new("aws_s3_metadata".to_string()))
809 .await
810 .unwrap();
811 assert!(!tables.is_empty());
812 }
813
814 #[tokio::test]
815 async fn test_s3tables_load_table() {
816 let catalog = match load_s3tables_catalog_from_env().await {
817 Ok(Some(catalog)) => catalog,
818 Ok(None) => return,
819 Err(e) => panic!("Error loading catalog: {e}"),
820 };
821
822 let _table = catalog
823 .load_table(&TableIdent::new(
824 NamespaceIdent::new("aws_s3_metadata".to_string()),
825 "query_storage_metadata".to_string(),
826 ))
827 .await
828 .expect("table that exists should be loaded");
829
830 let load_table_err = catalog
831 .load_table(&TableIdent::new(
832 NamespaceIdent::new("not_a_namespace".to_string()),
833 "not_a_table_name".to_string(),
834 ))
835 .await
836 .expect_err("loading a table that does not exist should fail");
837 assert_eq!(
838 load_table_err.kind(),
839 ErrorKind::TableNotFound,
840 "must return table not found error for non-existent table"
841 );
842 }
843
844 #[tokio::test]
845 async fn test_s3tables_create_delete_namespace() {
846 let catalog = match load_s3tables_catalog_from_env().await {
847 Ok(Some(catalog)) => catalog,
848 Ok(None) => return,
849 Err(e) => panic!("Error loading catalog: {e}"),
850 };
851
852 let namespace = NamespaceIdent::new("test_s3tables_create_delete_namespace".to_string());
853 catalog
854 .create_namespace(&namespace, HashMap::new())
855 .await
856 .unwrap();
857 assert!(catalog.namespace_exists(&namespace).await.unwrap());
858 catalog.drop_namespace(&namespace).await.unwrap();
859 assert!(!catalog.namespace_exists(&namespace).await.unwrap());
860 }
861
862 #[tokio::test]
863 async fn test_s3tables_create_delete_table() {
864 let catalog = match load_s3tables_catalog_from_env().await {
865 Ok(Some(catalog)) => catalog,
866 Ok(None) => return,
867 Err(e) => panic!("Error loading catalog: {e}"),
868 };
869
870 let creation = {
871 let schema = Schema::builder()
872 .with_schema_id(0)
873 .with_fields(vec![
874 NestedField::required(1, "foo", Type::Primitive(PrimitiveType::Int)).into(),
875 NestedField::required(2, "bar", Type::Primitive(PrimitiveType::String)).into(),
876 ])
877 .build()
878 .unwrap();
879 TableCreation::builder()
880 .name("test_s3tables_create_delete_table".to_string())
881 .properties(HashMap::new())
882 .schema(schema)
883 .build()
884 };
885
886 let namespace = NamespaceIdent::new("test_s3tables_create_delete_table".to_string());
887 let table_ident = TableIdent::new(
888 namespace.clone(),
889 "test_s3tables_create_delete_table".to_string(),
890 );
891 catalog.drop_namespace(&namespace).await.ok();
892 catalog.drop_table(&table_ident).await.ok();
893
894 catalog
895 .create_namespace(&namespace, HashMap::new())
896 .await
897 .unwrap();
898 catalog.create_table(&namespace, creation).await.unwrap();
899 assert!(catalog.table_exists(&table_ident).await.unwrap());
900 catalog.drop_table(&table_ident).await.unwrap();
901 assert!(!catalog.table_exists(&table_ident).await.unwrap());
902 catalog.drop_namespace(&namespace).await.unwrap();
903 }
904
905 #[tokio::test]
906 async fn test_s3tables_update_table() {
907 let catalog = match load_s3tables_catalog_from_env().await {
908 Ok(Some(catalog)) => catalog,
909 Ok(None) => return,
910 Err(e) => panic!("Error loading catalog: {e}"),
911 };
912
913 let namespace = NamespaceIdent::new("test_s3tables_update_table".to_string());
915 let table_ident =
916 TableIdent::new(namespace.clone(), "test_s3tables_update_table".to_string());
917
918 catalog.drop_table(&table_ident).await.ok();
920 catalog.drop_namespace(&namespace).await.ok();
921
922 catalog
924 .create_namespace(&namespace, HashMap::new())
925 .await
926 .unwrap();
927
928 let creation = {
929 let schema = Schema::builder()
930 .with_schema_id(0)
931 .with_fields(vec![
932 NestedField::required(1, "foo", Type::Primitive(PrimitiveType::Int)).into(),
933 NestedField::required(2, "bar", Type::Primitive(PrimitiveType::String)).into(),
934 ])
935 .build()
936 .unwrap();
937 TableCreation::builder()
938 .name(table_ident.name().to_string())
939 .properties(HashMap::new())
940 .schema(schema)
941 .build()
942 };
943
944 let table = catalog.create_table(&namespace, creation).await.unwrap();
945
946 let tx = Transaction::new(&table);
948
949 let original_metadata_location = table.metadata_location();
951
952 let tx = tx
954 .update_table_properties()
955 .set("test_property".to_string(), "test_value".to_string())
956 .apply(tx)
957 .unwrap();
958
959 let updated_table = tx.commit(&catalog).await.unwrap();
961
962 assert_eq!(
964 updated_table.metadata().properties().get("test_property"),
965 Some(&"test_value".to_string())
966 );
967
968 assert_ne!(
970 updated_table.metadata_location(),
971 original_metadata_location,
972 "Metadata location should be updated after commit"
973 );
974
975 let reloaded_table = catalog.load_table(&table_ident).await.unwrap();
977
978 assert_eq!(
980 reloaded_table.metadata().properties().get("test_property"),
981 Some(&"test_value".to_string())
982 );
983 assert_eq!(
984 reloaded_table.metadata_location(),
985 updated_table.metadata_location(),
986 "Reloaded table should have the same metadata location as the updated table"
987 );
988 }
989
990 #[tokio::test]
991 async fn test_builder_load_missing_bucket_arn() {
992 let builder = S3TablesCatalogBuilder::default();
993 let result = builder.load("s3tables", HashMap::new()).await;
994
995 assert!(result.is_err());
996 if let Err(err) = result {
997 assert_eq!(err.kind(), ErrorKind::DataInvalid);
998 assert_eq!(err.message(), "Table bucket ARN is required");
999 }
1000 }
1001
1002 #[tokio::test]
1003 async fn test_builder_with_endpoint_url_ok() {
1004 let builder = S3TablesCatalogBuilder::default().with_endpoint_url("http://localhost:4566");
1005
1006 let result = builder
1007 .load(
1008 "s3tables",
1009 HashMap::from([
1010 (
1011 S3TABLES_CATALOG_PROP_TABLE_BUCKET_ARN.to_string(),
1012 "arn:aws:s3tables:us-east-1:123456789012:bucket/test".to_string(),
1013 ),
1014 ("some_prop".to_string(), "some_value".to_string()),
1015 ]),
1016 )
1017 .await;
1018
1019 assert!(result.is_ok());
1020 }
1021
1022 #[tokio::test]
1023 async fn test_builder_with_client_ok() {
1024 use aws_config::BehaviorVersion;
1025
1026 let sdk_config = aws_config::defaults(BehaviorVersion::latest()).load().await;
1027 let client = aws_sdk_s3tables::Client::new(&sdk_config);
1028
1029 let builder = S3TablesCatalogBuilder::default().with_client(client);
1030 let result = builder
1031 .load(
1032 "s3tables",
1033 HashMap::from([(
1034 S3TABLES_CATALOG_PROP_TABLE_BUCKET_ARN.to_string(),
1035 "arn:aws:s3tables:us-east-1:123456789012:bucket/test".to_string(),
1036 )]),
1037 )
1038 .await;
1039
1040 assert!(result.is_ok());
1041 }
1042
1043 #[tokio::test]
1044 async fn test_builder_with_table_bucket_arn() {
1045 let test_arn = "arn:aws:s3tables:us-west-2:123456789012:bucket/test-bucket";
1046 let builder = S3TablesCatalogBuilder::default().with_table_bucket_arn(test_arn);
1047
1048 let result = builder.load("s3tables", HashMap::new()).await;
1049
1050 assert!(result.is_ok());
1051 let catalog = result.unwrap();
1052 assert_eq!(catalog.config.table_bucket_arn, test_arn);
1053 }
1054
1055 #[tokio::test]
1056 async fn test_builder_empty_table_bucket_arn_edge_cases() {
1057 let mut props = HashMap::new();
1058 props.insert(
1059 S3TABLES_CATALOG_PROP_TABLE_BUCKET_ARN.to_string(),
1060 "".to_string(),
1061 );
1062
1063 let builder = S3TablesCatalogBuilder::default();
1064 let result = builder.load("s3tables", props).await;
1065
1066 assert!(result.is_err());
1067 if let Err(err) = result {
1068 assert_eq!(err.kind(), ErrorKind::DataInvalid);
1069 assert_eq!(err.message(), "Table bucket ARN is required");
1070 }
1071 }
1072
1073 #[tokio::test]
1074 async fn test_endpoint_url_property_overrides_builder_method() {
1075 let test_arn = "arn:aws:s3tables:us-west-2:123456789012:bucket/test-bucket";
1076 let builder_endpoint = "http://localhost:4566";
1077 let property_endpoint = "http://localhost:8080";
1078
1079 let builder = S3TablesCatalogBuilder::default()
1080 .with_table_bucket_arn(test_arn)
1081 .with_endpoint_url(builder_endpoint);
1082
1083 let mut props = HashMap::new();
1084 props.insert(
1085 S3TABLES_CATALOG_PROP_ENDPOINT_URL.to_string(),
1086 property_endpoint.to_string(),
1087 );
1088
1089 let result = builder.load("s3tables", props).await;
1090
1091 assert!(result.is_ok());
1092 let catalog = result.unwrap();
1093
1094 assert_eq!(
1096 catalog.config.endpoint_url,
1097 Some(property_endpoint.to_string())
1098 );
1099 assert_ne!(
1100 catalog.config.endpoint_url,
1101 Some(builder_endpoint.to_string())
1102 );
1103 }
1104
1105 #[tokio::test]
1106 async fn test_endpoint_url_builder_method_only() {
1107 let test_arn = "arn:aws:s3tables:us-west-2:123456789012:bucket/test-bucket";
1108 let builder_endpoint = "http://localhost:4566";
1109
1110 let builder = S3TablesCatalogBuilder::default()
1111 .with_table_bucket_arn(test_arn)
1112 .with_endpoint_url(builder_endpoint);
1113
1114 let result = builder.load("s3tables", HashMap::new()).await;
1115
1116 assert!(result.is_ok());
1117 let catalog = result.unwrap();
1118
1119 assert_eq!(
1120 catalog.config.endpoint_url,
1121 Some(builder_endpoint.to_string())
1122 );
1123 }
1124
1125 #[tokio::test]
1126 async fn test_endpoint_url_property_only() {
1127 let test_arn = "arn:aws:s3tables:us-west-2:123456789012:bucket/test-bucket";
1128 let property_endpoint = "http://localhost:8080";
1129
1130 let builder = S3TablesCatalogBuilder::default().with_table_bucket_arn(test_arn);
1131
1132 let mut props = HashMap::new();
1133 props.insert(
1134 S3TABLES_CATALOG_PROP_ENDPOINT_URL.to_string(),
1135 property_endpoint.to_string(),
1136 );
1137
1138 let result = builder.load("s3tables", props).await;
1139
1140 assert!(result.is_ok());
1141 let catalog = result.unwrap();
1142
1143 assert_eq!(
1144 catalog.config.endpoint_url,
1145 Some(property_endpoint.to_string())
1146 );
1147 }
1148
1149 #[tokio::test]
1150 async fn test_table_bucket_arn_property_overrides_builder_method() {
1151 let builder_arn = "arn:aws:s3tables:us-west-2:123456789012:bucket/builder-bucket";
1152 let property_arn = "arn:aws:s3tables:us-east-1:987654321098:bucket/property-bucket";
1153
1154 let builder = S3TablesCatalogBuilder::default().with_table_bucket_arn(builder_arn);
1155
1156 let mut props = HashMap::new();
1157 props.insert(
1158 S3TABLES_CATALOG_PROP_TABLE_BUCKET_ARN.to_string(),
1159 property_arn.to_string(),
1160 );
1161
1162 let result = builder.load("s3tables", props).await;
1163
1164 assert!(result.is_ok());
1165 let catalog = result.unwrap();
1166
1167 assert_eq!(catalog.config.table_bucket_arn, property_arn);
1168 assert_ne!(catalog.config.table_bucket_arn, builder_arn);
1169 }
1170
1171 #[tokio::test]
1172 async fn test_table_bucket_arn_builder_method_only() {
1173 let builder_arn = "arn:aws:s3tables:us-west-2:123456789012:bucket/builder-bucket";
1174
1175 let builder = S3TablesCatalogBuilder::default().with_table_bucket_arn(builder_arn);
1176
1177 let result = builder.load("s3tables", HashMap::new()).await;
1178
1179 assert!(result.is_ok());
1180 let catalog = result.unwrap();
1181
1182 assert_eq!(catalog.config.table_bucket_arn, builder_arn);
1183 }
1184
1185 #[tokio::test]
1186 async fn test_table_bucket_arn_property_only() {
1187 let property_arn = "arn:aws:s3tables:us-east-1:987654321098:bucket/property-bucket";
1188
1189 let builder = S3TablesCatalogBuilder::default();
1190
1191 let mut props = HashMap::new();
1192 props.insert(
1193 S3TABLES_CATALOG_PROP_TABLE_BUCKET_ARN.to_string(),
1194 property_arn.to_string(),
1195 );
1196
1197 let result = builder.load("s3tables", props).await;
1198
1199 assert!(result.is_ok());
1200 let catalog = result.unwrap();
1201
1202 assert_eq!(catalog.config.table_bucket_arn, property_arn);
1203 }
1204
1205 #[tokio::test]
1206 async fn test_builder_empty_name_validation() {
1207 let test_arn = "arn:aws:s3tables:us-west-2:123456789012:bucket/test-bucket";
1208 let builder = S3TablesCatalogBuilder::default().with_table_bucket_arn(test_arn);
1209
1210 let result = builder.load("", HashMap::new()).await;
1211
1212 assert!(result.is_err());
1213 if let Err(err) = result {
1214 assert_eq!(err.kind(), ErrorKind::DataInvalid);
1215 assert_eq!(err.message(), "Catalog name cannot be empty");
1216 }
1217 }
1218
1219 #[tokio::test]
1220 async fn test_builder_whitespace_only_name_validation() {
1221 let test_arn = "arn:aws:s3tables:us-west-2:123456789012:bucket/test-bucket";
1222 let builder = S3TablesCatalogBuilder::default().with_table_bucket_arn(test_arn);
1223
1224 let result = builder.load(" \t\n ", HashMap::new()).await;
1225
1226 assert!(result.is_err());
1227 if let Err(err) = result {
1228 assert_eq!(err.kind(), ErrorKind::DataInvalid);
1229 assert_eq!(err.message(), "Catalog name cannot be empty");
1230 }
1231 }
1232
1233 #[tokio::test]
1234 async fn test_builder_name_validation_with_missing_arn() {
1235 let builder = S3TablesCatalogBuilder::default();
1236
1237 let result = builder.load("", HashMap::new()).await;
1238
1239 assert!(result.is_err());
1240 if let Err(err) = result {
1241 assert_eq!(err.kind(), ErrorKind::DataInvalid);
1242 assert_eq!(err.message(), "Catalog name cannot be empty");
1243 }
1244 }
1245
1246 #[tokio::test]
1248 async fn test_s3tables_create_table_write_load_table_read() {
1249 use iceberg::writer::base_writer::data_file_writer::DataFileWriterBuilder;
1250 use iceberg::writer::file_writer::ParquetWriterBuilder;
1251 use iceberg::writer::file_writer::location_generator::{
1252 DefaultFileNameGenerator, DefaultLocationGenerator,
1253 };
1254 use iceberg::writer::file_writer::rolling_writer::RollingFileWriterBuilder;
1255 use iceberg::writer::{IcebergWriter, IcebergWriterBuilder};
1256
1257 let catalog = match load_s3tables_catalog_from_env().await {
1258 Ok(Some(c)) => c,
1259 Ok(None) => return,
1260 Err(e) => panic!("Error loading catalog: {e}"),
1261 };
1262
1263 let ns = NamespaceIdent::new(format!("test_rw_{}", uuid::Uuid::new_v4().simple()));
1264 catalog.create_namespace(&ns, HashMap::new()).await.unwrap();
1265
1266 let table_name = String::from("table");
1267
1268 let schema = Schema::builder()
1269 .with_fields(vec![
1270 NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
1271 ])
1272 .build()
1273 .unwrap();
1274 let creation = TableCreation::builder()
1275 .name(table_name.clone())
1276 .schema(schema)
1277 .build();
1278
1279 let table = catalog.create_table(&ns, creation).await.unwrap();
1280
1281 let arrow_schema: Arc<arrow_schema::Schema> = Arc::new(
1283 table
1284 .metadata()
1285 .current_schema()
1286 .as_ref()
1287 .try_into()
1288 .unwrap(),
1289 );
1290 let batch = arrow_array::RecordBatch::try_new(arrow_schema, vec![Arc::new(
1291 arrow_array::Int32Array::from(vec![42]),
1292 )])
1293 .unwrap();
1294
1295 let location_generator = DefaultLocationGenerator::new(table.metadata()).unwrap();
1297 let file_name_generator = DefaultFileNameGenerator::new(
1298 "test".to_string(),
1299 None,
1300 iceberg::spec::DataFileFormat::Parquet,
1301 );
1302 let parquet_writer_builder = ParquetWriterBuilder::new(
1303 parquet::file::properties::WriterProperties::default(),
1304 table.metadata().current_schema().clone(),
1305 );
1306 let rw = RollingFileWriterBuilder::new_with_default_file_size(
1307 parquet_writer_builder,
1308 table.file_io().clone(),
1309 location_generator,
1310 file_name_generator,
1311 );
1312 let mut writer = DataFileWriterBuilder::new(rw).build(None).await.unwrap();
1313 writer.write(batch.clone()).await.unwrap();
1314 let data_files = writer.close().await.unwrap();
1315
1316 let tx = Transaction::new(&table);
1317 let tx = tx
1318 .fast_append()
1319 .add_data_files(data_files)
1320 .apply(tx)
1321 .unwrap();
1322 tx.commit(&catalog).await.unwrap();
1323
1324 let table_ident = TableIdent::new(ns.clone(), table_name.clone());
1326 let reloaded = catalog.load_table(&table_ident).await.unwrap();
1327 let batches: Vec<arrow_array::RecordBatch> = reloaded
1328 .scan()
1329 .select_all()
1330 .build()
1331 .expect("scan to be valid (snapshot exists, schema is OK)")
1332 .to_arrow()
1333 .await
1334 .expect("scan tasks should be OK")
1335 .try_collect()
1336 .await
1337 .expect("scan should complete successfully");
1338
1339 assert_eq!(batches.len(), 1);
1340 assert_eq!(
1341 batches[0], batch,
1342 "read records should match records written earlier"
1343 );
1344
1345 catalog.purge_table(&table_ident).await.ok();
1347 catalog.drop_namespace(&ns).await.ok();
1348 }
1349}