1use std::collections::{HashMap, HashSet};
21use std::fmt::{Debug, Formatter};
22use std::future::Future;
23use std::num::NonZeroU32;
24use std::str::FromStr;
25use std::sync::{Arc, OnceLock};
26use std::time::Duration;
27
28use async_trait::async_trait;
29use iceberg::encryption::kms::{KeyManagementClient, KmsClientFactory};
30use iceberg::io::{
31 ADLS_SAS_TOKEN, FileIO, FileIOBuilder, GCS_TOKEN, S3_SESSION_TOKEN, StorageFactory,
32};
33use iceberg::table::Table;
34use iceberg::{
35 Catalog, CatalogBuilder, Error, ErrorKind, Namespace, NamespaceIdent, Result, Runtime,
36 SessionCatalog, SessionContext, TableCommit, TableCreation, TableIdent,
37};
38use itertools::Itertools;
39use reqwest::header::{
40 HeaderMap, HeaderName, HeaderValue, {self},
41};
42use reqwest::{Client, Method, StatusCode, Url};
43use tokio::sync::OnceCell;
44use typed_builder::TypedBuilder;
45
46use crate::auth::{
47 AUTH_TYPE_NONE, AUTH_TYPE_OAUTH2, AuthManager, AuthSession, NoopAuthManager, OAuth2Manager,
48};
49use crate::client::{
50 HttpClient, deserialize_catalog_response, deserialize_unexpected_catalog_error,
51};
52use crate::endpoint::{Endpoint, V1_NAMESPACE_EXISTS, V1_TABLE_EXISTS};
53use crate::request::HttpRequest;
54use crate::response::HttpResponse;
55use crate::types::{
56 CatalogConfig, CommitTableRequest, CommitTableResponse, CreateNamespaceRequest,
57 CreateTableRequest, ListNamespaceResponse, ListTablesResponse, LoadTableResult,
58 NamespaceResponse, RegisterTableRequest, RenameTableRequest,
59};
60
61pub const REST_CATALOG_PROP_URI: &str = "uri";
63pub const REST_CATALOG_PROP_WAREHOUSE: &str = "warehouse";
65pub const REST_CATALOG_PROP_PAGE_SIZE: &str = "rest-page-size";
68pub const REST_CATALOG_PROP_DISABLE_HEADER_REDACTION: &str = "disable-header-redaction";
71pub const REST_CATALOG_PROP_AUTH_TYPE: &str = "rest.auth.type";
75pub const REST_CATALOG_PROP_TABLE_CACHE_MAX_ENTRIES: &str = "rest-table-cache.max-entries";
81pub const REST_CATALOG_PROP_TABLE_CACHE_EXPIRE_AFTER_WRITE_MS: &str =
84 "rest-table-cache.expire-after-write-ms";
85
86const TABLE_CACHE_MAX_ENTRIES_DEFAULT: u64 = 100;
87const TABLE_CACHE_EXPIRE_AFTER_WRITE_MS_DEFAULT: u64 = 5 * 60 * 1000;
88const TABLE_CACHE_EXPIRE_AFTER_WRITE_MAX: Duration = Duration::from_secs(1000 * 365 * 24 * 3600);
90
91type TableCache = moka::sync::Cache<(String, TableIdent), (String, Table)>;
95
96fn is_short_lived_credential(key: &str) -> bool {
98 let key = key.to_ascii_lowercase();
99 key == S3_SESSION_TOKEN || key == GCS_TOKEN || key.starts_with(ADLS_SAS_TOKEN)
100}
101
102const ICEBERG_REST_SPEC_VERSION: &str = "0.14.1";
103const CARGO_PKG_VERSION: &str = env!("CARGO_PKG_VERSION");
104const PATH_V1: &str = "v1";
105
106#[derive(Debug, Default)]
112pub struct RestCatalogBuilder {
113 session_context: Option<SessionContext>,
114 inner: RestSessionCatalogBuilder,
115}
116
117impl CatalogBuilder for RestCatalogBuilder {
118 type C = RestCatalog;
119
120 fn with_storage_factory(mut self, storage_factory: Arc<dyn StorageFactory>) -> Self {
121 self.inner = self.inner.with_storage_factory(storage_factory);
122 self
123 }
124
125 fn with_kms_client_factory(mut self, kms_client_factory: Arc<dyn KmsClientFactory>) -> Self {
126 self.inner = self.inner.with_kms_client_factory(kms_client_factory);
127 self
128 }
129
130 fn with_runtime(mut self, runtime: Runtime) -> Self {
131 self.inner = self.inner.with_runtime(runtime);
132 self
133 }
134
135 fn load(
136 self,
137 name: impl Into<String>,
138 props: HashMap<String, String>,
139 ) -> impl Future<Output = Result<Self::C>> + Send {
140 let name = name.into();
141 async move {
142 let context = self.session_context.unwrap_or_else(SessionContext::empty);
143 let session_catalog = Arc::new(self.inner.load(name, props).await?);
144
145 Ok(RestCatalog::from_session_catalog(context, session_catalog))
146 }
147 }
148}
149
150impl RestCatalogBuilder {
151 pub fn with_page_size(mut self, page_size: NonZeroU32) -> Self {
154 self.inner = self.inner.with_page_size(page_size);
155 self
156 }
157
158 pub fn with_client(mut self, client: Client) -> Self {
160 self.inner = self.inner.with_client(client);
161 self
162 }
163
164 pub fn with_session_context(mut self, context: SessionContext) -> Self {
169 self.session_context = Some(context);
170 self
171 }
172
173 pub fn with_auth_manager<M>(mut self, auth_manager: M) -> Self
178 where M: AuthManager + 'static {
179 self.inner = self.inner.with_auth_manager(auth_manager);
180 self
181 }
182}
183
184#[derive(Clone, TypedBuilder)]
186pub(crate) struct RestCatalogConfig {
187 #[builder(default, setter(strip_option))]
188 name: Option<String>,
189
190 uri: String,
191
192 #[builder(default, setter(strip_option(fallback = warehouse_opt)))]
193 warehouse: Option<String>,
194
195 #[builder(default)]
196 props: HashMap<String, String>,
197
198 #[builder(default)]
199 client: Option<Client>,
200
201 #[builder(default)]
205 default_client: Arc<OnceLock<Client>>,
206}
207
208fn is_sensitive_prop(key: &str) -> bool {
211 key.contains("token")
212 || key.contains("credential")
213 || key.contains("secret")
214 || key.contains("password")
215 || key.contains("key")
216 || key.contains("connection-string")
217 || key.starts_with("header.")
218}
219
220impl Debug for RestCatalogConfig {
223 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
224 let props: HashMap<&str, &str> = self
225 .props
226 .iter()
227 .map(|(key, value)| {
228 let value = if is_sensitive_prop(key) {
229 "[REDACTED]"
230 } else {
231 value.as_str()
232 };
233 (key.as_str(), value)
234 })
235 .collect();
236 f.debug_struct("RestCatalogConfig")
237 .field("name", &self.name)
238 .field("uri", &self.uri)
239 .field("warehouse", &self.warehouse)
240 .field("props", &props)
241 .finish_non_exhaustive()
242 }
243}
244
245impl RestCatalogConfig {
246 fn url_prefixed(&self, parts: &[&str]) -> String {
247 [&self.uri, PATH_V1]
248 .into_iter()
249 .chain(self.props.get("prefix").map(|s| &**s))
250 .chain(parts.iter().cloned())
251 .join("/")
252 }
253
254 fn config_endpoint(&self) -> String {
255 [&self.uri, PATH_V1, "config"].join("/")
256 }
257
258 pub(crate) fn get_token_endpoint(&self) -> String {
259 self.explicit_oauth2_server_uri()
260 .unwrap_or_else(|| default_token_endpoint(&self.uri))
261 }
262
263 pub(crate) fn explicit_oauth2_server_uri(&self) -> Option<String> {
265 self.props.get("oauth2-server-uri").cloned()
266 }
267
268 fn page_size(&self) -> Result<Option<NonZeroU32>> {
270 self.props
271 .get(REST_CATALOG_PROP_PAGE_SIZE)
272 .map(|value| {
273 let invalid_page_size = || {
274 Error::new(
275 ErrorKind::DataInvalid,
276 format!("Invalid {REST_CATALOG_PROP_PAGE_SIZE}: {value}"),
277 )
278 };
279 let page_size = value
280 .parse::<NonZeroU32>()
281 .map_err(|err| invalid_page_size().with_source(err))?;
282 if page_size.get() > i32::MAX as u32 {
283 return Err(invalid_page_size());
284 }
285 Ok(page_size)
286 })
287 .transpose()
288 }
289
290 fn namespaces_endpoint(&self) -> String {
291 self.url_prefixed(&["namespaces"])
292 }
293
294 fn namespace_endpoint(&self, ns: &NamespaceIdent) -> String {
295 self.url_prefixed(&["namespaces", &ns.to_url_string()])
296 }
297
298 fn tables_endpoint(&self, ns: &NamespaceIdent) -> String {
299 self.url_prefixed(&["namespaces", &ns.to_url_string(), "tables"])
300 }
301
302 fn rename_table_endpoint(&self) -> String {
303 self.url_prefixed(&["tables", "rename"])
304 }
305
306 fn register_table_endpoint(&self, ns: &NamespaceIdent) -> String {
307 self.url_prefixed(&["namespaces", &ns.to_url_string(), "register"])
308 }
309
310 fn table_endpoint(&self, table: &TableIdent) -> String {
311 self.url_prefixed(&[
312 "namespaces",
313 &table.namespace.to_url_string(),
314 "tables",
315 &table.name,
316 ])
317 }
318
319 pub(crate) fn client(&self) -> Client {
323 self.client
324 .clone()
325 .unwrap_or_else(|| self.default_client.get_or_init(Client::default).clone())
326 }
327
328 pub(crate) fn token(&self) -> Option<String> {
332 self.props.get("token").cloned()
333 }
334
335 pub(crate) fn credential(&self) -> Option<(Option<String>, String)> {
338 credential_from_props(&self.props)
339 }
340
341 pub(crate) fn extra_headers(&self) -> Result<HeaderMap> {
343 extra_headers_from_props(&self.props)
344 }
345
346 pub(crate) fn extra_oauth_params(&self) -> HashMap<String, String> {
348 oauth_params_from_props(&self.props)
349 }
350
351 pub(crate) fn disable_header_redaction(&self) -> bool {
356 self.props
357 .get(REST_CATALOG_PROP_DISABLE_HEADER_REDACTION)
358 .map(|v| v.eq_ignore_ascii_case("true"))
359 .unwrap_or(false)
360 }
361
362 fn table_cache(&self) -> Result<Option<TableCache>> {
365 let parse = |key: &str, default: u64| -> Result<u64> {
366 self.props.get(key).map_or(Ok(default), |v| {
367 v.parse().map_err(|e| {
368 Error::new(ErrorKind::DataInvalid, format!("Invalid {key}: {v}")).with_source(e)
369 })
370 })
371 };
372 let expire_after_write_ms = parse(
373 REST_CATALOG_PROP_TABLE_CACHE_EXPIRE_AFTER_WRITE_MS,
374 TABLE_CACHE_EXPIRE_AFTER_WRITE_MS_DEFAULT,
375 )?;
376 let expire_after_write = Duration::from_millis(expire_after_write_ms);
377 if expire_after_write.is_zero() || expire_after_write > TABLE_CACHE_EXPIRE_AFTER_WRITE_MAX {
379 return Err(Error::new(
380 ErrorKind::DataInvalid,
381 format!(
382 "Invalid {REST_CATALOG_PROP_TABLE_CACHE_EXPIRE_AFTER_WRITE_MS}: {expire_after_write_ms} (must be between 1 ms and 1000 years)"
383 ),
384 ));
385 }
386 let max_entries = parse(
387 REST_CATALOG_PROP_TABLE_CACHE_MAX_ENTRIES,
388 TABLE_CACHE_MAX_ENTRIES_DEFAULT,
389 )?;
390 Ok((max_entries > 0).then(|| {
391 TableCache::builder()
392 .max_capacity(max_entries)
393 .time_to_live(expire_after_write)
394 .build()
395 }))
396 }
397
398 pub(crate) fn merge_with_config(mut self, mut config: CatalogConfig) -> Self {
400 if let Some(uri) = config.overrides.remove(REST_CATALOG_PROP_URI) {
401 self.uri = uri;
402 }
403
404 let mut props = config.defaults;
405 props.extend(self.props);
406 if let Some(warehouse) = &self.warehouse {
409 props.insert(REST_CATALOG_PROP_WAREHOUSE.to_string(), warehouse.clone());
410 }
411 props.extend(config.overrides);
412
413 self.props = props;
414 self
415 }
416}
417
418pub(crate) fn credential_from_props(
426 props: &HashMap<String, String>,
427) -> Option<(Option<String>, String)> {
428 let cred = props.get("credential")?;
429
430 match cred.split_once(':') {
431 Some((client_id, client_secret)) => {
432 Some((Some(client_id.to_string()), client_secret.to_string()))
433 }
434 None => Some((None, cred.to_string())),
435 }
436}
437
438pub(crate) fn extra_headers_from_props(props: &HashMap<String, String>) -> Result<HeaderMap> {
445 let mut headers = HeaderMap::from_iter([
446 (
447 header::CONTENT_TYPE,
448 HeaderValue::from_static("application/json"),
449 ),
450 (
451 HeaderName::from_static("x-client-version"),
452 HeaderValue::from_static(ICEBERG_REST_SPEC_VERSION),
453 ),
454 (
455 header::USER_AGENT,
456 HeaderValue::from_str(&format!("iceberg-rs/{CARGO_PKG_VERSION}")).unwrap(),
457 ),
458 ]);
459
460 headers.extend(explicit_headers_from_props(props)?);
461
462 Ok(headers)
463}
464
465pub(crate) fn default_token_endpoint(uri: &str) -> String {
467 [uri, PATH_V1, "oauth", "tokens"].join("/")
468}
469
470pub(crate) fn explicit_headers_from_props(props: &HashMap<String, String>) -> Result<HeaderMap> {
472 let mut headers = HeaderMap::new();
473 for (key, value) in props
474 .iter()
475 .filter_map(|(k, v)| k.strip_prefix("header.").map(|k| (k, v)))
476 {
477 headers.insert(
478 HeaderName::from_str(key).map_err(|e| {
479 Error::new(
480 ErrorKind::DataInvalid,
481 format!("Invalid header name: {key}"),
482 )
483 .with_source(e)
484 })?,
485 HeaderValue::from_str(value).map_err(|e| {
486 Error::new(
487 ErrorKind::DataInvalid,
488 format!("Invalid value for header: {key}"),
490 )
491 .with_source(e)
492 })?,
493 );
494 }
495
496 Ok(headers)
497}
498
499pub(crate) fn oauth_params_from_props(props: &HashMap<String, String>) -> HashMap<String, String> {
501 let mut params = HashMap::new();
502
503 if let Some(scope) = props.get("scope") {
504 params.insert("scope".to_string(), scope.to_string());
505 } else {
506 params.insert("scope".to_string(), "catalog".to_string());
507 }
508
509 let optional_params = ["audience", "resource"];
510 for param_name in optional_params {
511 if let Some(value) = props.get(param_name) {
512 params.insert(param_name.to_string(), value.to_string());
513 }
514 }
515
516 params
517}
518
519struct RestCatalogClient {
520 auth_manager: Arc<dyn AuthManager>,
523 catalog_session: Arc<dyn AuthSession>,
525 http_client: HttpClient,
528 config: RestCatalogConfig,
532 endpoints: HashSet<Endpoint>,
534 page_size: Option<NonZeroU32>,
536 table_cache: Option<TableCache>,
538}
539
540impl Debug for RestCatalogClient {
541 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
542 f.debug_struct("RestCatalogClient")
545 .field("http_client", &self.http_client)
546 .field("config", &self.config)
547 .field("endpoints", &self.endpoints)
548 .finish_non_exhaustive()
549 }
550}
551
552impl RestCatalogClient {
553 async fn init(
556 user_config: &RestCatalogConfig,
557 auth_manager: Arc<dyn AuthManager>,
558 ) -> Result<Self> {
559 let http_client = HttpClient::new(user_config)?;
560 let catalog_config = {
564 let init_session = auth_manager
565 .init_session(&http_client, &Self::auth_props(user_config))
566 .await?;
567 Self::load_config(&http_client, init_session.as_ref(), user_config).await?
568 };
569 let endpoints = match &catalog_config.endpoints {
572 Some(advertised) if !advertised.is_empty() => advertised.iter().cloned().collect(),
573 _ => crate::endpoint::DEFAULT_ENDPOINTS.clone(),
574 };
575 let config = user_config.clone().merge_with_config(catalog_config);
576 let page_size = config.page_size()?;
577 let http_client = http_client.update_with(&config)?;
578 let catalog_session = auth_manager
581 .catalog_session(&http_client, &Self::auth_props(&config))
582 .await?;
583
584 Ok(Self {
585 auth_manager,
586 catalog_session,
587 table_cache: config.table_cache()?,
588 config,
589 http_client,
590 endpoints,
591 page_size,
592 })
593 }
594
595 fn invalidate_table(&self, context: &SessionContext, table: &TableIdent) {
596 if let Some(cache) = &self.table_cache {
597 cache.invalidate(&(context.session_id().to_string(), table.clone()));
598 }
599 }
600
601 #[cfg(test)]
603 async fn token(&self) -> Option<String> {
604 self.http_client.token(self.catalog_session.as_ref()).await
605 }
606
607 async fn query_catalog(
609 &self,
610 context: &SessionContext,
611 request: HttpRequest,
612 ) -> Result<HttpResponse> {
613 let session = self
614 .auth_manager
615 .contextual_session(context, Arc::clone(&self.catalog_session))
616 .await?;
617 self.http_client
618 .query_catalog(session.as_ref(), request)
619 .await
620 }
621
622 fn auth_props(config: &RestCatalogConfig) -> HashMap<String, String> {
625 let mut props = config.props.clone();
630 props.insert(REST_CATALOG_PROP_URI.to_string(), config.uri.clone());
631 if let Some(warehouse) = &config.warehouse {
632 props
635 .entry(REST_CATALOG_PROP_WAREHOUSE.to_string())
636 .or_insert_with(|| warehouse.clone());
637 }
638 props
639 }
640
641 async fn load_config(
645 http_client: &HttpClient,
646 auth_session: &dyn AuthSession,
647 user_config: &RestCatalogConfig,
648 ) -> Result<CatalogConfig> {
649 let mut request_builder = http_client.request(Method::GET, user_config.config_endpoint());
650
651 if let Some(warehouse_location) = &user_config.warehouse {
652 request_builder = request_builder.query(&[("warehouse", warehouse_location)]);
653 }
654
655 let request = HttpRequest::build(request_builder)?;
656
657 let http_response = http_client.query_catalog(auth_session, request).await?;
658
659 match http_response.status() {
660 StatusCode::OK => deserialize_catalog_response(http_response),
661 _ => Err(deserialize_unexpected_catalog_error(
662 http_response,
663 http_client.disable_header_redaction(),
664 )),
665 }
666 }
667}
668
669#[derive(Debug)]
675pub struct RestCatalog {
676 session_context: SessionContext,
677 inner: Arc<RestSessionCatalog>,
678}
679
680impl RestCatalog {
681 #[cfg(test)]
683 fn new(
684 context: SessionContext,
685 config: RestCatalogConfig,
686 auth_manager: Option<Box<dyn AuthManager>>,
687 storage_factory: Option<Arc<dyn StorageFactory>>,
688 runtime: Runtime,
689 kms_client: Option<Arc<dyn KeyManagementClient>>,
690 ) -> Self {
691 let session_catalog = Arc::new(RestSessionCatalog::new(
692 config,
693 auth_manager,
694 storage_factory,
695 runtime,
696 kms_client,
697 ));
698
699 Self::from_session_catalog(context, session_catalog)
700 }
701
702 fn from_session_catalog(context: SessionContext, inner: Arc<RestSessionCatalog>) -> Self {
703 Self {
704 session_context: context,
705 inner,
706 }
707 }
708
709 #[cfg(test)]
710 async fn client(&self) -> Result<&RestCatalogClient> {
711 self.inner.client().await
712 }
713}
714
715#[async_trait]
719impl Catalog for RestCatalog {
720 async fn list_namespaces(
721 &self,
722 parent: Option<&NamespaceIdent>,
723 ) -> Result<Vec<NamespaceIdent>> {
724 self.inner
725 .list_namespaces(&self.session_context, parent)
726 .await
727 }
728
729 async fn create_namespace(
730 &self,
731 namespace: &NamespaceIdent,
732 properties: HashMap<String, String>,
733 ) -> Result<Namespace> {
734 self.inner
735 .create_namespace(&self.session_context, namespace, properties)
736 .await
737 }
738
739 async fn get_namespace(&self, namespace: &NamespaceIdent) -> Result<Namespace> {
740 self.inner
741 .get_namespace(&self.session_context, namespace)
742 .await
743 }
744
745 async fn namespace_exists(&self, ns: &NamespaceIdent) -> Result<bool> {
746 self.inner.namespace_exists(&self.session_context, ns).await
747 }
748
749 async fn update_namespace(
750 &self,
751 namespace: &NamespaceIdent,
752 properties: HashMap<String, String>,
753 ) -> Result<()> {
754 self.inner
755 .update_namespace(&self.session_context, namespace, properties)
756 .await
757 }
758
759 async fn drop_namespace(&self, namespace: &NamespaceIdent) -> Result<()> {
760 self.inner
761 .drop_namespace(&self.session_context, namespace)
762 .await
763 }
764
765 async fn list_tables(&self, namespace: &NamespaceIdent) -> Result<Vec<TableIdent>> {
766 self.inner
767 .list_tables(&self.session_context, namespace)
768 .await
769 }
770
771 async fn create_table(
772 &self,
773 namespace: &NamespaceIdent,
774 creation: TableCreation,
775 ) -> Result<Table> {
776 self.inner
777 .create_table(&self.session_context, namespace, creation)
778 .await
779 }
780
781 async fn load_table(&self, table_ident: &TableIdent) -> Result<Table> {
782 self.inner
783 .load_table(&self.session_context, table_ident)
784 .await
785 }
786
787 async fn drop_table(&self, table: &TableIdent) -> Result<()> {
788 self.inner.drop_table(&self.session_context, table).await
789 }
790
791 async fn purge_table(&self, table: &TableIdent) -> Result<()> {
792 self.inner.purge_table(&self.session_context, table).await
793 }
794
795 async fn table_exists(&self, table: &TableIdent) -> Result<bool> {
796 self.inner.table_exists(&self.session_context, table).await
797 }
798
799 async fn rename_table(&self, src: &TableIdent, dest: &TableIdent) -> Result<()> {
800 self.inner
801 .rename_table(&self.session_context, src, dest)
802 .await
803 }
804
805 async fn register_table(
806 &self,
807 table_ident: &TableIdent,
808 metadata_location: String,
809 ) -> Result<Table> {
810 self.inner
811 .register_table(&self.session_context, table_ident, metadata_location)
812 .await
813 }
814
815 async fn update_table(&self, commit: TableCommit) -> Result<Table> {
816 self.inner.update_table(&self.session_context, commit).await
817 }
818}
819
820#[derive(Debug)]
826pub struct RestSessionCatalog {
827 auth_manager_override: Option<Arc<dyn AuthManager>>,
831 user_config: RestCatalogConfig,
835 client: OnceCell<RestCatalogClient>,
836 storage_factory: Option<Arc<dyn StorageFactory>>,
838 runtime: Runtime,
839 kms_client: Option<Arc<dyn KeyManagementClient>>,
841}
842
843impl RestSessionCatalog {
844 fn new(
846 config: RestCatalogConfig,
847 auth_manager_override: Option<Box<dyn AuthManager>>,
848 storage_factory: Option<Arc<dyn StorageFactory>>,
849 runtime: Runtime,
850 kms_client: Option<Arc<dyn KeyManagementClient>>,
851 ) -> Self {
852 Self {
853 auth_manager_override: auth_manager_override.map(Arc::from),
854 user_config: config,
855 client: OnceCell::new(),
856 storage_factory,
857 runtime,
858 kms_client,
859 }
860 }
861
862 async fn delete_table(
864 &self,
865 context: &SessionContext,
866 table: &TableIdent,
867 purge: bool,
868 ) -> Result<()> {
869 let client = self.client().await?;
870
871 let mut request_builder = client
872 .http_client
873 .request(Method::DELETE, client.config.table_endpoint(table));
874
875 if purge {
876 request_builder = request_builder.query(&[("purgeRequested", "true")]);
877 }
878
879 let request = HttpRequest::build(request_builder)?;
880 client.invalidate_table(context, table);
881 let http_response = client.query_catalog(context, request).await?;
882
883 match http_response.status() {
884 StatusCode::NO_CONTENT | StatusCode::OK => Ok(()),
885 StatusCode::NOT_FOUND => Err(Error::new(
886 ErrorKind::TableNotFound,
887 "Tried to drop a table that does not exist",
888 )),
889 _ => Err(deserialize_unexpected_catalog_error(
890 http_response,
891 client.http_client.disable_header_redaction(),
892 )),
893 }
894 }
895
896 fn auth_type(config: &RestCatalogConfig) -> String {
901 config
902 .props
903 .get(REST_CATALOG_PROP_AUTH_TYPE)
904 .map(|auth_type| auth_type.to_ascii_lowercase())
906 .unwrap_or_else(|| {
907 if config.token().is_some()
908 || config.credential().is_some()
909 || config.explicit_oauth2_server_uri().is_some()
910 {
911 AUTH_TYPE_OAUTH2.to_string()
912 } else {
913 AUTH_TYPE_NONE.to_string()
914 }
915 })
916 }
917
918 fn resolve_auth_manager(&self) -> Result<Arc<dyn AuthManager>> {
921 if let Some(auth_manager) = &self.auth_manager_override {
922 return Ok(auth_manager.clone());
923 }
924 let config = &self.user_config;
925 let auth_type = Self::auth_type(config);
926 if auth_type == AUTH_TYPE_OAUTH2 && !config.props.contains_key(REST_CATALOG_PROP_AUTH_TYPE)
929 {
930 tracing::warn!(
931 "Inferring {REST_CATALOG_PROP_AUTH_TYPE}={AUTH_TYPE_OAUTH2} from the configured \
932 OAuth properties; set it explicitly to avoid this warning"
933 );
934 }
935 match auth_type.as_str() {
936 AUTH_TYPE_NONE => Ok(Arc::new(NoopAuthManager)),
937 AUTH_TYPE_OAUTH2 => Ok(Arc::new(OAuth2Manager::from_config(config)?)),
938 other => Err(Error::new(
939 ErrorKind::DataInvalid,
940 format!(
941 "unknown '{REST_CATALOG_PROP_AUTH_TYPE}': {other}; use \
942 `RestSessionCatalogBuilder::with_auth_manager` or \
943 `RestCatalogBuilder::with_auth_manager` to inject a custom auth manager"
944 ),
945 )),
946 }
947 }
948
949 async fn client(&self) -> Result<&RestCatalogClient> {
951 self.client
952 .get_or_try_init(|| async {
953 RestCatalogClient::init(&self.user_config, self.resolve_auth_manager()?).await
954 })
955 .await
956 }
957
958 pub(crate) async fn supports_endpoint(&self, endpoint: &Endpoint) -> Result<bool> {
962 Ok(self.client().await?.endpoints.contains(endpoint))
963 }
964
965 async fn check_exists_via_head(
968 &self,
969 context: &SessionContext,
970 client: &RestCatalogClient,
971 url: String,
972 ) -> Result<bool> {
973 let request = HttpRequest::build(client.http_client.request(Method::HEAD, url))?;
974 let http_response = client.query_catalog(context, request).await?;
975
976 match http_response.status() {
977 StatusCode::NO_CONTENT | StatusCode::OK => Ok(true),
978 StatusCode::NOT_FOUND => Ok(false),
979 _ => Err(deserialize_unexpected_catalog_error(
980 http_response,
981 client.http_client.disable_header_redaction(),
982 )),
983 }
984 }
985
986 async fn load_file_io(
987 &self,
988 metadata_location: Option<&str>,
989 extra_config: Option<HashMap<String, String>>,
990 ) -> Result<FileIO> {
991 let mut props = self.client().await?.config.props.clone();
992 if let Some(config) = extra_config {
993 props.extend(config);
994 }
995
996 let warehouse_path = match self.client().await?.config.warehouse.as_deref() {
999 Some(url) if Url::parse(url).is_ok() => Some(url),
1000 Some(_) => None,
1001 None => None,
1002 };
1003
1004 if metadata_location.or(warehouse_path).is_none() {
1005 return Err(Error::new(
1006 ErrorKind::Unexpected,
1007 "Unable to load file io, neither warehouse nor metadata location is set!",
1008 ));
1009 }
1010
1011 let factory = self
1013 .storage_factory
1014 .clone()
1015 .ok_or_else(|| {
1016 Error::new(
1017 ErrorKind::Unexpected,
1018 "StorageFactory must be provided for REST catalog table operations. Use `with_storage_factory` to configure it.",
1019 )
1020 })?;
1021
1022 let file_io = FileIOBuilder::new(factory).with_props(props).build();
1023
1024 Ok(file_io)
1025 }
1026}
1027
1028#[async_trait]
1031impl SessionCatalog for RestSessionCatalog {
1032 async fn list_namespaces(
1033 &self,
1034 context: &SessionContext,
1035 parent: Option<&NamespaceIdent>,
1036 ) -> Result<Vec<NamespaceIdent>> {
1037 let client = self.client().await?;
1038 let endpoint = client.config.namespaces_endpoint();
1039 let mut namespaces = Vec::new();
1040 let mut next_token = client.page_size.map(|_| String::new());
1042
1043 loop {
1044 let mut request = client.http_client.request(Method::GET, endpoint.clone());
1045
1046 if let Some(page_size) = client.page_size {
1047 request = request.query(&[("pageSize", page_size)]);
1048 }
1049
1050 if let Some(ns) = parent {
1052 request = request.query(&[("parent", ns.to_url_string())]);
1053 }
1054
1055 if let Some(token) = next_token {
1056 request = request.query(&[("pageToken", token)]);
1057 }
1058
1059 let http_response = client
1060 .query_catalog(context, HttpRequest::build(request)?)
1061 .await?;
1062
1063 match http_response.status() {
1064 StatusCode::OK => {
1065 let response =
1066 deserialize_catalog_response::<ListNamespaceResponse>(http_response)?;
1067
1068 namespaces.extend(response.namespaces);
1069
1070 match response.next_page_token {
1071 Some(token) => next_token = Some(token),
1072 None => break,
1073 }
1074 }
1075 StatusCode::NOT_FOUND => {
1076 return Err(Error::new(
1077 ErrorKind::NamespaceNotFound,
1078 "The parent parameter of the namespace provided does not exist",
1079 ));
1080 }
1081 _ => {
1082 return Err(deserialize_unexpected_catalog_error(
1083 http_response,
1084 client.http_client.disable_header_redaction(),
1085 ));
1086 }
1087 }
1088 }
1089
1090 Ok(namespaces)
1091 }
1092
1093 async fn create_namespace(
1094 &self,
1095 context: &SessionContext,
1096 namespace: &NamespaceIdent,
1097 properties: HashMap<String, String>,
1098 ) -> Result<Namespace> {
1099 let client = self.client().await?;
1100
1101 let request = HttpRequest::build(
1102 client
1103 .http_client
1104 .request(Method::POST, client.config.namespaces_endpoint())
1105 .json(&CreateNamespaceRequest {
1106 namespace: namespace.clone(),
1107 properties,
1108 }),
1109 )?;
1110
1111 let http_response = client.query_catalog(context, request).await?;
1112
1113 match http_response.status() {
1114 StatusCode::OK => {
1115 let response = deserialize_catalog_response::<NamespaceResponse>(http_response)?;
1116 Ok(Namespace::from(response))
1117 }
1118 StatusCode::CONFLICT => Err(Error::new(
1119 ErrorKind::NamespaceAlreadyExists,
1120 "Tried to create a namespace that already exists",
1121 )),
1122 _ => Err(deserialize_unexpected_catalog_error(
1123 http_response,
1124 client.http_client.disable_header_redaction(),
1125 )),
1126 }
1127 }
1128
1129 async fn get_namespace(
1130 &self,
1131 context: &SessionContext,
1132 namespace: &NamespaceIdent,
1133 ) -> Result<Namespace> {
1134 let client = self.client().await?;
1135
1136 let request = HttpRequest::build(
1137 client
1138 .http_client
1139 .request(Method::GET, client.config.namespace_endpoint(namespace)),
1140 )?;
1141
1142 let http_response = client.query_catalog(context, request).await?;
1143
1144 match http_response.status() {
1145 StatusCode::OK => {
1146 let response = deserialize_catalog_response::<NamespaceResponse>(http_response)?;
1147 Ok(Namespace::from(response))
1148 }
1149 StatusCode::NOT_FOUND => Err(Error::new(
1150 ErrorKind::NamespaceNotFound,
1151 "Tried to get a namespace that does not exist",
1152 )),
1153 _ => Err(deserialize_unexpected_catalog_error(
1154 http_response,
1155 client.http_client.disable_header_redaction(),
1156 )),
1157 }
1158 }
1159
1160 async fn namespace_exists(
1161 &self,
1162 context: &SessionContext,
1163 ns: &NamespaceIdent,
1164 ) -> Result<bool> {
1165 if !self.supports_endpoint(&V1_NAMESPACE_EXISTS).await? {
1170 return match self.get_namespace(context, ns).await {
1171 Ok(_) => Ok(true),
1172 Err(e) if e.kind() == ErrorKind::NamespaceNotFound => Ok(false),
1173 Err(e) => Err(e),
1174 };
1175 }
1176
1177 let client = self.client().await?;
1178 self.check_exists_via_head(context, client, client.config.namespace_endpoint(ns))
1179 .await
1180 }
1181
1182 async fn update_namespace(
1183 &self,
1184 _context: &SessionContext,
1185 _namespace: &NamespaceIdent,
1186 _properties: HashMap<String, String>,
1187 ) -> Result<()> {
1188 Err(Error::new(
1189 ErrorKind::FeatureUnsupported,
1190 "Updating namespace not supported yet!",
1191 ))
1192 }
1193
1194 async fn drop_namespace(
1195 &self,
1196 context: &SessionContext,
1197 namespace: &NamespaceIdent,
1198 ) -> Result<()> {
1199 let client = self.client().await?;
1200
1201 let request = HttpRequest::build(
1202 client
1203 .http_client
1204 .request(Method::DELETE, client.config.namespace_endpoint(namespace)),
1205 )?;
1206
1207 let http_response = client.query_catalog(context, request).await?;
1208
1209 match http_response.status() {
1210 StatusCode::NO_CONTENT | StatusCode::OK => Ok(()),
1211 StatusCode::NOT_FOUND => Err(Error::new(
1212 ErrorKind::NamespaceNotFound,
1213 "Tried to drop a namespace that does not exist",
1214 )),
1215 _ => Err(deserialize_unexpected_catalog_error(
1216 http_response,
1217 client.http_client.disable_header_redaction(),
1218 )),
1219 }
1220 }
1221
1222 async fn list_tables(
1223 &self,
1224 context: &SessionContext,
1225 namespace: &NamespaceIdent,
1226 ) -> Result<Vec<TableIdent>> {
1227 let client = self.client().await?;
1228 let endpoint = client.config.tables_endpoint(namespace);
1229 let mut identifiers = Vec::new();
1230 let mut next_token = client.page_size.map(|_| String::new());
1232
1233 loop {
1234 let mut request = client.http_client.request(Method::GET, endpoint.clone());
1235
1236 if let Some(page_size) = client.page_size {
1237 request = request.query(&[("pageSize", page_size)]);
1238 }
1239
1240 if let Some(token) = next_token {
1241 request = request.query(&[("pageToken", token)]);
1242 }
1243
1244 let http_response = client
1245 .query_catalog(context, HttpRequest::build(request)?)
1246 .await?;
1247
1248 match http_response.status() {
1249 StatusCode::OK => {
1250 let response =
1251 deserialize_catalog_response::<ListTablesResponse>(http_response)?;
1252
1253 identifiers.extend(response.identifiers);
1254
1255 match response.next_page_token {
1256 Some(token) => next_token = Some(token),
1257 None => break,
1258 }
1259 }
1260 StatusCode::NOT_FOUND => {
1261 return Err(Error::new(
1262 ErrorKind::NamespaceNotFound,
1263 "Tried to list tables of a namespace that does not exist",
1264 ));
1265 }
1266 _ => {
1267 return Err(deserialize_unexpected_catalog_error(
1268 http_response,
1269 client.http_client.disable_header_redaction(),
1270 ));
1271 }
1272 }
1273 }
1274
1275 Ok(identifiers)
1276 }
1277
1278 async fn create_table(
1285 &self,
1286 context: &SessionContext,
1287 namespace: &NamespaceIdent,
1288 creation: TableCreation,
1289 ) -> Result<Table> {
1290 let client = self.client().await?;
1291
1292 let table_ident = TableIdent::new(namespace.clone(), creation.name.clone());
1293
1294 let request = HttpRequest::build(
1295 client
1296 .http_client
1297 .request(Method::POST, client.config.tables_endpoint(namespace))
1298 .json(&CreateTableRequest {
1299 name: creation.name,
1300 location: creation.location,
1301 schema: creation.schema,
1302 partition_spec: creation.partition_spec,
1303 write_order: creation.sort_order,
1304 stage_create: Some(false),
1305 properties: creation.properties,
1306 }),
1307 )?;
1308
1309 let http_response = client.query_catalog(context, request).await?;
1310
1311 let response = match http_response.status() {
1312 StatusCode::OK => deserialize_catalog_response::<LoadTableResult>(http_response)?,
1313 StatusCode::NOT_FOUND => {
1314 return Err(Error::new(
1315 ErrorKind::NamespaceNotFound,
1316 "Tried to create a table under a namespace that does not exist",
1317 ));
1318 }
1319 StatusCode::CONFLICT => {
1320 return Err(Error::new(
1321 ErrorKind::TableAlreadyExists,
1322 "The table already exists",
1323 ));
1324 }
1325 _ => {
1326 return Err(deserialize_unexpected_catalog_error(
1327 http_response,
1328 client.http_client.disable_header_redaction(),
1329 ));
1330 }
1331 };
1332
1333 let metadata_location = response.metadata_location.as_ref().ok_or(Error::new(
1334 ErrorKind::DataInvalid,
1335 "Metadata location missing in `create_table` response!",
1336 ))?;
1337
1338 let config = response
1339 .config
1340 .into_iter()
1341 .chain(self.user_config.props.clone())
1342 .collect();
1343
1344 let file_io = self
1345 .load_file_io(Some(metadata_location), Some(config))
1346 .await?;
1347
1348 let mut table_builder = Table::builder()
1349 .identifier(table_ident.clone())
1350 .file_io(file_io)
1351 .metadata(response.metadata)
1352 .runtime(self.runtime.clone());
1353 if let Some(kms_client) = self.kms_client.clone() {
1354 table_builder = table_builder.kms_client(kms_client);
1355 }
1356
1357 if let Some(metadata_location) = response.metadata_location {
1358 table_builder.metadata_location(metadata_location).build()
1359 } else {
1360 table_builder.build()
1361 }
1362 }
1363
1364 async fn load_table(
1370 &self,
1371 context: &SessionContext,
1372 table_ident: &TableIdent,
1373 ) -> Result<Table> {
1374 let client = self.client().await?;
1375 let cache = client.table_cache.as_ref().map(|cache| {
1376 let key = (context.session_id().to_string(), table_ident.clone());
1377 (cache, key)
1378 });
1379 let cached = cache.as_ref().and_then(|(cache, key)| cache.get(key));
1380
1381 let mut request_builder = client
1382 .http_client
1383 .request(Method::GET, client.config.table_endpoint(table_ident));
1384 if let Some((etag, _)) = &cached {
1385 request_builder = request_builder.header(header::IF_NONE_MATCH, etag);
1386 }
1387 let request = HttpRequest::build(request_builder)?;
1388
1389 let http_response = client.query_catalog(context, request).await?;
1390
1391 let (response, etag) = match http_response.status() {
1392 StatusCode::OK => {
1393 let etag = http_response
1394 .headers()
1395 .get(header::ETAG)
1396 .and_then(|v| v.to_str().ok())
1397 .map(str::to_string);
1398 (
1399 deserialize_catalog_response::<LoadTableResult>(http_response)?,
1400 etag,
1401 )
1402 }
1403 StatusCode::NOT_MODIFIED => {
1404 return cached.map(|(_, table)| table).ok_or_else(|| {
1405 Error::new(
1406 ErrorKind::Unexpected,
1407 "Received 304 Not Modified for a table load without If-None-Match",
1408 )
1409 });
1410 }
1411 StatusCode::NOT_FOUND => {
1412 client.invalidate_table(context, table_ident);
1413 return Err(Error::new(
1414 ErrorKind::TableNotFound,
1415 "Tried to load a table that does not exist",
1416 ));
1417 }
1418 _ => {
1419 return Err(deserialize_unexpected_catalog_error(
1420 http_response,
1421 client.http_client.disable_header_redaction(),
1422 ));
1423 }
1424 };
1425
1426 let vends_credentials = response
1427 .storage_credentials
1428 .as_ref()
1429 .is_some_and(|credentials| !credentials.is_empty())
1430 || response
1431 .config
1432 .keys()
1433 .any(|key| is_short_lived_credential(key));
1434
1435 let config = response
1436 .config
1437 .into_iter()
1438 .chain(self.user_config.props.clone())
1439 .collect();
1440
1441 let file_io = self
1442 .load_file_io(response.metadata_location.as_deref(), Some(config))
1443 .await?;
1444
1445 let mut table_builder = Table::builder()
1446 .identifier(table_ident.clone())
1447 .file_io(file_io)
1448 .metadata(response.metadata)
1449 .runtime(self.runtime.clone());
1450 if let Some(kms_client) = self.kms_client.clone() {
1451 table_builder = table_builder.kms_client(kms_client);
1452 }
1453
1454 let table = if let Some(metadata_location) = response.metadata_location {
1455 table_builder.metadata_location(metadata_location).build()
1456 } else {
1457 table_builder.build()
1458 }?;
1459
1460 if let Some((cache, key)) = cache {
1463 match etag {
1464 Some(etag) if !vends_credentials => cache.insert(key, (etag, table.clone())),
1465 _ => cache.invalidate(&key),
1466 }
1467 }
1468 Ok(table)
1469 }
1470
1471 async fn drop_table(&self, context: &SessionContext, table: &TableIdent) -> Result<()> {
1473 self.delete_table(context, table, false).await
1474 }
1475
1476 async fn purge_table(&self, context: &SessionContext, table: &TableIdent) -> Result<()> {
1479 self.delete_table(context, table, true).await
1480 }
1481
1482 async fn table_exists(&self, context: &SessionContext, table: &TableIdent) -> Result<bool> {
1484 if !self.supports_endpoint(&V1_TABLE_EXISTS).await? {
1488 return match self.load_table(context, table).await {
1489 Ok(_) => Ok(true),
1490 Err(e) if e.kind() == ErrorKind::TableNotFound => Ok(false),
1491 Err(e) => Err(e),
1492 };
1493 }
1494
1495 let client = self.client().await?;
1496 self.check_exists_via_head(context, client, client.config.table_endpoint(table))
1497 .await
1498 }
1499
1500 async fn rename_table(
1502 &self,
1503 context: &SessionContext,
1504 src: &TableIdent,
1505 dest: &TableIdent,
1506 ) -> Result<()> {
1507 let client = self.client().await?;
1508
1509 let request = HttpRequest::build(
1510 client
1511 .http_client
1512 .request(Method::POST, client.config.rename_table_endpoint())
1513 .json(&RenameTableRequest {
1514 source: src.clone(),
1515 destination: dest.clone(),
1516 }),
1517 )?;
1518
1519 client.invalidate_table(context, src);
1520 let http_response = client.query_catalog(context, request).await?;
1521
1522 match http_response.status() {
1523 StatusCode::NO_CONTENT | StatusCode::OK => Ok(()),
1524 StatusCode::NOT_FOUND => Err(Error::new(
1525 ErrorKind::TableNotFound,
1526 "Tried to rename a table that does not exist (is the namespace correct?)",
1527 )),
1528 StatusCode::CONFLICT => Err(Error::new(
1529 ErrorKind::TableAlreadyExists,
1530 "Tried to rename a table to a name that already exists",
1531 )),
1532 _ => Err(deserialize_unexpected_catalog_error(
1533 http_response,
1534 client.http_client.disable_header_redaction(),
1535 )),
1536 }
1537 }
1538
1539 async fn register_table(
1540 &self,
1541 context: &SessionContext,
1542 table_ident: &TableIdent,
1543 metadata_location: String,
1544 ) -> Result<Table> {
1545 let client = self.client().await?;
1546
1547 let request = HttpRequest::build(
1548 client
1549 .http_client
1550 .request(
1551 Method::POST,
1552 client
1553 .config
1554 .register_table_endpoint(table_ident.namespace()),
1555 )
1556 .json(&RegisterTableRequest {
1557 name: table_ident.name.clone(),
1558 metadata_location: metadata_location.clone(),
1559 overwrite: Some(false),
1560 }),
1561 )?;
1562
1563 client.invalidate_table(context, table_ident);
1564 let http_response = client.query_catalog(context, request).await?;
1565
1566 let response: LoadTableResult = match http_response.status() {
1567 StatusCode::OK => deserialize_catalog_response::<LoadTableResult>(http_response)?,
1568 StatusCode::NOT_FOUND => {
1569 return Err(Error::new(
1570 ErrorKind::NamespaceNotFound,
1571 "The namespace specified does not exist.",
1572 ));
1573 }
1574 StatusCode::CONFLICT => {
1575 return Err(Error::new(
1576 ErrorKind::TableAlreadyExists,
1577 "The given table already exists.",
1578 ));
1579 }
1580 _ => {
1581 return Err(deserialize_unexpected_catalog_error(
1582 http_response,
1583 client.http_client.disable_header_redaction(),
1584 ));
1585 }
1586 };
1587
1588 let metadata_location = response.metadata_location.as_ref().ok_or(Error::new(
1589 ErrorKind::DataInvalid,
1590 "Metadata location missing in `register_table` response!",
1591 ))?;
1592
1593 let file_io = self.load_file_io(Some(metadata_location), None).await?;
1594
1595 let mut table_builder = Table::builder()
1596 .identifier(table_ident.clone())
1597 .file_io(file_io)
1598 .metadata(response.metadata)
1599 .metadata_location(metadata_location.clone())
1600 .runtime(self.runtime.clone());
1601 if let Some(kms_client) = self.kms_client.clone() {
1602 table_builder = table_builder.kms_client(kms_client);
1603 }
1604 table_builder.build()
1605 }
1606
1607 async fn update_table(
1608 &self,
1609 context: &SessionContext,
1610 mut commit: TableCommit,
1611 ) -> Result<Table> {
1612 let client = self.client().await?;
1613
1614 let request = HttpRequest::build(
1615 client
1616 .http_client
1617 .request(
1618 Method::POST,
1619 client.config.table_endpoint(commit.identifier()),
1620 )
1621 .json(&CommitTableRequest {
1622 identifier: Some(commit.identifier().clone()),
1623 requirements: commit.take_requirements(),
1624 updates: commit.take_updates(),
1625 }),
1626 )?;
1627
1628 client.invalidate_table(context, commit.identifier());
1629 let http_response = client.query_catalog(context, request).await?;
1630
1631 let response: CommitTableResponse = match http_response.status() {
1632 StatusCode::OK => deserialize_catalog_response(http_response)?,
1633 StatusCode::NOT_FOUND => {
1634 return Err(Error::new(
1635 ErrorKind::TableNotFound,
1636 "Tried to update a table that does not exist",
1637 ));
1638 }
1639 StatusCode::CONFLICT => {
1640 return Err(Error::new(
1641 ErrorKind::CatalogCommitConflicts,
1642 "CatalogCommitConflicts, one or more requirements failed. The client may retry.",
1643 )
1644 .with_retryable(true));
1645 }
1646 StatusCode::INTERNAL_SERVER_ERROR => {
1647 return Err(Error::new(
1648 ErrorKind::Unexpected,
1649 "An unknown server-side problem occurred; the commit state is unknown.",
1650 ));
1651 }
1652 StatusCode::BAD_GATEWAY => {
1653 return Err(Error::new(
1654 ErrorKind::Unexpected,
1655 "A gateway or proxy received an invalid response from the upstream server; the commit state is unknown.",
1656 ));
1657 }
1658 StatusCode::GATEWAY_TIMEOUT => {
1659 return Err(Error::new(
1660 ErrorKind::Unexpected,
1661 "A server-side gateway timeout occurred; the commit state is unknown.",
1662 ));
1663 }
1664 _ => {
1665 return Err(deserialize_unexpected_catalog_error(
1666 http_response,
1667 client.http_client.disable_header_redaction(),
1668 ));
1669 }
1670 };
1671
1672 let file_io = self
1673 .load_file_io(Some(&response.metadata_location), None)
1674 .await?;
1675
1676 let mut table_builder = Table::builder()
1677 .identifier(commit.identifier().clone())
1678 .file_io(file_io)
1679 .metadata(response.metadata)
1680 .metadata_location(response.metadata_location)
1681 .runtime(self.runtime.clone());
1682 if let Some(kms_client) = self.kms_client.clone() {
1683 table_builder = table_builder.kms_client(kms_client);
1684 }
1685 table_builder.build()
1686 }
1687}
1688
1689#[derive(Debug)]
1694pub struct RestSessionCatalogBuilder {
1695 config: RestCatalogConfig,
1696 auth_manager: Option<Box<dyn AuthManager>>,
1697 storage_factory: Option<Arc<dyn StorageFactory>>,
1698 kms_client_factory: Option<Arc<dyn KmsClientFactory>>,
1699 runtime: Option<Runtime>,
1700}
1701
1702impl Default for RestSessionCatalogBuilder {
1703 fn default() -> Self {
1704 Self {
1705 config: RestCatalogConfig {
1706 name: None,
1707 uri: "".to_string(),
1708 warehouse: None,
1709 props: HashMap::new(),
1710 client: None,
1711 default_client: Arc::new(OnceLock::new()),
1712 },
1713 auth_manager: None,
1714 storage_factory: None,
1715 kms_client_factory: None,
1716 runtime: None,
1717 }
1718 }
1719}
1720
1721impl RestSessionCatalogBuilder {
1722 pub fn with_page_size(mut self, page_size: NonZeroU32) -> Self {
1730 self.config.props.insert(
1731 REST_CATALOG_PROP_PAGE_SIZE.to_string(),
1732 page_size.to_string(),
1733 );
1734 self
1735 }
1736
1737 pub fn with_client(mut self, client: Client) -> Self {
1739 self.config.client = Some(client);
1740 self
1741 }
1742
1743 pub fn with_auth_manager<M>(mut self, auth_manager: M) -> Self
1748 where M: AuthManager + 'static {
1749 self.auth_manager = Some(Box::new(auth_manager));
1750 self
1751 }
1752
1753 pub fn with_storage_factory(mut self, storage_factory: Arc<dyn StorageFactory>) -> Self {
1778 self.storage_factory = Some(storage_factory);
1779 self
1780 }
1781
1782 pub fn with_kms_client_factory(
1803 mut self,
1804 kms_client_factory: Arc<dyn KmsClientFactory>,
1805 ) -> Self {
1806 self.kms_client_factory = Some(kms_client_factory);
1807 self
1808 }
1809
1810 pub fn with_runtime(mut self, runtime: Runtime) -> Self {
1816 self.runtime = Some(runtime);
1817 self
1818 }
1819
1820 pub fn load(
1825 mut self,
1826 name: impl Into<String>,
1827 props: HashMap<String, String>,
1828 ) -> impl Future<Output = Result<RestSessionCatalog>> + Send {
1829 self.config.name = Some(name.into());
1830
1831 if props.contains_key(REST_CATALOG_PROP_URI) {
1832 self.config.uri = props
1833 .get(REST_CATALOG_PROP_URI)
1834 .cloned()
1835 .unwrap_or_default();
1836 }
1837
1838 if props.contains_key(REST_CATALOG_PROP_WAREHOUSE) {
1839 self.config.warehouse = props.get(REST_CATALOG_PROP_WAREHOUSE).cloned()
1840 }
1841
1842 self.config.props.extend(props);
1844 self.config.props.remove(REST_CATALOG_PROP_URI);
1845 self.config.props.remove(REST_CATALOG_PROP_WAREHOUSE);
1846
1847 async move {
1848 if self.config.name.is_none() {
1849 Err(Error::new(
1850 ErrorKind::DataInvalid,
1851 "Catalog name is required",
1852 ))
1853 } else if self.config.uri.is_empty() {
1854 Err(Error::new(
1855 ErrorKind::DataInvalid,
1856 "Catalog uri is required",
1857 ))
1858 } else {
1859 let runtime = self.runtime.unwrap_or_else(Runtime::current);
1860 let kms_client = match self.kms_client_factory {
1861 Some(factory) => Some(factory.create_kms_client(&self.config.props).await?),
1862 None => None,
1863 };
1864
1865 Ok(RestSessionCatalog::new(
1866 self.config,
1867 self.auth_manager,
1868 self.storage_factory,
1869 runtime,
1870 kms_client,
1871 ))
1872 }
1873 }
1874 }
1875}
1876
1877#[cfg(test)]
1878mod tests {
1879 use std::fs::File;
1880 use std::io::BufReader;
1881 use std::sync::{Arc, Mutex};
1882
1883 use chrono::{TimeZone, Utc};
1884 use iceberg::io::LocalFsStorageFactory;
1885 use iceberg::spec::{
1886 FormatVersion, NestedField, NullOrder, Operation, PrimitiveType, Schema, Snapshot,
1887 SnapshotLog, SortDirection, SortField, SortOrder, Summary, Transform, Type,
1888 UnboundPartitionField, UnboundPartitionSpec,
1889 };
1890 use iceberg::transaction::{ApplyTransactionAction, Transaction};
1891 use iceberg_test_utils::test_runtime;
1892 use mockito::{Matcher, Mock, Server, ServerGuard};
1893 use serde_json::json;
1894 use uuid::uuid;
1895
1896 use super::*;
1897 use crate::auth::AuthSession;
1898 use crate::request::HttpRequest;
1899
1900 #[derive(Debug)]
1901 struct PlainContextSession;
1902
1903 #[async_trait]
1904 impl AuthSession for PlainContextSession {
1905 async fn authenticate(&self, _request: &mut HttpRequest) -> Result<()> {
1906 Ok(())
1907 }
1908 }
1909
1910 #[derive(Debug)]
1911 struct ContextSession(String);
1912
1913 #[async_trait]
1914 impl AuthSession for ContextSession {
1915 async fn authenticate(&self, request: &mut HttpRequest) -> Result<()> {
1916 request.headers_mut().insert(
1917 "x-session-id",
1918 HeaderValue::from_str(&self.0).expect("valid test session ID"),
1919 );
1920 Ok(())
1921 }
1922 }
1923
1924 #[derive(Debug)]
1925 struct ContextManager {
1926 catalog_session: Arc<dyn AuthSession>,
1927 seen_session_ids: Arc<Mutex<Vec<String>>>,
1928 fail_contextual_session: bool,
1929 }
1930
1931 #[async_trait]
1932 impl AuthManager for ContextManager {
1933 async fn init_session(
1934 &self,
1935 _client: &HttpClient,
1936 _props: &HashMap<String, String>,
1937 ) -> Result<Box<dyn AuthSession>> {
1938 Ok(Box::new(PlainContextSession))
1939 }
1940
1941 async fn catalog_session(
1942 &self,
1943 _client: &HttpClient,
1944 _props: &HashMap<String, String>,
1945 ) -> Result<Arc<dyn AuthSession>> {
1946 Ok(self.catalog_session.clone())
1947 }
1948
1949 async fn contextual_session(
1950 &self,
1951 context: &SessionContext,
1952 catalog_session: Arc<dyn AuthSession>,
1953 ) -> Result<Arc<dyn AuthSession>> {
1954 assert!(Arc::ptr_eq(&catalog_session, &self.catalog_session));
1955 self.seen_session_ids
1956 .lock()
1957 .unwrap()
1958 .push(context.session_id().to_string());
1959 if self.fail_contextual_session {
1960 return Err(Error::new(
1961 ErrorKind::Unexpected,
1962 "contextual session failure",
1963 ));
1964 }
1965 Ok(Arc::new(ContextSession(context.session_id().to_string())))
1966 }
1967 }
1968
1969 fn test_catalog(config: RestCatalogConfig) -> RestSessionCatalog {
1970 RestSessionCatalog::new(config, None, None, Runtime::current(), None)
1971 }
1972
1973 fn test_catalog_with<M>(config: RestCatalogConfig, auth_manager: M) -> RestSessionCatalog
1974 where M: AuthManager + 'static {
1975 RestSessionCatalog::new(
1976 config,
1977 Some(Box::new(auth_manager)),
1978 None,
1979 Runtime::current(),
1980 None,
1981 )
1982 }
1983
1984 fn context_catalog(
1985 config: RestCatalogConfig,
1986 fail_contextual_session: bool,
1987 ) -> (RestSessionCatalog, Arc<Mutex<Vec<String>>>) {
1988 let seen_session_ids = Arc::new(Mutex::new(Vec::new()));
1989 let catalog_session: Arc<dyn AuthSession> = Arc::new(PlainContextSession);
1990 let catalog = test_catalog_with(config, ContextManager {
1991 catalog_session,
1992 seen_session_ids: seen_session_ids.clone(),
1993 fail_contextual_session,
1994 });
1995 (catalog, seen_session_ids)
1996 }
1997
1998 fn test_client() -> HttpClient {
1999 HttpClient::new(
2000 &RestCatalogConfig::builder()
2001 .uri("http://localhost".to_string())
2002 .build(),
2003 )
2004 .unwrap()
2005 }
2006
2007 fn session_catalog(config: RestCatalogConfig) -> RestSessionCatalog {
2009 RestSessionCatalog::new(
2010 config,
2011 None,
2012 Some(Arc::new(LocalFsStorageFactory)),
2013 Runtime::current(),
2014 None,
2015 )
2016 }
2017
2018 #[tokio::test]
2019 async fn test_page_size_precedence() {
2020 for (default, builder, property, override_value, expected) in [
2022 (Some("10"), None, None, None, 10),
2023 (Some("10"), Some(20), None, None, 20),
2024 (None, Some(20), None, None, 20),
2025 (Some("10"), None, Some("30"), None, 30),
2026 (Some("10"), Some(20), Some("30"), None, 30),
2027 (Some("10"), Some(20), Some("30"), Some("40"), 40),
2028 (None, None, None, Some("40"), 40),
2029 (None, Some(1), None, None, 1),
2030 (None, Some(i32::MAX as u32), None, None, i32::MAX as u32),
2031 (Some("invalid"), Some(20), None, None, 20),
2033 (None, Some(20), Some("0"), Some("40"), 40),
2034 (None, Some(u32::MAX), None, Some("40"), 40),
2035 ] {
2036 let mut server = Server::new_async().await;
2037 let page_size_props = |value: Option<&str>| -> HashMap<String, String> {
2038 value
2039 .map(|value| (REST_CATALOG_PROP_PAGE_SIZE.to_string(), value.to_string()))
2040 .into_iter()
2041 .collect()
2042 };
2043 let config_mock = server
2044 .mock("GET", "/v1/config")
2045 .with_status(200)
2046 .with_body(
2047 json!({
2048 "defaults": page_size_props(default),
2049 "overrides": page_size_props(override_value),
2050 })
2051 .to_string(),
2052 )
2053 .create_async()
2054 .await;
2055 let list_mock = server
2056 .mock("GET", "/v1/namespaces")
2057 .match_query(Matcher::UrlEncoded("pageSize".into(), expected.to_string()))
2058 .with_status(200)
2059 .with_body(r#"{"namespaces": []}"#)
2060 .expect(1)
2061 .create_async()
2062 .await;
2063
2064 let mut builder_config = RestCatalogBuilder::default();
2065 if let Some(page_size) = builder {
2066 builder_config = builder_config.with_page_size(NonZeroU32::new(page_size).unwrap());
2067 }
2068 let mut props = page_size_props(property);
2069 props.insert(REST_CATALOG_PROP_URI.to_string(), server.url());
2070 let catalog = builder_config.load("test", props).await.unwrap();
2071
2072 catalog.list_namespaces(None).await.unwrap_or_else(|err| {
2073 panic!(
2074 "default={default:?}, builder={builder:?}, property={property:?}, \
2075 override={override_value:?}, expected={expected}: {err}"
2076 )
2077 });
2078 config_mock.assert_async().await;
2079 list_mock.assert_async().await;
2080 }
2081 }
2082
2083 #[tokio::test]
2084 async fn test_page_size_pagination() {
2085 let mut server = Server::new_async().await;
2086 let config_mock = create_config_mock(&mut server).await;
2087 let catalog = RestCatalogBuilder::default()
2088 .with_page_size(NonZeroU32::MIN)
2089 .load(
2090 "test",
2091 HashMap::from([(REST_CATALOG_PROP_URI.to_string(), server.url())]),
2092 )
2093 .await
2094 .unwrap();
2095
2096 let mut mocks = Vec::new();
2097 for (endpoint, parent, first_body, empty_body, last_body) in [
2098 (
2099 "/v1/namespaces",
2100 Some("parent\u{1f}child"),
2101 json!({"namespaces": [["ns1"]], "next-page-token": "next/+"}),
2102 json!({"namespaces": [], "next-page-token": "last"}),
2103 json!({"namespaces": [["ns2"]], "next-page-token": null}),
2104 ),
2105 (
2106 "/v1/namespaces/ns1/tables",
2107 None,
2108 json!({"identifiers": [{"namespace": ["ns1"], "name": "t1"}], "next-page-token": "next/+"}),
2109 json!({"identifiers": [], "next-page-token": "last"}),
2110 json!({"identifiers": [{"namespace": ["ns1"], "name": "t2"}]}),
2111 ),
2112 ] {
2113 for (token, body) in [
2117 ("", first_body),
2118 ("next/+", empty_body),
2119 ("last", last_body),
2120 ] {
2121 let mut query = vec![
2122 Matcher::UrlEncoded("pageSize".into(), "1".into()),
2123 Matcher::UrlEncoded("pageToken".into(), token.into()),
2124 ];
2125 if let Some(parent) = parent {
2126 query.push(Matcher::UrlEncoded("parent".into(), parent.into()));
2127 }
2128 mocks.push(
2129 server
2130 .mock("GET", endpoint)
2131 .match_query(Matcher::AllOf(query))
2132 .with_status(200)
2133 .with_body(body.to_string())
2134 .expect(1)
2135 .create_async()
2136 .await,
2137 );
2138 }
2139 }
2140
2141 let parent = NamespaceIdent::from_vec(vec!["parent".into(), "child".into()]).unwrap();
2142 assert_eq!(catalog.list_namespaces(Some(&parent)).await.unwrap(), vec![
2143 NamespaceIdent::new("ns1".into()),
2144 NamespaceIdent::new("ns2".into())
2145 ]);
2146 let namespace = NamespaceIdent::new("ns1".into());
2147 assert_eq!(catalog.list_tables(&namespace).await.unwrap(), vec![
2148 TableIdent::new(namespace.clone(), "t1".into()),
2149 TableIdent::new(namespace, "t2".into())
2150 ]);
2151 config_mock.assert_async().await;
2152 for mock in mocks {
2153 mock.assert_async().await;
2154 }
2155 }
2156
2157 #[tokio::test]
2158 async fn test_page_size_is_not_a_result_limit() {
2159 let mut server = Server::new_async().await;
2160 let config_mock = create_config_mock(&mut server).await;
2161 let catalog = RestCatalogBuilder::default()
2162 .with_page_size(NonZeroU32::MIN)
2163 .load(
2164 "test",
2165 HashMap::from([(REST_CATALOG_PROP_URI.to_string(), server.url())]),
2166 )
2167 .await
2168 .unwrap();
2169
2170 let mut mocks = Vec::new();
2172 for (endpoint, body) in [
2173 ("/v1/namespaces", json!({"namespaces": [["ns1"], ["ns2"]]})),
2174 (
2175 "/v1/namespaces/ns1/tables",
2176 json!({"identifiers": [
2177 {"namespace": ["ns1"], "name": "t1"},
2178 {"namespace": ["ns1"], "name": "t2"},
2179 ]}),
2180 ),
2181 ] {
2182 mocks.push(
2183 server
2184 .mock("GET", endpoint)
2185 .match_query(Matcher::AllOf(vec![
2186 Matcher::UrlEncoded("pageSize".into(), "1".into()),
2187 Matcher::UrlEncoded("pageToken".into(), "".into()),
2188 ]))
2189 .with_status(200)
2190 .with_body(body.to_string())
2191 .expect(1)
2192 .create_async()
2193 .await,
2194 );
2195 }
2196
2197 assert_eq!(catalog.list_namespaces(None).await.unwrap(), vec![
2198 NamespaceIdent::new("ns1".into()),
2199 NamespaceIdent::new("ns2".into())
2200 ]);
2201 let namespace = NamespaceIdent::new("ns1".into());
2202 assert_eq!(catalog.list_tables(&namespace).await.unwrap(), vec![
2203 TableIdent::new(namespace.clone(), "t1".into()),
2204 TableIdent::new(namespace, "t2".into())
2205 ]);
2206 config_mock.assert_async().await;
2207 for mock in mocks {
2208 mock.assert_async().await;
2209 }
2210 }
2211
2212 #[tokio::test]
2213 async fn test_page_size_absent() {
2214 let mut server = Server::new_async().await;
2215 let config_mock = create_config_mock(&mut server).await;
2216 let namespaces_mock = server
2217 .mock("GET", "/v1/namespaces")
2218 .match_query(Matcher::Missing)
2219 .with_status(200)
2220 .with_body(r#"{"namespaces": []}"#)
2221 .create_async()
2222 .await;
2223 let tables_mock = server
2224 .mock("GET", "/v1/namespaces/ns1/tables")
2225 .match_query(Matcher::Missing)
2226 .with_status(200)
2227 .with_body(r#"{"identifiers": []}"#)
2228 .create_async()
2229 .await;
2230 let catalog = RestCatalogBuilder::default()
2231 .load(
2232 "test",
2233 HashMap::from([(REST_CATALOG_PROP_URI.to_string(), server.url())]),
2234 )
2235 .await
2236 .unwrap();
2237
2238 assert!(catalog.list_namespaces(None).await.unwrap().is_empty());
2239 assert!(
2240 catalog
2241 .list_tables(&NamespaceIdent::new("ns1".into()))
2242 .await
2243 .unwrap()
2244 .is_empty()
2245 );
2246 config_mock.assert_async().await;
2247 namespaces_mock.assert_async().await;
2248 tables_mock.assert_async().await;
2249 }
2250
2251 #[tokio::test]
2252 async fn test_invalid_page_size() {
2253 for value in [
2254 "",
2255 "0",
2256 "-1",
2257 "abc",
2258 "1.5",
2259 "2147483648",
2260 "4294967295",
2261 "4294967296",
2262 ] {
2263 for source in ["defaults", "client", "overrides"] {
2264 let mut server = Server::new_async().await;
2265 let mut config = json!({"defaults": {}, "overrides": {}});
2266 let mut props = HashMap::from([(REST_CATALOG_PROP_URI.to_string(), server.url())]);
2267 if source == "client" {
2268 config["defaults"][REST_CATALOG_PROP_PAGE_SIZE] = json!("1000");
2270 props.insert(REST_CATALOG_PROP_PAGE_SIZE.to_string(), value.to_string());
2271 } else {
2272 config[source][REST_CATALOG_PROP_PAGE_SIZE] = json!(value);
2273 }
2274 let config_mock = server
2275 .mock("GET", "/v1/config")
2276 .with_status(200)
2277 .with_body(config.to_string())
2278 .create_async()
2279 .await;
2280 let catalog = RestSessionCatalogBuilder::default()
2281 .load("test", props)
2282 .await
2283 .unwrap();
2284 let err = catalog
2285 .list_namespaces(&SessionContext::empty(), None)
2286 .await
2287 .unwrap_err();
2288 assert_eq!(err.kind(), ErrorKind::DataInvalid);
2289 assert_eq!(
2290 err.message(),
2291 format!("Invalid {REST_CATALOG_PROP_PAGE_SIZE}: {value}")
2292 );
2293 config_mock.assert_async().await;
2294 }
2295 }
2296 }
2297
2298 #[tokio::test]
2299 async fn test_update_config() {
2300 let mut server = Server::new_async().await;
2301
2302 let config_mock = server
2303 .mock("GET", "/v1/config")
2304 .with_status(200)
2305 .with_body(
2306 r#"{
2307 "overrides": {
2308 "warehouse": "s3://iceberg-catalog"
2309 },
2310 "defaults": {}
2311 }"#,
2312 )
2313 .create_async()
2314 .await;
2315
2316 let catalog = session_catalog(RestCatalogConfig::builder().uri(server.url()).build());
2317
2318 assert_eq!(
2319 catalog
2320 .client()
2321 .await
2322 .unwrap()
2323 .config
2324 .props
2325 .get("warehouse"),
2326 Some(&"s3://iceberg-catalog".to_string())
2327 );
2328
2329 config_mock.assert_async().await;
2330 }
2331
2332 async fn create_config_mock(server: &mut ServerGuard) -> Mock {
2333 server
2334 .mock("GET", "/v1/config")
2335 .with_status(200)
2336 .with_body(
2337 r#"{
2338 "overrides": {
2339 "warehouse": "s3://iceberg-catalog"
2340 },
2341 "defaults": {}
2342 }"#,
2343 )
2344 .create_async()
2345 .await
2346 }
2347
2348 async fn create_config_mock_with_exists_endpoints(server: &mut ServerGuard) -> Mock {
2351 server
2352 .mock("GET", "/v1/config")
2353 .with_status(200)
2354 .with_body(
2355 r#"{
2356 "overrides": { "warehouse": "s3://iceberg-catalog" },
2357 "defaults": {},
2358 "endpoints": [
2359 "HEAD /v1/{prefix}/namespaces/{namespace}",
2360 "HEAD /v1/{prefix}/namespaces/{namespace}/tables/{table}"
2361 ]
2362 }"#,
2363 )
2364 .create_async()
2365 .await
2366 }
2367
2368 #[tokio::test]
2369 async fn test_config_advertised_endpoints() {
2370 let mut server = Server::new_async().await;
2371
2372 let config_mock = server
2373 .mock("GET", "/v1/config")
2374 .with_status(200)
2375 .with_body(
2376 r#"{
2377 "overrides": {},
2378 "defaults": {},
2379 "endpoints": [
2380 "GET /v1/{prefix}/namespaces",
2381 "POST /v1/{prefix}/namespaces/{namespace}/tables/{table}/plan"
2382 ]
2383 }"#,
2384 )
2385 .create_async()
2386 .await;
2387
2388 let catalog = RestCatalog::new(
2389 SessionContext::empty(),
2390 RestCatalogConfig::builder().uri(server.url()).build(),
2391 None,
2392 Some(Arc::new(LocalFsStorageFactory)),
2393 Runtime::current(),
2394 None,
2395 );
2396
2397 let plan = "POST /v1/{prefix}/namespaces/{namespace}/tables/{table}/plan"
2398 .parse::<Endpoint>()
2399 .unwrap();
2400 assert!(catalog.inner.supports_endpoint(&plan).await.unwrap());
2401 let delete_ns = "DELETE /v1/{prefix}/namespaces/{namespace}"
2403 .parse::<Endpoint>()
2404 .unwrap();
2405 assert!(!catalog.inner.supports_endpoint(&delete_ns).await.unwrap());
2406
2407 config_mock.assert_async().await;
2408 }
2409
2410 #[tokio::test]
2411 async fn test_config_without_endpoints_falls_back_to_default_set() {
2412 let mut server = Server::new_async().await;
2413
2414 let config_mock = server
2415 .mock("GET", "/v1/config")
2416 .with_status(200)
2417 .with_body(r#"{ "overrides": {}, "defaults": {} }"#)
2418 .create_async()
2419 .await;
2420
2421 let catalog = RestCatalog::new(
2422 SessionContext::empty(),
2423 RestCatalogConfig::builder().uri(server.url()).build(),
2424 None,
2425 Some(Arc::new(LocalFsStorageFactory)),
2426 Runtime::current(),
2427 None,
2428 );
2429
2430 let load_table = "GET /v1/{prefix}/namespaces/{namespace}/tables/{table}"
2433 .parse::<Endpoint>()
2434 .unwrap();
2435 assert!(catalog.inner.supports_endpoint(&load_table).await.unwrap());
2436 let plan = "POST /v1/{prefix}/namespaces/{namespace}/tables/{table}/plan"
2438 .parse::<Endpoint>()
2439 .unwrap();
2440 assert!(!catalog.inner.supports_endpoint(&plan).await.unwrap());
2441
2442 config_mock.assert_async().await;
2443 }
2444
2445 #[tokio::test]
2446 async fn test_config_with_empty_endpoints_falls_back_to_default_set() {
2447 let mut server = Server::new_async().await;
2448
2449 let config_mock = server
2452 .mock("GET", "/v1/config")
2453 .with_status(200)
2454 .with_body(r#"{ "overrides": {}, "defaults": {}, "endpoints": [] }"#)
2455 .create_async()
2456 .await;
2457
2458 let catalog = RestCatalog::new(
2459 SessionContext::empty(),
2460 RestCatalogConfig::builder().uri(server.url()).build(),
2461 None,
2462 Some(Arc::new(LocalFsStorageFactory)),
2463 Runtime::current(),
2464 None,
2465 );
2466
2467 let load_table = "GET /v1/{prefix}/namespaces/{namespace}/tables/{table}"
2468 .parse::<Endpoint>()
2469 .unwrap();
2470 assert!(catalog.inner.supports_endpoint(&load_table).await.unwrap());
2471
2472 config_mock.assert_async().await;
2473 }
2474
2475 async fn create_oauth_mock(server: &mut ServerGuard) -> Mock {
2476 create_oauth_mock_with_path(server, "/v1/oauth/tokens", "ey000000000000", 200).await
2477 }
2478
2479 async fn create_oauth_mock_with_path(
2480 server: &mut ServerGuard,
2481 path: &str,
2482 token: &str,
2483 status: usize,
2484 ) -> Mock {
2485 let body = format!(
2486 r#"{{
2487 "access_token": "{token}",
2488 "token_type": "Bearer",
2489 "issued_token_type": "urn:ietf:params:oauth:token-type:access_token",
2490 "expires_in": 86400
2491 }}"#
2492 );
2493 server
2494 .mock("POST", path)
2495 .with_status(status)
2496 .with_body(body)
2497 .expect(1)
2498 .create_async()
2499 .await
2500 }
2501
2502 #[tokio::test]
2503 async fn test_oauth() {
2504 let mut server = Server::new_async().await;
2505 let oauth_mock = create_oauth_mock(&mut server).await;
2506 let config_mock = create_config_mock(&mut server).await;
2507
2508 let mut props = HashMap::new();
2509 props.insert("credential".to_string(), "client1:secret1".to_string());
2510
2511 let catalog = session_catalog(
2512 RestCatalogConfig::builder()
2513 .uri(server.url())
2514 .props(props)
2515 .build(),
2516 );
2517
2518 let token = catalog.client().await.unwrap().token().await;
2519 oauth_mock.assert_async().await;
2520 config_mock.assert_async().await;
2521 assert_eq!(token, Some("ey000000000000".to_string()));
2522 }
2523
2524 #[tokio::test]
2525 async fn test_oauth_with_optional_param() {
2526 let mut props = HashMap::new();
2527 props.insert("credential".to_string(), "client1:secret1".to_string());
2528 props.insert("scope".to_string(), "custom_scope".to_string());
2529 props.insert("audience".to_string(), "custom_audience".to_string());
2530 props.insert("resource".to_string(), "custom_resource".to_string());
2531
2532 let mut server = Server::new_async().await;
2533 let oauth_mock = server
2534 .mock("POST", "/v1/oauth/tokens")
2535 .match_body(Matcher::Regex("scope=custom_scope".to_string()))
2536 .match_body(Matcher::Regex("audience=custom_audience".to_string()))
2537 .match_body(Matcher::Regex("resource=custom_resource".to_string()))
2538 .with_status(200)
2539 .with_body(
2540 r#"{
2541 "access_token": "ey000000000000",
2542 "token_type": "Bearer",
2543 "issued_token_type": "urn:ietf:params:oauth:token-type:access_token",
2544 "expires_in": 86400
2545 }"#,
2546 )
2547 .expect(1)
2548 .create_async()
2549 .await;
2550
2551 let config_mock = create_config_mock(&mut server).await;
2552
2553 let catalog = session_catalog(
2554 RestCatalogConfig::builder()
2555 .uri(server.url())
2556 .props(props)
2557 .build(),
2558 );
2559
2560 let token = catalog.client().await.unwrap().token().await;
2561
2562 oauth_mock.assert_async().await;
2563 config_mock.assert_async().await;
2564 assert_eq!(token, Some("ey000000000000".to_string()));
2565 }
2566
2567 #[tokio::test]
2568 async fn test_http_headers() {
2569 let server = Server::new_async().await;
2570 let mut props = HashMap::new();
2571 props.insert("credential".to_string(), "client1:secret1".to_string());
2572
2573 let config = RestCatalogConfig::builder()
2574 .uri(server.url())
2575 .props(props)
2576 .build();
2577 let headers: HeaderMap = config.extra_headers().unwrap();
2578
2579 let expected_headers = HeaderMap::from_iter([
2580 (
2581 header::CONTENT_TYPE,
2582 HeaderValue::from_static("application/json"),
2583 ),
2584 (
2585 HeaderName::from_static("x-client-version"),
2586 HeaderValue::from_static(ICEBERG_REST_SPEC_VERSION),
2587 ),
2588 (
2589 header::USER_AGENT,
2590 HeaderValue::from_str(&format!("iceberg-rs/{CARGO_PKG_VERSION}")).unwrap(),
2591 ),
2592 ]);
2593 assert_eq!(headers, expected_headers);
2594 }
2595
2596 #[tokio::test]
2597 async fn test_http_headers_with_custom_headers() {
2598 let server = Server::new_async().await;
2599 let mut props = HashMap::new();
2600 props.insert("credential".to_string(), "client1:secret1".to_string());
2601 props.insert(
2602 "header.content-type".to_string(),
2603 "application/yaml".to_string(),
2604 );
2605 props.insert(
2606 "header.customized-header".to_string(),
2607 "some/value".to_string(),
2608 );
2609
2610 let config = RestCatalogConfig::builder()
2611 .uri(server.url())
2612 .props(props)
2613 .build();
2614 let headers: HeaderMap = config.extra_headers().unwrap();
2615
2616 let expected_headers = HeaderMap::from_iter([
2617 (
2618 header::CONTENT_TYPE,
2619 HeaderValue::from_static("application/yaml"),
2620 ),
2621 (
2622 HeaderName::from_static("x-client-version"),
2623 HeaderValue::from_static(ICEBERG_REST_SPEC_VERSION),
2624 ),
2625 (
2626 header::USER_AGENT,
2627 HeaderValue::from_str(&format!("iceberg-rs/{CARGO_PKG_VERSION}")).unwrap(),
2628 ),
2629 (
2630 HeaderName::from_static("customized-header"),
2631 HeaderValue::from_static("some/value"),
2632 ),
2633 ]);
2634 assert_eq!(headers, expected_headers);
2635 }
2636
2637 #[tokio::test]
2638 async fn test_oauth_with_oauth2_server_uri() {
2639 let mut server = Server::new_async().await;
2640 let config_mock = create_config_mock(&mut server).await;
2641
2642 let mut auth_server = Server::new_async().await;
2643 let auth_server_path = "/some/path";
2644 let oauth_mock =
2645 create_oauth_mock_with_path(&mut auth_server, auth_server_path, "ey000000000000", 200)
2646 .await;
2647
2648 let mut props = HashMap::new();
2649 props.insert("credential".to_string(), "client1:secret1".to_string());
2650 props.insert(
2651 "oauth2-server-uri".to_string(),
2652 format!("{}{}", auth_server.url(), auth_server_path).to_string(),
2653 );
2654
2655 let catalog = session_catalog(
2656 RestCatalogConfig::builder()
2657 .uri(server.url())
2658 .props(props)
2659 .build(),
2660 );
2661
2662 let token = catalog.client().await.unwrap().token().await;
2663
2664 oauth_mock.assert_async().await;
2665 config_mock.assert_async().await;
2666 assert_eq!(token, Some("ey000000000000".to_string()));
2667 }
2668
2669 #[tokio::test]
2670 async fn test_config_override() {
2671 let mut server = Server::new_async().await;
2672 let mut redirect_server = Server::new_async().await;
2673 let new_uri = redirect_server.url();
2674
2675 let config_mock = server
2676 .mock("GET", "/v1/config")
2677 .with_status(200)
2678 .with_body(
2679 json!(
2680 {
2681 "overrides": {
2682 "uri": new_uri,
2683 "warehouse": "s3://iceberg-catalog",
2684 "prefix": "ice/warehouses/my"
2685 },
2686 "defaults": {},
2687 }
2688 )
2689 .to_string(),
2690 )
2691 .create_async()
2692 .await;
2693
2694 let list_ns_mock = redirect_server
2695 .mock("GET", "/v1/ice/warehouses/my/namespaces")
2696 .with_body(
2697 r#"{
2698 "namespaces": []
2699 }"#,
2700 )
2701 .create_async()
2702 .await;
2703
2704 let catalog = session_catalog(RestCatalogConfig::builder().uri(server.url()).build());
2705
2706 let _namespaces = catalog
2707 .list_namespaces(&SessionContext::empty(), None)
2708 .await
2709 .unwrap();
2710
2711 config_mock.assert_async().await;
2712 list_ns_mock.assert_async().await;
2713 }
2714
2715 #[tokio::test]
2716 async fn test_list_namespace() {
2717 let mut server = Server::new_async().await;
2718
2719 let config_mock = create_config_mock(&mut server).await;
2720
2721 let list_ns_mock = server
2722 .mock("GET", "/v1/namespaces")
2723 .with_body(
2724 r#"{
2725 "namespaces": [
2726 ["ns1", "ns11"],
2727 ["ns2"]
2728 ]
2729 }"#,
2730 )
2731 .create_async()
2732 .await;
2733
2734 let catalog = session_catalog(RestCatalogConfig::builder().uri(server.url()).build());
2735
2736 let namespaces = catalog
2737 .list_namespaces(&SessionContext::empty(), None)
2738 .await
2739 .unwrap();
2740
2741 let expected_ns = vec![
2742 NamespaceIdent::from_vec(vec!["ns1".to_string(), "ns11".to_string()]).unwrap(),
2743 NamespaceIdent::from_vec(vec!["ns2".to_string()]).unwrap(),
2744 ];
2745
2746 assert_eq!(expected_ns, namespaces);
2747
2748 config_mock.assert_async().await;
2749 list_ns_mock.assert_async().await;
2750 }
2751
2752 #[tokio::test]
2753 async fn test_auth_type_none_disables_auth() {
2754 let props = HashMap::from([
2756 (REST_CATALOG_PROP_AUTH_TYPE.to_string(), "none".to_string()),
2757 ("token".to_string(), "some-oauth-token".to_string()),
2758 ]);
2759 let config = RestCatalogConfig::builder()
2760 .uri("http://localhost".to_string())
2761 .props(props)
2762 .build();
2763
2764 let session = test_catalog(config)
2765 .resolve_auth_manager()
2766 .unwrap()
2767 .init_session(&test_client(), &HashMap::new())
2768 .await
2769 .unwrap();
2770 let mut req = HttpRequest::new(
2771 Client::new()
2772 .get("https://rest.example.com/v1/config")
2773 .build()
2774 .unwrap(),
2775 );
2776 session.authenticate(&mut req).await.unwrap();
2777 assert!(req.headers().get("authorization").is_none());
2778 }
2779
2780 #[tokio::test]
2781 async fn test_header_prop_overrides_token_on_the_wire() {
2782 let mut server = Server::new_async().await;
2786 let config_mock = create_config_mock(&mut server).await;
2787 let list_ns_mock = server
2788 .mock("GET", "/v1/namespaces")
2789 .match_header("authorization", "Basic xyz")
2790 .with_body(r#"{"namespaces": []}"#)
2791 .create_async()
2792 .await;
2793
2794 let props = HashMap::from([
2795 ("token".to_string(), "some-oauth-token".to_string()),
2796 ("header.authorization".to_string(), "Basic xyz".to_string()),
2797 ]);
2798 let catalog = RestCatalog::new(
2799 SessionContext::empty(),
2800 RestCatalogConfig::builder()
2801 .uri(server.url())
2802 .props(props)
2803 .build(),
2804 None,
2805 Some(Arc::new(LocalFsStorageFactory)),
2806 Runtime::current(),
2807 None,
2808 );
2809
2810 catalog.list_namespaces(None).await.unwrap();
2811 config_mock.assert_async().await;
2812 list_ns_mock.assert_async().await;
2813 }
2814
2815 #[tokio::test]
2816 async fn test_builtin_oauth_endpoint_follows_uri_override() {
2817 let mut bootstrap = Server::new_async().await;
2821 let overridden = Server::new_async().await;
2822
2823 let config_mock = bootstrap
2824 .mock("GET", "/v1/config")
2825 .with_status(200)
2826 .with_body(format!(
2827 r#"{{"overrides": {{"uri": "{}"}}, "defaults": {{}}}}"#,
2828 overridden.url()
2829 ))
2830 .create_async()
2831 .await;
2832 let bootstrap_oauth_mock =
2834 create_oauth_mock_with_path(&mut bootstrap, "/v1/oauth/tokens", "tok-boot", 200).await;
2835
2836 let props = HashMap::from([("credential".to_string(), "client1:secret1".to_string())]);
2837 let catalog = RestCatalog::new(
2838 SessionContext::empty(),
2839 RestCatalogConfig::builder()
2840 .uri(bootstrap.url())
2841 .props(props)
2842 .build(),
2843 None,
2844 Some(Arc::new(LocalFsStorageFactory)),
2845 Runtime::current(),
2846 None,
2847 );
2848
2849 let client = catalog.client().await.unwrap();
2850 config_mock.assert_async().await;
2851 bootstrap_oauth_mock.assert_async().await;
2852 let session_debug = format!("{:?}", client.catalog_session);
2855 assert!(session_debug.contains(&format!("{}/v1/oauth/tokens", overridden.url())));
2856 }
2857
2858 #[tokio::test]
2859 async fn test_concurrent_authenticate_single_token_exchange() {
2860 let mut server = Server::new_async().await;
2863 let oauth_mock =
2865 create_oauth_mock_with_path(&mut server, "/v1/oauth/tokens", "tok-once", 200).await;
2866
2867 let manager = OAuth2Manager::new(format!("{}/v1/oauth/tokens", server.url()))
2868 .with_credential(Some("client1".to_string()), "secret1".to_string());
2869 let session: Arc<dyn AuthSession> = Arc::from(
2870 manager
2871 .init_session(&test_client(), &HashMap::new())
2872 .await
2873 .unwrap(),
2874 );
2875
2876 let client = Client::new();
2877 let attempts = (0..8).map(|_| {
2878 let session = session.clone();
2879 let client = client.clone();
2880 async move {
2881 let mut req = HttpRequest::new(
2882 client
2883 .get("https://rest.example.com/v1/config")
2884 .build()
2885 .unwrap(),
2886 );
2887 session.authenticate(&mut req).await.unwrap();
2888 req.headers()
2889 .get("authorization")
2890 .unwrap()
2891 .to_str()
2892 .unwrap()
2893 .to_string()
2894 }
2895 });
2896 let bearers = futures::future::join_all(attempts).await;
2897
2898 oauth_mock.assert_async().await;
2899 assert!(bearers.iter().all(|b| b == "Bearer tok-once"));
2900 }
2901
2902 #[tokio::test]
2903 async fn test_seeded_token_takes_precedence_over_credential() {
2904 let mut server = Server::new_async().await;
2907 let oauth_mock = server
2908 .mock("POST", "/v1/oauth/tokens")
2909 .expect(0)
2910 .create_async()
2911 .await;
2912
2913 let manager = OAuth2Manager::new(format!("{}/v1/oauth/tokens", server.url()))
2914 .with_token("tok-seed")
2915 .with_credential(Some("client1".to_string()), "secret1".to_string());
2916 let session = manager
2917 .init_session(&test_client(), &HashMap::new())
2918 .await
2919 .unwrap();
2920
2921 let mut req = HttpRequest::new(
2922 Client::new()
2923 .get("https://rest.example.com/v1/config")
2924 .build()
2925 .unwrap(),
2926 );
2927 session.authenticate(&mut req).await.unwrap();
2928 assert_eq!(
2929 req.headers().get("authorization").unwrap(),
2930 "Bearer tok-seed"
2931 );
2932
2933 oauth_mock.assert_async().await;
2934 }
2935
2936 #[tokio::test]
2937 async fn test_injected_oauth_manager_keeps_endpoint_and_options() {
2938 let mut server = Server::new_async().await;
2942 let config_mock = server
2945 .mock("GET", "/v1/config")
2946 .with_status(200)
2947 .with_body(r#"{"defaults": {"credential": "client1:secret1"}, "overrides": {}}"#)
2948 .create_async()
2949 .await;
2950
2951 let default_endpoint_mock = server
2953 .mock("POST", "/v1/oauth/tokens")
2954 .expect(0)
2955 .create_async()
2956 .await;
2957 let custom_endpoint_mock = server
2960 .mock("POST", "/custom/oauth/tokens")
2961 .match_header("x-tenant", "t1")
2962 .match_body(Matcher::AllOf(vec![
2965 Matcher::Regex("scope=catalog".to_string()),
2966 Matcher::Regex("audience=aud-1".to_string()),
2967 ]))
2968 .with_status(200)
2969 .with_body(
2970 r#"{
2971 "access_token": "ey000000000000",
2972 "token_type": "Bearer",
2973 "issued_token_type": "urn:ietf:params:oauth:token-type:access_token",
2974 "expires_in": 86400
2975 }"#,
2976 )
2977 .create_async()
2978 .await;
2979 let ns_mock = server
2980 .mock("GET", "/v1/namespaces")
2981 .match_header("authorization", "Bearer ey000000000000")
2982 .with_body(r#"{"namespaces": []}"#)
2983 .create_async()
2984 .await;
2985
2986 let manager = OAuth2Manager::new(format!("{}/custom/oauth/tokens", server.url()))
2987 .with_extra_headers(HeaderMap::from_iter([(
2988 HeaderName::from_static("x-tenant"),
2989 HeaderValue::from_static("t1"),
2990 )]))
2991 .with_extra_oauth_params(HashMap::from([(
2992 "audience".to_string(),
2993 "aud-1".to_string(),
2994 )]));
2995 let catalog = RestCatalog::new(
2996 SessionContext::empty(),
2997 RestCatalogConfig::builder().uri(server.url()).build(),
2998 Some(Box::new(manager)),
2999 Some(Arc::new(LocalFsStorageFactory)),
3000 Runtime::current(),
3001 None,
3002 );
3003
3004 catalog.list_namespaces(None).await.unwrap();
3005
3006 config_mock.assert_async().await;
3007 custom_endpoint_mock.assert_async().await;
3008 default_endpoint_mock.assert_async().await;
3009 ns_mock.assert_async().await;
3010 }
3011
3012 #[tokio::test]
3013 async fn test_props_token_takes_precedence_over_props_credential() {
3014 let mut server = Server::new_async().await;
3016 let oauth_mock = server
3017 .mock("POST", "/v1/oauth/tokens")
3018 .expect(0)
3019 .create_async()
3020 .await;
3021
3022 let manager = OAuth2Manager::new(format!("{}/v1/oauth/tokens", server.url()));
3023 let session = manager
3024 .init_session(
3025 &test_client(),
3026 &HashMap::from([
3027 ("token".to_string(), "tok-props".to_string()),
3028 ("credential".to_string(), "client1:secret1".to_string()),
3029 ]),
3030 )
3031 .await
3032 .unwrap();
3033
3034 let mut req = HttpRequest::new(
3035 Client::new()
3036 .get("https://rest.example.com/v1/config")
3037 .build()
3038 .unwrap(),
3039 );
3040 session.authenticate(&mut req).await.unwrap();
3041 assert_eq!(
3042 req.headers().get("authorization").unwrap(),
3043 "Bearer tok-props"
3044 );
3045 oauth_mock.assert_async().await;
3046 }
3047
3048 #[tokio::test]
3049 async fn test_manager_exchanges_over_the_catalog_client() {
3050 let mut server = Server::new_async().await;
3054 let token_mock = server
3055 .mock("POST", "/v1/oauth/tokens")
3056 .match_header("x-from", "manager")
3057 .match_header("x-catalog-only", Matcher::Missing)
3058 .with_status(200)
3059 .with_body(r#"{"access_token": "tok", "token_type": "Bearer"}"#)
3060 .create_async()
3061 .await;
3062
3063 let catalog_client = HttpClient::new(
3064 &RestCatalogConfig::builder()
3065 .uri(server.url())
3066 .props(HashMap::from([(
3067 "header.x-catalog-only".to_string(),
3068 "not-on-token-requests".to_string(),
3069 )]))
3070 .build(),
3071 )
3072 .unwrap();
3073
3074 let manager = OAuth2Manager::new(format!("{}/v1/oauth/tokens", server.url()))
3075 .with_credential(Some("client1".to_string()), "secret1".to_string())
3076 .with_extra_headers(HeaderMap::from_iter([(
3077 HeaderName::from_static("x-from"),
3078 HeaderValue::from_static("manager"),
3079 )]));
3080 let session = manager
3081 .init_session(&catalog_client, &HashMap::new())
3082 .await
3083 .unwrap();
3084
3085 let mut req = HttpRequest::new(
3086 Client::new()
3087 .get("https://rest.example.com/v1/namespaces")
3088 .build()
3089 .unwrap(),
3090 );
3091 session.authenticate(&mut req).await.unwrap();
3092 token_mock.assert_async().await;
3093 }
3094
3095 #[tokio::test]
3096 async fn test_handshake_is_authenticated_by_the_init_session() {
3097 let mut server = Server::new_async().await;
3100 let config_mock = server
3101 .mock("GET", "/v1/config")
3102 .match_header("authorization", "Bearer tok-init")
3103 .with_status(200)
3104 .with_body(r#"{"defaults": {}, "overrides": {}}"#)
3105 .create_async()
3106 .await;
3107
3108 let catalog = RestCatalog::new(
3109 SessionContext::empty(),
3110 RestCatalogConfig::builder()
3111 .uri(server.url())
3112 .props(HashMap::from([(
3113 "token".to_string(),
3114 "tok-init".to_string(),
3115 )]))
3116 .build(),
3117 None,
3118 Some(Arc::new(LocalFsStorageFactory)),
3119 Runtime::current(),
3120 None,
3121 );
3122
3123 catalog.client().await.unwrap();
3124 config_mock.assert_async().await;
3125 }
3126
3127 #[tokio::test]
3128 async fn test_init_session_receives_user_props() {
3129 use tokio::sync::Mutex as AsyncMutex;
3130
3131 #[derive(Debug)]
3134 struct PlainSession;
3135 #[async_trait]
3136 impl AuthSession for PlainSession {
3137 async fn authenticate(&self, _request: &mut HttpRequest) -> Result<()> {
3138 Ok(())
3139 }
3140 }
3141
3142 #[derive(Debug)]
3143 struct CapturingManager(Arc<AsyncMutex<Option<HashMap<String, String>>>>);
3144 #[async_trait]
3145 impl AuthManager for CapturingManager {
3146 async fn init_session(
3147 &self,
3148 _client: &HttpClient,
3149 props: &HashMap<String, String>,
3150 ) -> Result<Box<dyn AuthSession>> {
3151 *self.0.lock().await = Some(props.clone());
3152 Ok(Box::new(PlainSession))
3153 }
3154 async fn catalog_session(
3155 &self,
3156 _client: &HttpClient,
3157 _props: &HashMap<String, String>,
3158 ) -> Result<Arc<dyn AuthSession>> {
3159 Ok(Arc::new(PlainSession))
3160 }
3161 }
3162
3163 let mut server = Server::new_async().await;
3164 let config_mock = create_config_mock(&mut server).await;
3165 let captured = Arc::new(AsyncMutex::new(None));
3166 let catalog = RestCatalog::new(
3167 SessionContext::empty(),
3168 RestCatalogConfig::builder()
3169 .uri(server.url())
3170 .props(HashMap::from([(
3171 "token".to_string(),
3172 "tok-user".to_string(),
3173 )]))
3174 .build(),
3175 Some(Box::new(CapturingManager(captured.clone()))),
3176 Some(Arc::new(LocalFsStorageFactory)),
3177 Runtime::current(),
3178 None,
3179 );
3180
3181 catalog.client().await.unwrap();
3182 config_mock.assert_async().await;
3183 let props = captured.lock().await.clone().unwrap();
3184 assert_eq!(props.get("token").map(String::as_str), Some("tok-user"));
3185 assert_eq!(
3186 props.get(REST_CATALOG_PROP_URI).map(String::as_str),
3187 Some(server.url().as_str())
3188 );
3189 }
3190
3191 #[tokio::test]
3192 async fn test_catalog_session_receives_resolved_warehouse() {
3193 use tokio::sync::Mutex as AsyncMutex;
3194
3195 #[derive(Debug)]
3199 struct PlainSession;
3200 #[async_trait]
3201 impl AuthSession for PlainSession {
3202 async fn authenticate(&self, _request: &mut HttpRequest) -> Result<()> {
3203 Ok(())
3204 }
3205 }
3206
3207 #[derive(Debug)]
3208 struct CapturingManager(Arc<AsyncMutex<Option<HashMap<String, String>>>>);
3209 #[async_trait]
3210 impl AuthManager for CapturingManager {
3211 async fn init_session(
3212 &self,
3213 _client: &HttpClient,
3214 _props: &HashMap<String, String>,
3215 ) -> Result<Box<dyn AuthSession>> {
3216 Ok(Box::new(PlainSession))
3217 }
3218 async fn catalog_session(
3219 &self,
3220 _client: &HttpClient,
3221 props: &HashMap<String, String>,
3222 ) -> Result<Arc<dyn AuthSession>> {
3223 *self.0.lock().await = Some(props.clone());
3224 Ok(Arc::new(PlainSession))
3225 }
3226 }
3227
3228 let mut server = Server::new_async().await;
3230 let config_mock = server
3231 .mock("GET", "/v1/config")
3232 .match_query(Matcher::UrlEncoded(
3233 "warehouse".to_string(),
3234 "client-wh".to_string(),
3235 ))
3236 .with_status(200)
3237 .with_body(r#"{"defaults": {"warehouse": "default-wh"}, "overrides": {}}"#)
3238 .create_async()
3239 .await;
3240 let captured = Arc::new(AsyncMutex::new(None));
3241 let catalog = RestCatalog::new(
3242 SessionContext::empty(),
3243 RestCatalogConfig::builder()
3244 .uri(server.url())
3245 .warehouse("client-wh".to_string())
3246 .build(),
3247 Some(Box::new(CapturingManager(captured.clone()))),
3248 Some(Arc::new(LocalFsStorageFactory)),
3249 Runtime::current(),
3250 None,
3251 );
3252 catalog.client().await.unwrap();
3253 config_mock.assert_async().await;
3254 let props = captured.lock().await.clone().unwrap();
3255 assert_eq!(
3256 props.get("warehouse").map(String::as_str),
3257 Some("client-wh")
3258 );
3259
3260 let mut server = Server::new_async().await;
3262 let config_mock = server
3263 .mock("GET", "/v1/config")
3264 .match_query(Matcher::UrlEncoded(
3265 "warehouse".to_string(),
3266 "client-wh".to_string(),
3267 ))
3268 .with_status(200)
3269 .with_body(r#"{"defaults": {}, "overrides": {"warehouse": "override-wh"}}"#)
3270 .create_async()
3271 .await;
3272 let captured = Arc::new(AsyncMutex::new(None));
3273 let catalog = RestCatalog::new(
3274 SessionContext::empty(),
3275 RestCatalogConfig::builder()
3276 .uri(server.url())
3277 .warehouse("client-wh".to_string())
3278 .build(),
3279 Some(Box::new(CapturingManager(captured.clone()))),
3280 Some(Arc::new(LocalFsStorageFactory)),
3281 Runtime::current(),
3282 None,
3283 );
3284 catalog.client().await.unwrap();
3285 config_mock.assert_async().await;
3286 let props = captured.lock().await.clone().unwrap();
3287 assert_eq!(
3288 props.get("warehouse").map(String::as_str),
3289 Some("override-wh")
3290 );
3291 }
3292
3293 #[tokio::test]
3294 async fn test_init_session_dropped_before_catalog_session() {
3295 use std::sync::atomic::{AtomicBool, Ordering};
3296
3297 #[derive(Debug)]
3300 struct GuardSession(Arc<AtomicBool>);
3301 impl Drop for GuardSession {
3302 fn drop(&mut self) {
3303 self.0.store(true, Ordering::SeqCst);
3304 }
3305 }
3306 #[async_trait]
3307 impl AuthSession for GuardSession {
3308 async fn authenticate(&self, _request: &mut HttpRequest) -> Result<()> {
3309 Ok(())
3310 }
3311 }
3312
3313 #[derive(Debug)]
3314 struct PlainSession;
3315 #[async_trait]
3316 impl AuthSession for PlainSession {
3317 async fn authenticate(&self, _request: &mut HttpRequest) -> Result<()> {
3318 Ok(())
3319 }
3320 }
3321
3322 #[derive(Debug)]
3323 struct GuardManager(Arc<AtomicBool>);
3324 #[async_trait]
3325 impl AuthManager for GuardManager {
3326 async fn init_session(
3327 &self,
3328 _client: &HttpClient,
3329 _props: &HashMap<String, String>,
3330 ) -> Result<Box<dyn AuthSession>> {
3331 Ok(Box::new(GuardSession(self.0.clone())))
3332 }
3333 async fn catalog_session(
3334 &self,
3335 _client: &HttpClient,
3336 _props: &HashMap<String, String>,
3337 ) -> Result<Arc<dyn AuthSession>> {
3338 if !self.0.load(Ordering::SeqCst) {
3339 return Err(Error::new(
3340 ErrorKind::Unexpected,
3341 "init session must be dropped before catalog_session",
3342 ));
3343 }
3344 Ok(Arc::new(PlainSession))
3345 }
3346 }
3347
3348 let mut server = Server::new_async().await;
3349 let config_mock = create_config_mock(&mut server).await;
3350
3351 let dropped = Arc::new(AtomicBool::new(false));
3352 let catalog = RestCatalog::new(
3353 SessionContext::empty(),
3354 RestCatalogConfig::builder().uri(server.url()).build(),
3355 Some(Box::new(GuardManager(dropped.clone()))),
3356 Some(Arc::new(LocalFsStorageFactory)),
3357 Runtime::current(),
3358 None,
3359 );
3360
3361 catalog.client().await.unwrap();
3362 config_mock.assert_async().await;
3363 assert!(dropped.load(Ordering::SeqCst));
3364 }
3365
3366 #[tokio::test]
3367 async fn test_contextual_session_authenticates_each_catalog_request() {
3368 let mut server = Server::new_async().await;
3369 let config_mock = create_config_mock(&mut server).await;
3370 let first_page = server
3371 .mock("GET", "/v1/namespaces")
3372 .match_header("x-session-id", "session-123")
3373 .with_body(r#"{"namespaces": [["ns1"]], "next-page-token": "next"}"#)
3374 .create_async()
3375 .await;
3376 let second_page = server
3377 .mock("GET", "/v1/namespaces?pageToken=next")
3378 .match_header("x-session-id", "session-123")
3379 .with_body(r#"{"namespaces": [["ns2"]]}"#)
3380 .create_async()
3381 .await;
3382
3383 let (catalog, seen_session_ids) = context_catalog(
3384 RestCatalogConfig::builder().uri(server.url()).build(),
3385 false,
3386 );
3387 let context = SessionContext::builder()
3388 .session_id("session-123".to_string())
3389 .build();
3390
3391 let namespaces = catalog.list_namespaces(&context, None).await.unwrap();
3392
3393 assert_eq!(namespaces, vec![
3394 NamespaceIdent::new("ns1".to_string()),
3395 NamespaceIdent::new("ns2".to_string()),
3396 ]);
3397 assert_eq!(*seen_session_ids.lock().unwrap(), vec![
3398 "session-123",
3399 "session-123"
3400 ]);
3401 config_mock.assert_async().await;
3402 first_page.assert_async().await;
3403 second_page.assert_async().await;
3404 }
3405
3406 #[tokio::test]
3407 async fn test_contextual_session_authenticates_namespace_exists_head() {
3408 let mut server = Server::new_async().await;
3409 let config_mock = create_config_mock_with_exists_endpoints(&mut server).await;
3410 let exists_mock = server
3411 .mock("HEAD", "/v1/namespaces/ns1")
3412 .match_header("x-session-id", "namespace-session")
3413 .with_status(204)
3414 .create_async()
3415 .await;
3416 let (catalog, seen_session_ids) = context_catalog(
3417 RestCatalogConfig::builder().uri(server.url()).build(),
3418 false,
3419 );
3420 let context = SessionContext::builder()
3421 .session_id("namespace-session".to_string())
3422 .build();
3423
3424 assert!(
3425 catalog
3426 .namespace_exists(&context, &NamespaceIdent::new("ns1".to_string()))
3427 .await
3428 .unwrap()
3429 );
3430 assert_eq!(*seen_session_ids.lock().unwrap(), vec!["namespace-session"]);
3431 config_mock.assert_async().await;
3432 exists_mock.assert_async().await;
3433 }
3434
3435 #[tokio::test]
3436 async fn test_contextual_session_authenticates_table_exists_head() {
3437 let mut server = Server::new_async().await;
3438 let config_mock = create_config_mock_with_exists_endpoints(&mut server).await;
3439 let exists_mock = server
3440 .mock("HEAD", "/v1/namespaces/ns1/tables/table1")
3441 .match_header("x-session-id", "table-session")
3442 .with_status(204)
3443 .create_async()
3444 .await;
3445 let (catalog, seen_session_ids) = context_catalog(
3446 RestCatalogConfig::builder().uri(server.url()).build(),
3447 false,
3448 );
3449 let context = SessionContext::builder()
3450 .session_id("table-session".to_string())
3451 .build();
3452
3453 assert!(
3454 catalog
3455 .table_exists(
3456 &context,
3457 &TableIdent::new(NamespaceIdent::new("ns1".to_string()), "table1".to_string(),),
3458 )
3459 .await
3460 .unwrap()
3461 );
3462 assert_eq!(*seen_session_ids.lock().unwrap(), vec!["table-session"]);
3463 config_mock.assert_async().await;
3464 exists_mock.assert_async().await;
3465 }
3466
3467 #[tokio::test]
3468 async fn test_contextual_session_authenticates_write_operation() {
3469 let mut server = Server::new_async().await;
3470 let config_mock = create_config_mock(&mut server).await;
3471 let drop_mock = server
3472 .mock("DELETE", "/v1/namespaces/ns1")
3473 .match_header("x-session-id", "write-session")
3474 .with_status(204)
3475 .create_async()
3476 .await;
3477 let (catalog, seen_session_ids) = context_catalog(
3478 RestCatalogConfig::builder().uri(server.url()).build(),
3479 false,
3480 );
3481 let context = SessionContext::builder()
3482 .session_id("write-session".to_string())
3483 .build();
3484
3485 catalog
3486 .drop_namespace(&context, &NamespaceIdent::new("ns1".to_string()))
3487 .await
3488 .unwrap();
3489 assert_eq!(*seen_session_ids.lock().unwrap(), vec!["write-session"]);
3490 config_mock.assert_async().await;
3491 drop_mock.assert_async().await;
3492 }
3493
3494 #[tokio::test]
3495 async fn test_contextual_session_error_prevents_operation_request() {
3496 let mut server = Server::new_async().await;
3497 let config_mock = create_config_mock(&mut server).await;
3498 let list_mock = server
3499 .mock("GET", "/v1/namespaces")
3500 .expect(0)
3501 .create_async()
3502 .await;
3503 let (catalog, seen_session_ids) =
3504 context_catalog(RestCatalogConfig::builder().uri(server.url()).build(), true);
3505 let context = SessionContext::builder()
3506 .session_id("failing-session".to_string())
3507 .build();
3508
3509 let error = catalog.list_namespaces(&context, None).await.unwrap_err();
3510 assert_eq!(error.kind(), ErrorKind::Unexpected);
3511 assert_eq!(error.message(), "contextual session failure");
3512 assert_eq!(*seen_session_ids.lock().unwrap(), vec!["failing-session"]);
3513 config_mock.assert_async().await;
3514 list_mock.assert_async().await;
3515 }
3516
3517 #[test]
3518 fn test_config_debug_redacts_secrets() {
3519 let config = RestCatalogConfig::builder()
3520 .uri("http://localhost".to_string())
3521 .props(HashMap::from([
3522 ("token".to_string(), "tok-secret".to_string()),
3523 ("credential".to_string(), "id:cred-secret".to_string()),
3524 ("header.authorization".to_string(), "Basic xyz".to_string()),
3525 ("adls.account-key".to_string(), "adls-secret".to_string()),
3526 ("s3.sse.key".to_string(), "sse-secret".to_string()),
3527 (
3528 "adls.connection-string".to_string(),
3529 "cs-secret".to_string(),
3530 ),
3531 ("warehouse".to_string(), "wh1".to_string()),
3532 ]))
3533 .build();
3534
3535 let out = format!("{config:?}");
3536 assert!(!out.contains("tok-secret"));
3537 assert!(!out.contains("cred-secret"));
3538 assert!(!out.contains("Basic xyz"));
3539 assert!(!out.contains("adls-secret"));
3540 assert!(!out.contains("sse-secret"));
3541 assert!(!out.contains("cs-secret"));
3542 assert!(out.contains("[REDACTED]"));
3543 assert!(out.contains("wh1"));
3544 }
3545
3546 #[tokio::test]
3547 async fn test_auth_type_defaults() {
3548 let bare = RestCatalogConfig::builder()
3551 .uri("http://localhost".to_string())
3552 .build();
3553 assert!(
3554 format!("{:?}", test_catalog(bare).resolve_auth_manager().unwrap())
3555 .contains("NoopAuthManager")
3556 );
3557
3558 let with_token = RestCatalogConfig::builder()
3559 .uri("http://localhost".to_string())
3560 .props(HashMap::from([("token".to_string(), "tok".to_string())]))
3561 .build();
3562 assert!(
3563 format!(
3564 "{:?}",
3565 test_catalog(with_token).resolve_auth_manager().unwrap()
3566 )
3567 .contains("OAuth2Manager")
3568 );
3569
3570 let mixed_case = RestCatalogConfig::builder()
3572 .uri("http://localhost".to_string())
3573 .props(HashMap::from([(
3574 REST_CATALOG_PROP_AUTH_TYPE.to_string(),
3575 "OAuth2".to_string(),
3576 )]))
3577 .build();
3578 assert!(
3579 format!(
3580 "{:?}",
3581 test_catalog(mixed_case).resolve_auth_manager().unwrap()
3582 )
3583 .contains("OAuth2Manager")
3584 );
3585
3586 let with_endpoint = RestCatalogConfig::builder()
3589 .uri("http://localhost".to_string())
3590 .props(HashMap::from([(
3591 "oauth2-server-uri".to_string(),
3592 "http://auth.example.com/tokens".to_string(),
3593 )]))
3594 .build();
3595 assert!(
3596 format!(
3597 "{:?}",
3598 test_catalog(with_endpoint).resolve_auth_manager().unwrap()
3599 )
3600 .contains("OAuth2Manager")
3601 );
3602 }
3603
3604 #[tokio::test]
3605 async fn test_unknown_auth_type_is_rejected() {
3606 let props = HashMap::from([(
3607 REST_CATALOG_PROP_AUTH_TYPE.to_string(),
3608 "kerberos".to_string(),
3609 )]);
3610 let config = RestCatalogConfig::builder()
3611 .uri("http://localhost".to_string())
3612 .props(props)
3613 .build();
3614
3615 let err = test_catalog(config).resolve_auth_manager().unwrap_err();
3616 assert!(err.message().contains(REST_CATALOG_PROP_AUTH_TYPE));
3617 }
3618
3619 #[tokio::test]
3620 async fn test_with_auth_manager_overrides_config() {
3621 #[derive(Debug)]
3623 struct StubAuthManager;
3624 #[async_trait]
3625 impl AuthManager for StubAuthManager {
3626 async fn init_session(
3627 &self,
3628 _client: &HttpClient,
3629 _props: &HashMap<String, String>,
3630 ) -> Result<Box<dyn AuthSession>> {
3631 unimplemented!()
3632 }
3633 async fn catalog_session(
3634 &self,
3635 _client: &HttpClient,
3636 _props: &HashMap<String, String>,
3637 ) -> Result<Arc<dyn AuthSession>> {
3638 unimplemented!()
3639 }
3640 }
3641
3642 let config = RestCatalogConfig::builder()
3643 .uri("http://localhost".to_string())
3644 .props(HashMap::from([(
3645 REST_CATALOG_PROP_AUTH_TYPE.to_string(),
3646 "kerberos".to_string(),
3647 )]))
3648 .build();
3649
3650 let catalog = test_catalog_with(config, StubAuthManager);
3652 assert!(catalog.resolve_auth_manager().is_ok());
3653 }
3654
3655 #[tokio::test]
3656 async fn test_list_namespace_with_pagination() {
3657 let mut server = Server::new_async().await;
3658
3659 let config_mock = create_config_mock(&mut server).await;
3660
3661 let list_ns_mock_page1 = server
3662 .mock("GET", "/v1/namespaces")
3663 .with_body(
3664 r#"{
3665 "namespaces": [
3666 ["ns1", "ns11"],
3667 ["ns2"]
3668 ],
3669 "next-page-token": "token123"
3670 }"#,
3671 )
3672 .create_async()
3673 .await;
3674
3675 let list_ns_mock_page2 = server
3676 .mock("GET", "/v1/namespaces?pageToken=token123")
3677 .with_body(
3678 r#"{
3679 "namespaces": [
3680 ["ns3"],
3681 ["ns4", "ns41"]
3682 ]
3683 }"#,
3684 )
3685 .create_async()
3686 .await;
3687
3688 let catalog = session_catalog(RestCatalogConfig::builder().uri(server.url()).build());
3689
3690 let namespaces = catalog
3691 .list_namespaces(&SessionContext::empty(), None)
3692 .await
3693 .unwrap();
3694
3695 let expected_ns = vec![
3696 NamespaceIdent::from_vec(vec!["ns1".to_string(), "ns11".to_string()]).unwrap(),
3697 NamespaceIdent::from_vec(vec!["ns2".to_string()]).unwrap(),
3698 NamespaceIdent::from_vec(vec!["ns3".to_string()]).unwrap(),
3699 NamespaceIdent::from_vec(vec!["ns4".to_string(), "ns41".to_string()]).unwrap(),
3700 ];
3701
3702 assert_eq!(expected_ns, namespaces);
3703
3704 config_mock.assert_async().await;
3705 list_ns_mock_page1.assert_async().await;
3706 list_ns_mock_page2.assert_async().await;
3707 }
3708
3709 #[tokio::test]
3710 async fn test_list_namespace_with_multiple_pages() {
3711 let mut server = Server::new_async().await;
3712
3713 let config_mock = create_config_mock(&mut server).await;
3714
3715 let list_ns_mock_page1 = server
3717 .mock("GET", "/v1/namespaces")
3718 .with_body(
3719 r#"{
3720 "namespaces": [
3721 ["ns1", "ns11"],
3722 ["ns2"]
3723 ],
3724 "next-page-token": "page2"
3725 }"#,
3726 )
3727 .create_async()
3728 .await;
3729
3730 let list_ns_mock_page2 = server
3732 .mock("GET", "/v1/namespaces?pageToken=page2")
3733 .with_body(
3734 r#"{
3735 "namespaces": [
3736 ["ns3"],
3737 ["ns4", "ns41"]
3738 ],
3739 "next-page-token": "page3"
3740 }"#,
3741 )
3742 .create_async()
3743 .await;
3744
3745 let list_ns_mock_page3 = server
3747 .mock("GET", "/v1/namespaces?pageToken=page3")
3748 .with_body(
3749 r#"{
3750 "namespaces": [
3751 ["ns5", "ns51", "ns511"]
3752 ],
3753 "next-page-token": "page4"
3754 }"#,
3755 )
3756 .create_async()
3757 .await;
3758
3759 let list_ns_mock_page4 = server
3761 .mock("GET", "/v1/namespaces?pageToken=page4")
3762 .with_body(
3763 r#"{
3764 "namespaces": [
3765 ["ns6"],
3766 ["ns7"]
3767 ],
3768 "next-page-token": "page5"
3769 }"#,
3770 )
3771 .create_async()
3772 .await;
3773
3774 let list_ns_mock_page5 = server
3776 .mock("GET", "/v1/namespaces?pageToken=page5")
3777 .with_body(
3778 r#"{
3779 "namespaces": [
3780 ["ns8", "ns81"]
3781 ]
3782 }"#,
3783 )
3784 .create_async()
3785 .await;
3786
3787 let catalog = session_catalog(RestCatalogConfig::builder().uri(server.url()).build());
3788
3789 let namespaces = catalog
3790 .list_namespaces(&SessionContext::empty(), None)
3791 .await
3792 .unwrap();
3793
3794 let expected_ns = vec![
3795 NamespaceIdent::from_vec(vec!["ns1".to_string(), "ns11".to_string()]).unwrap(),
3796 NamespaceIdent::from_vec(vec!["ns2".to_string()]).unwrap(),
3797 NamespaceIdent::from_vec(vec!["ns3".to_string()]).unwrap(),
3798 NamespaceIdent::from_vec(vec!["ns4".to_string(), "ns41".to_string()]).unwrap(),
3799 NamespaceIdent::from_vec(vec![
3800 "ns5".to_string(),
3801 "ns51".to_string(),
3802 "ns511".to_string(),
3803 ])
3804 .unwrap(),
3805 NamespaceIdent::from_vec(vec!["ns6".to_string()]).unwrap(),
3806 NamespaceIdent::from_vec(vec!["ns7".to_string()]).unwrap(),
3807 NamespaceIdent::from_vec(vec!["ns8".to_string(), "ns81".to_string()]).unwrap(),
3808 ];
3809
3810 assert_eq!(expected_ns, namespaces);
3811
3812 config_mock.assert_async().await;
3814 list_ns_mock_page1.assert_async().await;
3815 list_ns_mock_page2.assert_async().await;
3816 list_ns_mock_page3.assert_async().await;
3817 list_ns_mock_page4.assert_async().await;
3818 list_ns_mock_page5.assert_async().await;
3819 }
3820
3821 #[tokio::test]
3822 async fn test_create_namespace() {
3823 let mut server = Server::new_async().await;
3824
3825 let config_mock = create_config_mock(&mut server).await;
3826
3827 let create_ns_mock = server
3828 .mock("POST", "/v1/namespaces")
3829 .with_body(
3830 r#"{
3831 "namespace": [ "ns1", "ns11"],
3832 "properties" : {
3833 "key1": "value1"
3834 }
3835 }"#,
3836 )
3837 .create_async()
3838 .await;
3839
3840 let catalog = session_catalog(RestCatalogConfig::builder().uri(server.url()).build());
3841
3842 let namespaces = catalog
3843 .create_namespace(
3844 &SessionContext::empty(),
3845 &NamespaceIdent::from_vec(vec!["ns1".to_string(), "ns11".to_string()]).unwrap(),
3846 HashMap::from([("key1".to_string(), "value1".to_string())]),
3847 )
3848 .await
3849 .unwrap();
3850
3851 let expected_ns = Namespace::with_properties(
3852 NamespaceIdent::from_vec(vec!["ns1".to_string(), "ns11".to_string()]).unwrap(),
3853 HashMap::from([("key1".to_string(), "value1".to_string())]),
3854 );
3855
3856 assert_eq!(expected_ns, namespaces);
3857
3858 config_mock.assert_async().await;
3859 create_ns_mock.assert_async().await;
3860 }
3861
3862 #[tokio::test]
3863 async fn test_get_namespace() {
3864 let mut server = Server::new_async().await;
3865
3866 let config_mock = create_config_mock(&mut server).await;
3867
3868 let get_ns_mock = server
3869 .mock("GET", "/v1/namespaces/ns1")
3870 .with_body(
3871 r#"{
3872 "namespace": [ "ns1"],
3873 "properties" : {
3874 "key1": "value1"
3875 }
3876 }"#,
3877 )
3878 .create_async()
3879 .await;
3880
3881 let catalog = session_catalog(RestCatalogConfig::builder().uri(server.url()).build());
3882
3883 let namespaces = catalog
3884 .get_namespace(
3885 &SessionContext::empty(),
3886 &NamespaceIdent::new("ns1".to_string()),
3887 )
3888 .await
3889 .unwrap();
3890
3891 let expected_ns = Namespace::with_properties(
3892 NamespaceIdent::new("ns1".to_string()),
3893 HashMap::from([("key1".to_string(), "value1".to_string())]),
3894 );
3895
3896 assert_eq!(expected_ns, namespaces);
3897
3898 config_mock.assert_async().await;
3899 get_ns_mock.assert_async().await;
3900 }
3901
3902 #[tokio::test]
3903 async fn check_namespace_exists() {
3904 let mut server = Server::new_async().await;
3905
3906 let config_mock = create_config_mock_with_exists_endpoints(&mut server).await;
3907
3908 let get_ns_mock = server
3909 .mock("HEAD", "/v1/namespaces/ns1")
3910 .with_status(204)
3911 .create_async()
3912 .await;
3913
3914 let catalog = session_catalog(RestCatalogConfig::builder().uri(server.url()).build());
3915
3916 assert!(
3917 catalog
3918 .namespace_exists(
3919 &SessionContext::empty(),
3920 &NamespaceIdent::new("ns1".to_string())
3921 )
3922 .await
3923 .unwrap()
3924 );
3925
3926 config_mock.assert_async().await;
3927 get_ns_mock.assert_async().await;
3928 }
3929
3930 #[tokio::test]
3931 async fn test_namespace_exists_falls_back_to_get_when_head_not_advertised() {
3932 let mut server = Server::new_async().await;
3933
3934 let config_mock = create_config_mock(&mut server).await;
3937 let get_ns_mock = server
3938 .mock("GET", "/v1/namespaces/ns1")
3939 .with_status(200)
3940 .with_body(
3941 r#"{
3942 "namespace": ["ns1"],
3943 "properties": {}
3944 }"#,
3945 )
3946 .create_async()
3947 .await;
3948
3949 let catalog = RestCatalog::new(
3950 SessionContext::empty(),
3951 RestCatalogConfig::builder().uri(server.url()).build(),
3952 None,
3953 Some(Arc::new(LocalFsStorageFactory)),
3954 Runtime::current(),
3955 None,
3956 );
3957
3958 assert!(
3959 catalog
3960 .namespace_exists(&NamespaceIdent::new("ns1".to_string()))
3961 .await
3962 .unwrap()
3963 );
3964
3965 config_mock.assert_async().await;
3966 get_ns_mock.assert_async().await;
3967 }
3968
3969 #[tokio::test]
3970 async fn test_drop_namespace() {
3971 let mut server = Server::new_async().await;
3972
3973 let config_mock = create_config_mock(&mut server).await;
3974
3975 let drop_ns_mock = server
3976 .mock("DELETE", "/v1/namespaces/ns1")
3977 .with_status(204)
3978 .create_async()
3979 .await;
3980
3981 let catalog = session_catalog(RestCatalogConfig::builder().uri(server.url()).build());
3982
3983 catalog
3984 .drop_namespace(
3985 &SessionContext::empty(),
3986 &NamespaceIdent::new("ns1".to_string()),
3987 )
3988 .await
3989 .unwrap();
3990
3991 config_mock.assert_async().await;
3992 drop_ns_mock.assert_async().await;
3993 }
3994
3995 #[tokio::test]
3996 async fn test_list_tables() {
3997 let mut server = Server::new_async().await;
3998
3999 let config_mock = create_config_mock(&mut server).await;
4000
4001 let list_tables_mock = server
4002 .mock("GET", "/v1/namespaces/ns1/tables")
4003 .with_status(200)
4004 .with_body(
4005 r#"{
4006 "identifiers": [
4007 {
4008 "namespace": ["ns1"],
4009 "name": "table1"
4010 },
4011 {
4012 "namespace": ["ns1"],
4013 "name": "table2"
4014 }
4015 ]
4016 }"#,
4017 )
4018 .create_async()
4019 .await;
4020
4021 let catalog = session_catalog(RestCatalogConfig::builder().uri(server.url()).build());
4022
4023 let tables = catalog
4024 .list_tables(
4025 &SessionContext::empty(),
4026 &NamespaceIdent::new("ns1".to_string()),
4027 )
4028 .await
4029 .unwrap();
4030
4031 let expected_tables = vec![
4032 TableIdent::new(NamespaceIdent::new("ns1".to_string()), "table1".to_string()),
4033 TableIdent::new(NamespaceIdent::new("ns1".to_string()), "table2".to_string()),
4034 ];
4035
4036 assert_eq!(tables, expected_tables);
4037
4038 config_mock.assert_async().await;
4039 list_tables_mock.assert_async().await;
4040 }
4041
4042 #[tokio::test]
4043 async fn test_list_tables_with_pagination() {
4044 let mut server = Server::new_async().await;
4045
4046 let config_mock = create_config_mock(&mut server).await;
4047
4048 let list_tables_mock_page1 = server
4049 .mock("GET", "/v1/namespaces/ns1/tables")
4050 .with_status(200)
4051 .with_body(
4052 r#"{
4053 "identifiers": [
4054 {
4055 "namespace": ["ns1"],
4056 "name": "table1"
4057 },
4058 {
4059 "namespace": ["ns1"],
4060 "name": "table2"
4061 }
4062 ],
4063 "next-page-token": "token456"
4064 }"#,
4065 )
4066 .create_async()
4067 .await;
4068
4069 let list_tables_mock_page2 = server
4070 .mock("GET", "/v1/namespaces/ns1/tables?pageToken=token456")
4071 .with_status(200)
4072 .with_body(
4073 r#"{
4074 "identifiers": [
4075 {
4076 "namespace": ["ns1"],
4077 "name": "table3"
4078 },
4079 {
4080 "namespace": ["ns1"],
4081 "name": "table4"
4082 }
4083 ]
4084 }"#,
4085 )
4086 .create_async()
4087 .await;
4088
4089 let catalog = session_catalog(RestCatalogConfig::builder().uri(server.url()).build());
4090
4091 let tables = catalog
4092 .list_tables(
4093 &SessionContext::empty(),
4094 &NamespaceIdent::new("ns1".to_string()),
4095 )
4096 .await
4097 .unwrap();
4098
4099 let expected_tables = vec![
4100 TableIdent::new(NamespaceIdent::new("ns1".to_string()), "table1".to_string()),
4101 TableIdent::new(NamespaceIdent::new("ns1".to_string()), "table2".to_string()),
4102 TableIdent::new(NamespaceIdent::new("ns1".to_string()), "table3".to_string()),
4103 TableIdent::new(NamespaceIdent::new("ns1".to_string()), "table4".to_string()),
4104 ];
4105
4106 assert_eq!(tables, expected_tables);
4107
4108 config_mock.assert_async().await;
4109 list_tables_mock_page1.assert_async().await;
4110 list_tables_mock_page2.assert_async().await;
4111 }
4112
4113 #[tokio::test]
4114 async fn test_list_tables_with_multiple_pages() {
4115 let mut server = Server::new_async().await;
4116
4117 let config_mock = create_config_mock(&mut server).await;
4118
4119 let list_tables_mock_page1 = server
4121 .mock("GET", "/v1/namespaces/ns1/tables")
4122 .with_status(200)
4123 .with_body(
4124 r#"{
4125 "identifiers": [
4126 {
4127 "namespace": ["ns1"],
4128 "name": "table1"
4129 },
4130 {
4131 "namespace": ["ns1"],
4132 "name": "table2"
4133 }
4134 ],
4135 "next-page-token": "page2"
4136 }"#,
4137 )
4138 .create_async()
4139 .await;
4140
4141 let list_tables_mock_page2 = server
4143 .mock("GET", "/v1/namespaces/ns1/tables?pageToken=page2")
4144 .with_status(200)
4145 .with_body(
4146 r#"{
4147 "identifiers": [
4148 {
4149 "namespace": ["ns1"],
4150 "name": "table3"
4151 },
4152 {
4153 "namespace": ["ns1"],
4154 "name": "table4"
4155 }
4156 ],
4157 "next-page-token": "page3"
4158 }"#,
4159 )
4160 .create_async()
4161 .await;
4162
4163 let list_tables_mock_page3 = server
4165 .mock("GET", "/v1/namespaces/ns1/tables?pageToken=page3")
4166 .with_status(200)
4167 .with_body(
4168 r#"{
4169 "identifiers": [
4170 {
4171 "namespace": ["ns1"],
4172 "name": "table5"
4173 }
4174 ],
4175 "next-page-token": "page4"
4176 }"#,
4177 )
4178 .create_async()
4179 .await;
4180
4181 let list_tables_mock_page4 = server
4183 .mock("GET", "/v1/namespaces/ns1/tables?pageToken=page4")
4184 .with_status(200)
4185 .with_body(
4186 r#"{
4187 "identifiers": [
4188 {
4189 "namespace": ["ns1"],
4190 "name": "table6"
4191 },
4192 {
4193 "namespace": ["ns1"],
4194 "name": "table7"
4195 }
4196 ],
4197 "next-page-token": "page5"
4198 }"#,
4199 )
4200 .create_async()
4201 .await;
4202
4203 let list_tables_mock_page5 = server
4205 .mock("GET", "/v1/namespaces/ns1/tables?pageToken=page5")
4206 .with_status(200)
4207 .with_body(
4208 r#"{
4209 "identifiers": [
4210 {
4211 "namespace": ["ns1"],
4212 "name": "table8"
4213 }
4214 ]
4215 }"#,
4216 )
4217 .create_async()
4218 .await;
4219
4220 let catalog = session_catalog(RestCatalogConfig::builder().uri(server.url()).build());
4221
4222 let tables = catalog
4223 .list_tables(
4224 &SessionContext::empty(),
4225 &NamespaceIdent::new("ns1".to_string()),
4226 )
4227 .await
4228 .unwrap();
4229
4230 let expected_tables = vec![
4231 TableIdent::new(NamespaceIdent::new("ns1".to_string()), "table1".to_string()),
4232 TableIdent::new(NamespaceIdent::new("ns1".to_string()), "table2".to_string()),
4233 TableIdent::new(NamespaceIdent::new("ns1".to_string()), "table3".to_string()),
4234 TableIdent::new(NamespaceIdent::new("ns1".to_string()), "table4".to_string()),
4235 TableIdent::new(NamespaceIdent::new("ns1".to_string()), "table5".to_string()),
4236 TableIdent::new(NamespaceIdent::new("ns1".to_string()), "table6".to_string()),
4237 TableIdent::new(NamespaceIdent::new("ns1".to_string()), "table7".to_string()),
4238 TableIdent::new(NamespaceIdent::new("ns1".to_string()), "table8".to_string()),
4239 ];
4240
4241 assert_eq!(tables, expected_tables);
4242
4243 config_mock.assert_async().await;
4245 list_tables_mock_page1.assert_async().await;
4246 list_tables_mock_page2.assert_async().await;
4247 list_tables_mock_page3.assert_async().await;
4248 list_tables_mock_page4.assert_async().await;
4249 list_tables_mock_page5.assert_async().await;
4250 }
4251
4252 #[tokio::test]
4253 async fn test_drop_tables() {
4254 let mut server = Server::new_async().await;
4255
4256 let config_mock = create_config_mock(&mut server).await;
4257
4258 let delete_table_mock = server
4259 .mock("DELETE", "/v1/namespaces/ns1/tables/table1")
4260 .with_status(204)
4261 .create_async()
4262 .await;
4263
4264 let catalog = session_catalog(RestCatalogConfig::builder().uri(server.url()).build());
4265
4266 catalog
4267 .drop_table(
4268 &SessionContext::empty(),
4269 &TableIdent::new(NamespaceIdent::new("ns1".to_string()), "table1".to_string()),
4270 )
4271 .await
4272 .unwrap();
4273
4274 config_mock.assert_async().await;
4275 delete_table_mock.assert_async().await;
4276 }
4277
4278 #[tokio::test]
4279 async fn test_check_table_exists() {
4280 let mut server = Server::new_async().await;
4281
4282 let config_mock = create_config_mock_with_exists_endpoints(&mut server).await;
4283
4284 let check_table_exists_mock = server
4285 .mock("HEAD", "/v1/namespaces/ns1/tables/table1")
4286 .with_status(204)
4287 .create_async()
4288 .await;
4289
4290 let catalog = session_catalog(RestCatalogConfig::builder().uri(server.url()).build());
4291
4292 assert!(
4293 catalog
4294 .table_exists(
4295 &SessionContext::empty(),
4296 &TableIdent::new(NamespaceIdent::new("ns1".to_string()), "table1".to_string(),),
4297 )
4298 .await
4299 .unwrap()
4300 );
4301
4302 config_mock.assert_async().await;
4303 check_table_exists_mock.assert_async().await;
4304 }
4305
4306 #[tokio::test]
4307 async fn test_table_exists_falls_back_to_load_when_head_not_advertised() {
4308 let mut server = Server::new_async().await;
4309
4310 let config_mock = create_config_mock(&mut server).await;
4313 let load_table_mock = server
4314 .mock("GET", "/v1/namespaces/ns1/tables/table1")
4315 .with_status(200)
4316 .with_body_from_file(format!(
4317 "{}/testdata/{}",
4318 env!("CARGO_MANIFEST_DIR"),
4319 "load_table_response.json"
4320 ))
4321 .create_async()
4322 .await;
4323
4324 let catalog = RestCatalog::new(
4325 SessionContext::empty(),
4326 RestCatalogConfig::builder().uri(server.url()).build(),
4327 None,
4328 Some(Arc::new(LocalFsStorageFactory)),
4329 Runtime::current(),
4330 None,
4331 );
4332
4333 assert!(
4334 catalog
4335 .table_exists(&TableIdent::new(
4336 NamespaceIdent::new("ns1".to_string()),
4337 "table1".to_string(),
4338 ))
4339 .await
4340 .unwrap()
4341 );
4342
4343 config_mock.assert_async().await;
4344 load_table_mock.assert_async().await;
4345 }
4346
4347 #[tokio::test]
4348 async fn test_rename_table() {
4349 let mut server = Server::new_async().await;
4350
4351 let config_mock = create_config_mock(&mut server).await;
4352
4353 let rename_table_mock = server
4354 .mock("POST", "/v1/tables/rename")
4355 .with_status(204)
4356 .create_async()
4357 .await;
4358
4359 let catalog = session_catalog(RestCatalogConfig::builder().uri(server.url()).build());
4360
4361 catalog
4362 .rename_table(
4363 &SessionContext::empty(),
4364 &TableIdent::new(NamespaceIdent::new("ns1".to_string()), "table1".to_string()),
4365 &TableIdent::new(NamespaceIdent::new("ns1".to_string()), "table2".to_string()),
4366 )
4367 .await
4368 .unwrap();
4369
4370 config_mock.assert_async().await;
4371 rename_table_mock.assert_async().await;
4372 }
4373
4374 #[tokio::test]
4375 async fn test_load_table() {
4376 let mut server = Server::new_async().await;
4377
4378 let config_mock = create_config_mock(&mut server).await;
4379
4380 let rename_table_mock = server
4381 .mock("GET", "/v1/namespaces/ns1/tables/test1")
4382 .with_status(200)
4383 .with_body_from_file(format!(
4384 "{}/testdata/{}",
4385 env!("CARGO_MANIFEST_DIR"),
4386 "load_table_response.json"
4387 ))
4388 .create_async()
4389 .await;
4390
4391 let catalog = session_catalog(RestCatalogConfig::builder().uri(server.url()).build());
4392
4393 let table = catalog
4394 .load_table(
4395 &SessionContext::empty(),
4396 &TableIdent::new(NamespaceIdent::new("ns1".to_string()), "test1".to_string()),
4397 )
4398 .await
4399 .unwrap();
4400
4401 assert_eq!(
4402 &TableIdent::from_strs(vec!["ns1", "test1"]).unwrap(),
4403 table.identifier()
4404 );
4405 assert_eq!(
4406 "s3://warehouse/database/table/metadata/00001-5f2f8166-244c-4eae-ac36-384ecdec81fc.gz.metadata.json",
4407 table.metadata_location().unwrap()
4408 );
4409 assert_eq!(FormatVersion::V1, table.metadata().format_version());
4410 assert_eq!("s3://warehouse/database/table", table.metadata().location());
4411 assert_eq!(
4412 uuid!("b55d9dda-6561-423a-8bfc-787980ce421f"),
4413 table.metadata().uuid()
4414 );
4415 assert_eq!(
4416 Utc.timestamp_millis_opt(1646787054459).unwrap(),
4417 table.metadata().last_updated_timestamp().unwrap()
4418 );
4419 assert_eq!(
4420 vec![&Arc::new(
4421 Schema::builder()
4422 .with_fields(vec![
4423 NestedField::optional(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
4424 NestedField::optional(2, "data", Type::Primitive(PrimitiveType::String))
4425 .into(),
4426 ])
4427 .build()
4428 .unwrap()
4429 )],
4430 table.metadata().schemas_iter().collect::<Vec<_>>()
4431 );
4432 assert_eq!(
4433 &HashMap::from([
4434 ("owner".to_string(), "bryan".to_string()),
4435 (
4436 "write.metadata.compression-codec".to_string(),
4437 "gzip".to_string()
4438 )
4439 ]),
4440 table.metadata().properties()
4441 );
4442 assert_eq!(vec![&Arc::new(Snapshot::builder()
4443 .with_snapshot_id(3497810964824022504)
4444 .with_timestamp_ms(1646787054459)
4445 .with_manifest_list("s3://warehouse/database/table/metadata/snap-3497810964824022504-1-c4f68204-666b-4e50-a9df-b10c34bf6b82.avro")
4446 .with_sequence_number(0)
4447 .with_schema_id(0)
4448 .with_summary(Summary {
4449 operation: Operation::Append,
4450 additional_properties: HashMap::from_iter([
4451 ("spark.app.id", "local-1646787004168"),
4452 ("added-data-files", "1"),
4453 ("added-records", "1"),
4454 ("added-files-size", "697"),
4455 ("changed-partition-count", "1"),
4456 ("total-records", "1"),
4457 ("total-files-size", "697"),
4458 ("total-data-files", "1"),
4459 ("total-delete-files", "0"),
4460 ("total-position-deletes", "0"),
4461 ("total-equality-deletes", "0")
4462 ].iter().map(|p| (p.0.to_string(), p.1.to_string()))),
4463 }).build()
4464 )], table.metadata().snapshots().collect::<Vec<_>>());
4465 assert_eq!(
4466 &[SnapshotLog {
4467 timestamp_ms: 1646787054459,
4468 snapshot_id: 3497810964824022504,
4469 }],
4470 table.metadata().history()
4471 );
4472 assert_eq!(
4473 vec![&Arc::new(SortOrder {
4474 order_id: 0,
4475 fields: vec![],
4476 })],
4477 table.metadata().sort_orders_iter().collect::<Vec<_>>()
4478 );
4479
4480 config_mock.assert_async().await;
4481 rename_table_mock.assert_async().await;
4482 }
4483
4484 fn load_table_body(tag: &str) -> serde_json::Value {
4486 let file = File::open(format!(
4487 "{}/testdata/load_table_response.json",
4488 env!("CARGO_MANIFEST_DIR")
4489 ))
4490 .unwrap();
4491 let mut body: serde_json::Value = serde_json::from_reader(BufReader::new(file)).unwrap();
4492 body["metadata-location"] = json!(format!("s3://warehouse/{tag}.metadata.json"));
4493 body
4494 }
4495
4496 async fn load_table_mock(
4499 server: &mut ServerGuard,
4500 if_none_match: Matcher,
4501 status: usize,
4502 etag: Option<&str>,
4503 ) -> Mock {
4504 let mut mock = server
4505 .mock("GET", "/v1/namespaces/ns1/tables/test1")
4506 .match_header("if-none-match", if_none_match)
4507 .with_status(status);
4508 if let Some(etag) = etag {
4509 mock = mock.with_header("etag", etag);
4510 if status == 200 {
4511 mock = mock.with_body(load_table_body(etag.trim_matches('"')).to_string());
4512 }
4513 }
4514 mock.create_async().await
4515 }
4516
4517 fn test_table_ident() -> TableIdent {
4518 TableIdent::from_strs(["ns1", "test1"]).unwrap()
4519 }
4520
4521 #[tokio::test]
4522 async fn test_load_table_revalidates_cached_table_with_etag() {
4523 let mut server = Server::new_async().await;
4524 let _config = create_config_mock(&mut server).await;
4525 let _v1 = load_table_mock(&mut server, Matcher::Missing, 200, Some("\"v1\"")).await;
4526 let reuse_v1 = load_table_mock(&mut server, "\"v1\"".into(), 304, None).await;
4527 let catalog = session_catalog(RestCatalogConfig::builder().uri(server.url()).build());
4528 let context = SessionContext::empty();
4529 let location = |table: Table| table.metadata_location().unwrap().to_string();
4530
4531 let v1 = "s3://warehouse/v1.metadata.json";
4532 let v2 = "s3://warehouse/v2.metadata.json";
4533 let ident = test_table_ident();
4534 let load = || catalog.load_table(&context, &ident);
4535 assert_eq!(location(load().await.unwrap()), v1);
4536 assert_eq!(location(load().await.unwrap()), v1);
4537
4538 reuse_v1.remove_async().await;
4539 let _v2 = load_table_mock(&mut server, "\"v1\"".into(), 200, Some("\"v2\"")).await;
4540 let reuse_v2 = load_table_mock(&mut server, "\"v2\"".into(), 304, None).await;
4541 assert_eq!(location(load().await.unwrap()), v2);
4542 assert_eq!(location(load().await.unwrap()), v2);
4543 reuse_v2.assert_async().await;
4544 }
4545
4546 #[tokio::test]
4547 async fn test_load_table_not_modified_without_cached_table() {
4548 let mut server = Server::new_async().await;
4549 let _config = create_config_mock(&mut server).await;
4550 let _load = load_table_mock(&mut server, Matcher::Missing, 304, None).await;
4551 let catalog = session_catalog(RestCatalogConfig::builder().uri(server.url()).build());
4552
4553 let err = catalog
4554 .load_table(&SessionContext::empty(), &test_table_ident())
4555 .await
4556 .unwrap_err();
4557
4558 assert_eq!(err.kind(), ErrorKind::Unexpected);
4559 }
4560
4561 #[tokio::test]
4562 async fn test_table_cache_is_per_session() {
4563 let mut server = Server::new_async().await;
4564 let _config = create_config_mock(&mut server).await;
4565 let miss = load_table_mock(&mut server, Matcher::Missing, 200, Some("\"v1\""))
4566 .await
4567 .expect(2);
4568 let hit = load_table_mock(&mut server, "\"v1\"".into(), 304, None).await;
4569 let catalog = session_catalog(RestCatalogConfig::builder().uri(server.url()).build());
4570 let session = |id: &str| SessionContext::builder().session_id(id.to_string()).build();
4571
4572 for id in ["a", "b", "a"] {
4573 catalog
4574 .load_table(&session(id), &test_table_ident())
4575 .await
4576 .unwrap();
4577 }
4578
4579 miss.assert_async().await;
4580 hit.assert_async().await;
4581 }
4582
4583 #[tokio::test]
4584 async fn test_table_cache_disabled_with_zero_max_entries() {
4585 let mut server = Server::new_async().await;
4586 let _config = create_config_mock(&mut server).await;
4587 let load = load_table_mock(&mut server, Matcher::Missing, 200, Some("\"v1\""))
4588 .await
4589 .expect(2);
4590 let catalog = session_catalog(
4591 RestCatalogConfig::builder()
4592 .uri(server.url())
4593 .props(HashMap::from([(
4594 REST_CATALOG_PROP_TABLE_CACHE_MAX_ENTRIES.to_string(),
4595 "0".to_string(),
4596 )]))
4597 .build(),
4598 );
4599 let context = SessionContext::empty();
4600
4601 catalog
4602 .load_table(&context, &test_table_ident())
4603 .await
4604 .unwrap();
4605 catalog
4606 .load_table(&context, &test_table_ident())
4607 .await
4608 .unwrap();
4609
4610 load.assert_async().await;
4611 }
4612
4613 #[tokio::test]
4614 async fn test_table_changes_invalidate_cached_table() {
4615 let mut server = Server::new_async().await;
4616 let _config = create_config_mock(&mut server).await;
4617 let miss = load_table_mock(&mut server, Matcher::Missing, 200, Some("\"v1\""))
4618 .await
4619 .expect(5);
4620 let hit = load_table_mock(&mut server, "\"v1\"".into(), 304, None).await;
4621 let _drop = server
4622 .mock("DELETE", "/v1/namespaces/ns1/tables/test1")
4623 .with_status(204)
4624 .create_async()
4625 .await;
4626 let _rename = server
4627 .mock("POST", "/v1/tables/rename")
4628 .with_status(204)
4629 .create_async()
4630 .await;
4631 let _register = server
4632 .mock("POST", "/v1/namespaces/ns1/register")
4633 .with_status(200)
4634 .with_body(load_table_body("v1").to_string())
4635 .create_async()
4636 .await;
4637 let _commit = server
4638 .mock("POST", "/v1/namespaces/ns1/tables/test1")
4639 .with_status(200)
4640 .with_body_from_file(format!(
4641 "{}/testdata/update_table_response.json",
4642 env!("CARGO_MANIFEST_DIR")
4643 ))
4644 .create_async()
4645 .await;
4646 let catalog = RestCatalog::new(
4647 SessionContext::empty(),
4648 RestCatalogConfig::builder().uri(server.url()).build(),
4649 None,
4650 Some(Arc::new(LocalFsStorageFactory)),
4651 Runtime::current(),
4652 None,
4653 );
4654 let ident = test_table_ident();
4655
4656 catalog.load_table(&ident).await.unwrap();
4659 catalog.drop_table(&ident).await.unwrap();
4660 catalog.load_table(&ident).await.unwrap();
4661 let dest = TableIdent::from_strs(["ns1", "test2"]).unwrap();
4662 catalog.rename_table(&ident, &dest).await.unwrap();
4663 catalog.load_table(&ident).await.unwrap();
4664 let location = "s3://warehouse/v1.metadata.json".to_string();
4665 catalog.register_table(&ident, location).await.unwrap();
4666 let table = catalog.load_table(&ident).await.unwrap();
4667 let tx = Transaction::new(&table);
4668 tx.update_table_properties()
4669 .set("k".to_string(), "v".to_string())
4670 .apply(tx)
4671 .unwrap()
4672 .commit(&catalog)
4673 .await
4674 .unwrap();
4675 catalog.load_table(&ident).await.unwrap();
4676
4677 miss.assert_async().await;
4678 hit.assert_async().await;
4679 }
4680
4681 #[tokio::test]
4682 async fn test_load_table_404_invalidates_cached_table() {
4683 let mut server = Server::new_async().await;
4684 let _config = create_config_mock(&mut server).await;
4685 let miss = load_table_mock(&mut server, Matcher::Missing, 200, Some("\"v1\""))
4686 .await
4687 .expect(2);
4688 let _gone = load_table_mock(&mut server, "\"v1\"".into(), 404, None).await;
4689 let catalog = session_catalog(RestCatalogConfig::builder().uri(server.url()).build());
4690 let context = SessionContext::empty();
4691 let ident = test_table_ident();
4692
4693 catalog.load_table(&context, &ident).await.unwrap();
4694 let err = catalog.load_table(&context, &ident).await.unwrap_err();
4695 assert_eq!(err.kind(), ErrorKind::TableNotFound);
4696 catalog.load_table(&context, &ident).await.unwrap();
4697
4698 miss.assert_async().await;
4699 }
4700
4701 #[tokio::test]
4702 async fn test_load_table_vending_credentials_evicts_cached_table() {
4703 let mut server = Server::new_async().await;
4704 let _config = create_config_mock(&mut server).await;
4705 let miss = load_table_mock(&mut server, Matcher::Missing, 200, Some("\"v1\""))
4706 .await
4707 .expect(2);
4708 let mut body = load_table_body("v2");
4709 body["storage-credentials"] = json!([{
4710 "prefix": "s3://warehouse/",
4711 "config": { "s3.session-token": "token" },
4712 }]);
4713 let _vended = server
4714 .mock("GET", "/v1/namespaces/ns1/tables/test1")
4715 .match_header("if-none-match", "\"v1\"")
4716 .with_status(200)
4717 .with_header("etag", "\"v2\"")
4718 .with_body(body.to_string())
4719 .create_async()
4720 .await;
4721 let catalog = session_catalog(RestCatalogConfig::builder().uri(server.url()).build());
4722 let context = SessionContext::empty();
4723
4724 for _ in 0..3 {
4725 catalog
4726 .load_table(&context, &test_table_ident())
4727 .await
4728 .unwrap();
4729 }
4730
4731 miss.assert_async().await;
4732 }
4733
4734 #[tokio::test]
4735 async fn test_load_table_with_vended_credentials_is_not_cached() {
4736 let mut server = Server::new_async().await;
4737 let _config = create_config_mock(&mut server).await;
4738 let mut body = load_table_body("v1");
4739 body["config"] = json!({ "s3.session-token": "token" });
4740 let load = server
4741 .mock("GET", "/v1/namespaces/ns1/tables/test1")
4742 .match_header("if-none-match", Matcher::Missing)
4743 .with_status(200)
4744 .with_header("etag", "\"v1\"")
4745 .with_body(body.to_string())
4746 .expect(2)
4747 .create_async()
4748 .await;
4749 let catalog = session_catalog(RestCatalogConfig::builder().uri(server.url()).build());
4750 let context = SessionContext::empty();
4751
4752 catalog
4753 .load_table(&context, &test_table_ident())
4754 .await
4755 .unwrap();
4756 catalog
4757 .load_table(&context, &test_table_ident())
4758 .await
4759 .unwrap();
4760
4761 load.assert_async().await;
4762 }
4763
4764 #[test]
4765 fn test_table_cache_rejects_out_of_range_expiry() {
4766 for value in ["0", "31536000000001"] {
4767 let config = RestCatalogConfig::builder()
4768 .uri("http://localhost".to_string())
4769 .props(HashMap::from([(
4770 REST_CATALOG_PROP_TABLE_CACHE_EXPIRE_AFTER_WRITE_MS.to_string(),
4771 value.to_string(),
4772 )]))
4773 .build();
4774
4775 assert_eq!(
4776 config.table_cache().unwrap_err().kind(),
4777 ErrorKind::DataInvalid
4778 );
4779 }
4780 }
4781
4782 #[test]
4783 fn test_is_short_lived_credential() {
4784 for key in [
4785 "S3.Session-Token",
4786 "gcs.oauth2.token",
4787 "adls.sas-token.account",
4788 ] {
4789 assert!(is_short_lived_credential(key), "{key}");
4790 }
4791 for key in ["s3.sse.key", "s3.access-key-id", "encryption.key-id"] {
4792 assert!(!is_short_lived_credential(key), "{key}");
4793 }
4794 }
4795
4796 #[tokio::test]
4797 async fn test_failed_drop_still_invalidates_cached_table() {
4798 let mut server = Server::new_async().await;
4799 let _config = create_config_mock(&mut server).await;
4800 let load = load_table_mock(&mut server, Matcher::Missing, 200, Some("\"v1\""))
4801 .await
4802 .expect(2);
4803 let _drop = server
4805 .mock("DELETE", "/v1/namespaces/ns1/tables/test1")
4806 .with_chunked_body(|_| Err(std::io::Error::other("connection reset")))
4807 .create_async()
4808 .await;
4809 let catalog = session_catalog(RestCatalogConfig::builder().uri(server.url()).build());
4810 let context = SessionContext::empty();
4811 let ident = test_table_ident();
4812
4813 catalog.load_table(&context, &ident).await.unwrap();
4814 catalog.drop_table(&context, &ident).await.unwrap_err();
4815 catalog.load_table(&context, &ident).await.unwrap();
4816
4817 load.assert_async().await;
4818 }
4819
4820 #[tokio::test]
4821 async fn test_load_table_404() {
4822 let mut server = Server::new_async().await;
4823
4824 let config_mock = create_config_mock(&mut server).await;
4825
4826 let rename_table_mock = server
4827 .mock("GET", "/v1/namespaces/ns1/tables/test1")
4828 .with_status(404)
4829 .with_body(r#"
4830{
4831 "error": {
4832 "message": "Table does not exist: ns1.test1 in warehouse 8bcb0838-50fc-472d-9ddb-8feb89ef5f1e",
4833 "type": "NoSuchNamespaceErrorException",
4834 "code": 404
4835 }
4836}
4837 "#)
4838 .create_async()
4839 .await;
4840
4841 let catalog = session_catalog(RestCatalogConfig::builder().uri(server.url()).build());
4842
4843 let table = catalog
4844 .load_table(
4845 &SessionContext::empty(),
4846 &TableIdent::new(NamespaceIdent::new("ns1".to_string()), "test1".to_string()),
4847 )
4848 .await;
4849
4850 assert!(table.is_err());
4851 assert!(table.err().unwrap().message().contains("does not exist"));
4852
4853 config_mock.assert_async().await;
4854 rename_table_mock.assert_async().await;
4855 }
4856
4857 #[tokio::test]
4858 async fn test_create_table() {
4859 let mut server = Server::new_async().await;
4860
4861 let config_mock = create_config_mock(&mut server).await;
4862
4863 let create_table_mock = server
4864 .mock("POST", "/v1/namespaces/ns1/tables")
4865 .with_status(200)
4866 .with_body_from_file(format!(
4867 "{}/testdata/{}",
4868 env!("CARGO_MANIFEST_DIR"),
4869 "create_table_response.json"
4870 ))
4871 .create_async()
4872 .await;
4873
4874 let catalog = session_catalog(RestCatalogConfig::builder().uri(server.url()).build());
4875
4876 let table_creation = TableCreation::builder()
4877 .name("test1".to_string())
4878 .schema(
4879 Schema::builder()
4880 .with_fields(vec![
4881 NestedField::optional(1, "foo", Type::Primitive(PrimitiveType::String))
4882 .into(),
4883 NestedField::required(2, "bar", Type::Primitive(PrimitiveType::Int)).into(),
4884 NestedField::optional(3, "baz", Type::Primitive(PrimitiveType::Boolean))
4885 .into(),
4886 ])
4887 .with_schema_id(1)
4888 .with_identifier_field_ids(vec![2])
4889 .build()
4890 .unwrap(),
4891 )
4892 .properties(HashMap::from([("owner".to_string(), "testx".to_string())]))
4893 .partition_spec(
4894 UnboundPartitionSpec::builder()
4895 .add_partition_field(
4896 UnboundPartitionField::builder()
4897 .source_ids(vec![1])
4898 .name("id")
4899 .transform(Transform::Truncate(3))
4900 .build()
4901 .unwrap(),
4902 )
4903 .unwrap()
4904 .build(),
4905 )
4906 .sort_order(
4907 SortOrder::builder()
4908 .with_sort_field(
4909 SortField::builder()
4910 .source_id(2)
4911 .transform(Transform::Identity)
4912 .direction(SortDirection::Ascending)
4913 .null_order(NullOrder::First)
4914 .build(),
4915 )
4916 .build_unbound()
4917 .unwrap(),
4918 )
4919 .build();
4920
4921 let table = catalog
4922 .create_table(
4923 &SessionContext::empty(),
4924 &NamespaceIdent::from_strs(["ns1"]).unwrap(),
4925 table_creation,
4926 )
4927 .await
4928 .unwrap();
4929
4930 assert_eq!(
4931 &TableIdent::from_strs(vec!["ns1", "test1"]).unwrap(),
4932 table.identifier()
4933 );
4934 assert_eq!(
4935 "s3://warehouse/database/table/metadata.json",
4936 table.metadata_location().unwrap()
4937 );
4938 assert_eq!(FormatVersion::V1, table.metadata().format_version());
4939 assert_eq!("s3://warehouse/database/table", table.metadata().location());
4940 assert_eq!(
4941 uuid!("bf289591-dcc0-4234-ad4f-5c3eed811a29"),
4942 table.metadata().uuid()
4943 );
4944 assert_eq!(
4945 1657810967051,
4946 table
4947 .metadata()
4948 .last_updated_timestamp()
4949 .unwrap()
4950 .timestamp_millis()
4951 );
4952 assert_eq!(
4953 vec![&Arc::new(
4954 Schema::builder()
4955 .with_fields(vec![
4956 NestedField::optional(1, "foo", Type::Primitive(PrimitiveType::String))
4957 .into(),
4958 NestedField::required(2, "bar", Type::Primitive(PrimitiveType::Int)).into(),
4959 NestedField::optional(3, "baz", Type::Primitive(PrimitiveType::Boolean))
4960 .into(),
4961 ])
4962 .with_schema_id(0)
4963 .with_identifier_field_ids(vec![2])
4964 .build()
4965 .unwrap()
4966 )],
4967 table.metadata().schemas_iter().collect::<Vec<_>>()
4968 );
4969 assert_eq!(
4970 &HashMap::from([
4971 (
4972 "write.delete.parquet.compression-codec".to_string(),
4973 "zstd".to_string()
4974 ),
4975 (
4976 "write.metadata.compression-codec".to_string(),
4977 "gzip".to_string()
4978 ),
4979 (
4980 "write.summary.partition-limit".to_string(),
4981 "100".to_string()
4982 ),
4983 (
4984 "write.parquet.compression-codec".to_string(),
4985 "zstd".to_string()
4986 ),
4987 ]),
4988 table.metadata().properties()
4989 );
4990 assert!(table.metadata().current_snapshot().is_none());
4991 assert!(table.metadata().history().is_empty());
4992 assert_eq!(
4993 vec![&Arc::new(SortOrder {
4994 order_id: 0,
4995 fields: vec![],
4996 })],
4997 table.metadata().sort_orders_iter().collect::<Vec<_>>()
4998 );
4999
5000 config_mock.assert_async().await;
5001 create_table_mock.assert_async().await;
5002 }
5003
5004 #[tokio::test]
5005 async fn test_create_table_409() {
5006 let mut server = Server::new_async().await;
5007
5008 let config_mock = create_config_mock(&mut server).await;
5009
5010 let create_table_mock = server
5011 .mock("POST", "/v1/namespaces/ns1/tables")
5012 .with_status(409)
5013 .with_body(r#"
5014{
5015 "error": {
5016 "message": "Table already exists: ns1.test1 in warehouse 8bcb0838-50fc-472d-9ddb-8feb89ef5f1e",
5017 "type": "AlreadyExistsException",
5018 "code": 409
5019 }
5020}
5021 "#)
5022 .create_async()
5023 .await;
5024
5025 let catalog = session_catalog(RestCatalogConfig::builder().uri(server.url()).build());
5026
5027 let table_creation = TableCreation::builder()
5028 .name("test1".to_string())
5029 .schema(
5030 Schema::builder()
5031 .with_fields(vec![
5032 NestedField::optional(1, "foo", Type::Primitive(PrimitiveType::String))
5033 .into(),
5034 NestedField::required(2, "bar", Type::Primitive(PrimitiveType::Int)).into(),
5035 NestedField::optional(3, "baz", Type::Primitive(PrimitiveType::Boolean))
5036 .into(),
5037 ])
5038 .with_schema_id(1)
5039 .with_identifier_field_ids(vec![2])
5040 .build()
5041 .unwrap(),
5042 )
5043 .properties(HashMap::from([("owner".to_string(), "testx".to_string())]))
5044 .build();
5045
5046 let table_result = catalog
5047 .create_table(
5048 &SessionContext::empty(),
5049 &NamespaceIdent::from_strs(["ns1"]).unwrap(),
5050 table_creation,
5051 )
5052 .await;
5053
5054 assert!(table_result.is_err());
5055 assert!(
5056 table_result
5057 .err()
5058 .unwrap()
5059 .message()
5060 .contains("already exists")
5061 );
5062
5063 config_mock.assert_async().await;
5064 create_table_mock.assert_async().await;
5065 }
5066
5067 #[tokio::test]
5068 async fn test_update_table() {
5069 let mut server = Server::new_async().await;
5070
5071 let config_mock = create_config_mock(&mut server).await;
5072
5073 let load_table_mock = server
5074 .mock("GET", "/v1/namespaces/ns1/tables/test1")
5075 .with_status(200)
5076 .with_body_from_file(format!(
5077 "{}/testdata/{}",
5078 env!("CARGO_MANIFEST_DIR"),
5079 "load_table_response.json"
5080 ))
5081 .create_async()
5082 .await;
5083
5084 let update_table_mock = server
5085 .mock("POST", "/v1/namespaces/ns1/tables/test1")
5086 .with_status(200)
5087 .with_body_from_file(format!(
5088 "{}/testdata/{}",
5089 env!("CARGO_MANIFEST_DIR"),
5090 "update_table_response.json"
5091 ))
5092 .create_async()
5093 .await;
5094
5095 let catalog = RestCatalog::new(
5096 SessionContext::empty(),
5097 RestCatalogConfig::builder().uri(server.url()).build(),
5098 None,
5099 Some(Arc::new(LocalFsStorageFactory)),
5100 Runtime::current(),
5101 None,
5102 );
5103
5104 let table1 = {
5105 let file = File::open(format!(
5106 "{}/testdata/{}",
5107 env!("CARGO_MANIFEST_DIR"),
5108 "create_table_response.json"
5109 ))
5110 .unwrap();
5111 let reader = BufReader::new(file);
5112 let resp = serde_json::from_reader::<_, LoadTableResult>(reader).unwrap();
5113
5114 Table::builder()
5115 .metadata(resp.metadata)
5116 .metadata_location(resp.metadata_location.unwrap())
5117 .identifier(TableIdent::from_strs(["ns1", "test1"]).unwrap())
5118 .file_io(FileIO::new_with_fs())
5119 .runtime(test_runtime())
5120 .build()
5121 .unwrap()
5122 };
5123
5124 let tx = Transaction::new(&table1);
5125 let table = tx
5126 .upgrade_table_version()
5127 .set_format_version(FormatVersion::V2)
5128 .apply(tx)
5129 .unwrap()
5130 .commit(&catalog)
5131 .await
5132 .unwrap();
5133
5134 assert_eq!(
5135 &TableIdent::from_strs(vec!["ns1", "test1"]).unwrap(),
5136 table.identifier()
5137 );
5138 assert_eq!(
5139 "s3://warehouse/database/table/metadata.json",
5140 table.metadata_location().unwrap()
5141 );
5142 assert_eq!(FormatVersion::V2, table.metadata().format_version());
5143 assert_eq!("s3://warehouse/database/table", table.metadata().location());
5144 assert_eq!(
5145 uuid!("bf289591-dcc0-4234-ad4f-5c3eed811a29"),
5146 table.metadata().uuid()
5147 );
5148 assert_eq!(
5149 1657810967051,
5150 table
5151 .metadata()
5152 .last_updated_timestamp()
5153 .unwrap()
5154 .timestamp_millis()
5155 );
5156 assert_eq!(
5157 vec![&Arc::new(
5158 Schema::builder()
5159 .with_fields(vec![
5160 NestedField::optional(1, "foo", Type::Primitive(PrimitiveType::String))
5161 .into(),
5162 NestedField::required(2, "bar", Type::Primitive(PrimitiveType::Int)).into(),
5163 NestedField::optional(3, "baz", Type::Primitive(PrimitiveType::Boolean))
5164 .into(),
5165 ])
5166 .with_schema_id(0)
5167 .with_identifier_field_ids(vec![2])
5168 .build()
5169 .unwrap()
5170 )],
5171 table.metadata().schemas_iter().collect::<Vec<_>>()
5172 );
5173 assert_eq!(
5174 &HashMap::from([
5175 (
5176 "write.delete.parquet.compression-codec".to_string(),
5177 "zstd".to_string()
5178 ),
5179 (
5180 "write.metadata.compression-codec".to_string(),
5181 "gzip".to_string()
5182 ),
5183 (
5184 "write.summary.partition-limit".to_string(),
5185 "100".to_string()
5186 ),
5187 (
5188 "write.parquet.compression-codec".to_string(),
5189 "zstd".to_string()
5190 ),
5191 ]),
5192 table.metadata().properties()
5193 );
5194 assert!(table.metadata().current_snapshot().is_none());
5195 assert!(table.metadata().history().is_empty());
5196 assert_eq!(
5197 vec![&Arc::new(SortOrder {
5198 order_id: 0,
5199 fields: vec![],
5200 })],
5201 table.metadata().sort_orders_iter().collect::<Vec<_>>()
5202 );
5203
5204 config_mock.assert_async().await;
5205 update_table_mock.assert_async().await;
5206 load_table_mock.assert_async().await
5207 }
5208
5209 #[tokio::test]
5210 async fn test_update_table_404() {
5211 let mut server = Server::new_async().await;
5212
5213 let config_mock = create_config_mock(&mut server).await;
5214
5215 let load_table_mock = server
5216 .mock("GET", "/v1/namespaces/ns1/tables/test1")
5217 .with_status(200)
5218 .with_body_from_file(format!(
5219 "{}/testdata/{}",
5220 env!("CARGO_MANIFEST_DIR"),
5221 "load_table_response.json"
5222 ))
5223 .create_async()
5224 .await;
5225
5226 let update_table_mock = server
5227 .mock("POST", "/v1/namespaces/ns1/tables/test1")
5228 .with_status(404)
5229 .with_body(
5230 r#"
5231{
5232 "error": {
5233 "message": "The given table does not exist",
5234 "type": "NoSuchTableException",
5235 "code": 404
5236 }
5237}
5238 "#,
5239 )
5240 .create_async()
5241 .await;
5242
5243 let catalog = RestCatalog::new(
5244 SessionContext::empty(),
5245 RestCatalogConfig::builder().uri(server.url()).build(),
5246 None,
5247 Some(Arc::new(LocalFsStorageFactory)),
5248 Runtime::current(),
5249 None,
5250 );
5251
5252 let table1 = {
5253 let file = File::open(format!(
5254 "{}/testdata/{}",
5255 env!("CARGO_MANIFEST_DIR"),
5256 "create_table_response.json"
5257 ))
5258 .unwrap();
5259 let reader = BufReader::new(file);
5260 let resp = serde_json::from_reader::<_, LoadTableResult>(reader).unwrap();
5261
5262 Table::builder()
5263 .metadata(resp.metadata)
5264 .metadata_location(resp.metadata_location.unwrap())
5265 .identifier(TableIdent::from_strs(["ns1", "test1"]).unwrap())
5266 .file_io(FileIO::new_with_fs())
5267 .runtime(test_runtime())
5268 .build()
5269 .unwrap()
5270 };
5271
5272 let tx = Transaction::new(&table1);
5273 let table_result = tx
5274 .upgrade_table_version()
5275 .set_format_version(FormatVersion::V2)
5276 .apply(tx)
5277 .unwrap()
5278 .commit(&catalog)
5279 .await;
5280
5281 assert!(table_result.is_err());
5282 assert!(
5283 table_result
5284 .err()
5285 .unwrap()
5286 .message()
5287 .contains("does not exist")
5288 );
5289
5290 config_mock.assert_async().await;
5291 update_table_mock.assert_async().await;
5292 load_table_mock.assert_async().await;
5293 }
5294
5295 #[tokio::test]
5296 async fn test_register_table() {
5297 let mut server = Server::new_async().await;
5298
5299 let config_mock = create_config_mock(&mut server).await;
5300
5301 let register_table_mock = server
5302 .mock("POST", "/v1/namespaces/ns1/register")
5303 .with_status(200)
5304 .with_body_from_file(format!(
5305 "{}/testdata/{}",
5306 env!("CARGO_MANIFEST_DIR"),
5307 "load_table_response.json"
5308 ))
5309 .create_async()
5310 .await;
5311
5312 let catalog = session_catalog(RestCatalogConfig::builder().uri(server.url()).build());
5313 let table_ident =
5314 TableIdent::new(NamespaceIdent::new("ns1".to_string()), "test1".to_string());
5315 let metadata_location = String::from(
5316 "s3://warehouse/database/table/metadata/00001-5f2f8166-244c-4eae-ac36-384ecdec81fc.gz.metadata.json",
5317 );
5318
5319 let table = catalog
5320 .register_table(&SessionContext::empty(), &table_ident, metadata_location)
5321 .await
5322 .unwrap();
5323
5324 assert_eq!(
5325 &TableIdent::from_strs(vec!["ns1", "test1"]).unwrap(),
5326 table.identifier()
5327 );
5328 assert_eq!(
5329 "s3://warehouse/database/table/metadata/00001-5f2f8166-244c-4eae-ac36-384ecdec81fc.gz.metadata.json",
5330 table.metadata_location().unwrap()
5331 );
5332
5333 config_mock.assert_async().await;
5334 register_table_mock.assert_async().await;
5335 }
5336
5337 #[tokio::test]
5338 async fn test_register_table_404() {
5339 let mut server = Server::new_async().await;
5340
5341 let config_mock = create_config_mock(&mut server).await;
5342
5343 let register_table_mock = server
5344 .mock("POST", "/v1/namespaces/ns1/register")
5345 .with_status(404)
5346 .with_body(
5347 r#"
5348{
5349 "error": {
5350 "message": "The namespace specified does not exist",
5351 "type": "NoSuchNamespaceErrorException",
5352 "code": 404
5353 }
5354}
5355 "#,
5356 )
5357 .create_async()
5358 .await;
5359
5360 let catalog = session_catalog(RestCatalogConfig::builder().uri(server.url()).build());
5361
5362 let table_ident =
5363 TableIdent::new(NamespaceIdent::new("ns1".to_string()), "test1".to_string());
5364 let metadata_location = String::from(
5365 "s3://warehouse/database/table/metadata/00001-5f2f8166-244c-4eae-ac36-384ecdec81fc.gz.metadata.json",
5366 );
5367 let table = catalog
5368 .register_table(&SessionContext::empty(), &table_ident, metadata_location)
5369 .await;
5370
5371 assert!(table.is_err());
5372 assert!(table.err().unwrap().message().contains("does not exist"));
5373
5374 config_mock.assert_async().await;
5375 register_table_mock.assert_async().await;
5376 }
5377
5378 #[tokio::test]
5379 async fn test_create_rest_catalog() {
5380 let builder = RestCatalogBuilder::default().with_client(Client::new());
5381
5382 let catalog = builder
5383 .load(
5384 "test",
5385 HashMap::from([
5386 (
5387 REST_CATALOG_PROP_URI.to_string(),
5388 "http://localhost:8080".to_string(),
5389 ),
5390 ("a".to_string(), "b".to_string()),
5391 ]),
5392 )
5393 .await;
5394
5395 assert!(catalog.is_ok());
5396
5397 let catalog = catalog.unwrap();
5398 let catalog_config = &catalog.inner.user_config;
5399 assert_eq!(catalog_config.name.as_deref(), Some("test"));
5400 assert_eq!(catalog_config.uri, "http://localhost:8080");
5401 assert_eq!(catalog_config.warehouse, None);
5402 assert!(catalog_config.client.is_some());
5403
5404 assert_eq!(catalog_config.props.get("a"), Some(&"b".to_string()));
5405 assert!(!catalog_config.props.contains_key(REST_CATALOG_PROP_URI));
5406 }
5407
5408 #[tokio::test]
5409 async fn test_create_rest_catalog_no_uri() {
5410 let builder = RestCatalogBuilder::default();
5411
5412 let catalog = builder
5413 .load(
5414 "test",
5415 HashMap::from([(
5416 REST_CATALOG_PROP_WAREHOUSE.to_string(),
5417 "s3://warehouse".to_string(),
5418 )]),
5419 )
5420 .await;
5421
5422 assert!(catalog.is_err());
5423 if let Err(err) = catalog {
5424 assert_eq!(err.kind(), ErrorKind::DataInvalid);
5425 assert_eq!(err.message(), "Catalog uri is required");
5426 }
5427 }
5428
5429 #[tokio::test]
5430 async fn test_create_session_catalog() {
5431 let builder = RestSessionCatalogBuilder::default();
5432
5433 let result = builder
5434 .load(
5435 "test",
5436 HashMap::from([
5437 (
5438 REST_CATALOG_PROP_URI.to_string(),
5439 "http://localhost:8080".to_string(),
5440 ),
5441 ("a".to_string(), "b".to_string()),
5442 ]),
5443 )
5444 .await;
5445
5446 assert!(result.is_ok());
5447
5448 let catalog = result.unwrap();
5449
5450 let catalog_config = catalog.user_config;
5451 assert_eq!(catalog_config.name.as_deref(), Some("test"));
5452 assert_eq!(catalog_config.uri, "http://localhost:8080");
5453 assert_eq!(catalog_config.warehouse, None);
5454 assert!(catalog_config.client.is_none());
5456
5457 assert_eq!(catalog_config.props.get("a"), Some(&"b".to_string()));
5459 assert!(!catalog_config.props.contains_key(REST_CATALOG_PROP_URI));
5460 }
5461
5462 #[tokio::test]
5463 async fn test_create_rest_catalog_with_session() {
5464 let context = SessionContext::builder()
5465 .session_id("test-id".to_string())
5466 .build();
5467
5468 let result = RestCatalogBuilder::default()
5469 .with_session_context(context)
5470 .load(
5471 "test",
5472 HashMap::from([(
5473 REST_CATALOG_PROP_URI.to_string(),
5474 "http://localhost:8080".to_string(),
5475 )]),
5476 )
5477 .await;
5478
5479 assert!(result.is_ok());
5480
5481 let catalog = result.unwrap();
5483 assert_eq!(catalog.session_context.session_id(), "test-id");
5484 }
5485
5486 #[tokio::test]
5487 async fn test_create_rest_catalog_default_session() {
5488 let result = RestCatalogBuilder::default()
5489 .load(
5490 "test",
5491 HashMap::from([(
5492 REST_CATALOG_PROP_URI.to_string(),
5493 "http://localhost:8080".to_string(),
5494 )]),
5495 )
5496 .await;
5497
5498 assert!(result.is_ok());
5499
5500 let catalog = result.unwrap();
5503 assert!(uuid::Uuid::parse_str(catalog.session_context.session_id()).is_ok());
5504 }
5505
5506 #[tokio::test]
5508 async fn test_rest_catalog_delegates() {
5509 let mut server = Server::new_async().await;
5510 let config_mock = create_config_mock(&mut server).await;
5511 let list_ns_mock = server
5512 .mock("GET", "/v1/namespaces")
5513 .with_body(
5514 r#"{
5515 "namespaces": [["ns1"]]
5516 }"#,
5517 )
5518 .create_async()
5519 .await;
5520
5521 let catalog = RestCatalog::new(
5522 SessionContext::empty(),
5523 RestCatalogConfig::builder().uri(server.url()).build(),
5524 None,
5525 Some(Arc::new(LocalFsStorageFactory)),
5526 Runtime::current(),
5527 None,
5528 );
5529
5530 let namespaces = catalog.list_namespaces(None).await.unwrap();
5531
5532 assert_eq!(namespaces, vec![
5533 NamespaceIdent::from_vec(vec!["ns1".to_string()]).unwrap()
5534 ]);
5535
5536 config_mock.assert_async().await;
5537 list_ns_mock.assert_async().await;
5538 }
5539}