1use std::collections::{HashMap, HashSet};
21use std::fmt::{Debug, Formatter};
22use std::future::Future;
23use std::str::FromStr;
24use std::sync::{Arc, OnceLock};
25
26use async_trait::async_trait;
27use iceberg::encryption::kms::{KeyManagementClient, KmsClientFactory};
28use iceberg::io::{FileIO, FileIOBuilder, StorageFactory};
29use iceberg::table::Table;
30use iceberg::{
31 Catalog, CatalogBuilder, Error, ErrorKind, Namespace, NamespaceIdent, Result, Runtime,
32 SessionCatalog, SessionContext, TableCommit, TableCreation, TableIdent,
33};
34use itertools::Itertools;
35use reqwest::header::{
36 HeaderMap, HeaderName, HeaderValue, {self},
37};
38use reqwest::{Client, Method, StatusCode, Url};
39use tokio::sync::OnceCell;
40use typed_builder::TypedBuilder;
41
42use crate::auth::{
43 AUTH_TYPE_NONE, AUTH_TYPE_OAUTH2, AuthManager, AuthSession, NoopAuthManager, OAuth2Manager,
44};
45use crate::client::{
46 HttpClient, deserialize_catalog_response, deserialize_unexpected_catalog_error,
47};
48use crate::endpoint::{Endpoint, V1_NAMESPACE_EXISTS, V1_TABLE_EXISTS};
49use crate::request::HttpRequest;
50use crate::response::HttpResponse;
51use crate::types::{
52 CatalogConfig, CommitTableRequest, CommitTableResponse, CreateNamespaceRequest,
53 CreateTableRequest, ListNamespaceResponse, ListTablesResponse, LoadTableResult,
54 NamespaceResponse, RegisterTableRequest, RenameTableRequest,
55};
56
57pub const REST_CATALOG_PROP_URI: &str = "uri";
59pub const REST_CATALOG_PROP_WAREHOUSE: &str = "warehouse";
61pub const REST_CATALOG_PROP_DISABLE_HEADER_REDACTION: &str = "disable-header-redaction";
64pub const REST_CATALOG_PROP_AUTH_TYPE: &str = "rest.auth.type";
68
69const ICEBERG_REST_SPEC_VERSION: &str = "0.14.1";
70const CARGO_PKG_VERSION: &str = env!("CARGO_PKG_VERSION");
71const PATH_V1: &str = "v1";
72
73#[derive(Debug, Default)]
79pub struct RestCatalogBuilder {
80 session_context: Option<SessionContext>,
81 inner: RestSessionCatalogBuilder,
82}
83
84impl CatalogBuilder for RestCatalogBuilder {
85 type C = RestCatalog;
86
87 fn with_storage_factory(mut self, storage_factory: Arc<dyn StorageFactory>) -> Self {
88 self.inner = self.inner.with_storage_factory(storage_factory);
89 self
90 }
91
92 fn with_kms_client_factory(mut self, kms_client_factory: Arc<dyn KmsClientFactory>) -> Self {
93 self.inner = self.inner.with_kms_client_factory(kms_client_factory);
94 self
95 }
96
97 fn with_runtime(mut self, runtime: Runtime) -> Self {
98 self.inner = self.inner.with_runtime(runtime);
99 self
100 }
101
102 fn load(
103 self,
104 name: impl Into<String>,
105 props: HashMap<String, String>,
106 ) -> impl Future<Output = Result<Self::C>> + Send {
107 let name = name.into();
108 async move {
109 let context = self.session_context.unwrap_or_else(SessionContext::empty);
110 let session_catalog = Arc::new(self.inner.load(name, props).await?);
111
112 Ok(RestCatalog::from_session_catalog(context, session_catalog))
113 }
114 }
115}
116
117impl RestCatalogBuilder {
118 pub fn with_client(mut self, client: Client) -> Self {
120 self.inner = self.inner.with_client(client);
121 self
122 }
123
124 pub fn with_session_context(mut self, context: SessionContext) -> Self {
129 self.session_context = Some(context);
130 self
131 }
132
133 pub fn with_auth_manager<M>(mut self, auth_manager: M) -> Self
138 where M: AuthManager + 'static {
139 self.inner = self.inner.with_auth_manager(auth_manager);
140 self
141 }
142}
143
144#[derive(Clone, TypedBuilder)]
146pub(crate) struct RestCatalogConfig {
147 #[builder(default, setter(strip_option))]
148 name: Option<String>,
149
150 uri: String,
151
152 #[builder(default, setter(strip_option(fallback = warehouse_opt)))]
153 warehouse: Option<String>,
154
155 #[builder(default)]
156 props: HashMap<String, String>,
157
158 #[builder(default)]
159 client: Option<Client>,
160
161 #[builder(default)]
165 default_client: Arc<OnceLock<Client>>,
166}
167
168fn is_sensitive_prop(key: &str) -> bool {
171 key.contains("token")
172 || key.contains("credential")
173 || key.contains("secret")
174 || key.contains("password")
175 || key.contains("key")
176 || key.contains("connection-string")
177 || key.starts_with("header.")
178}
179
180impl Debug for RestCatalogConfig {
183 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
184 let props: HashMap<&str, &str> = self
185 .props
186 .iter()
187 .map(|(key, value)| {
188 let value = if is_sensitive_prop(key) {
189 "[REDACTED]"
190 } else {
191 value.as_str()
192 };
193 (key.as_str(), value)
194 })
195 .collect();
196 f.debug_struct("RestCatalogConfig")
197 .field("name", &self.name)
198 .field("uri", &self.uri)
199 .field("warehouse", &self.warehouse)
200 .field("props", &props)
201 .finish_non_exhaustive()
202 }
203}
204
205impl RestCatalogConfig {
206 fn url_prefixed(&self, parts: &[&str]) -> String {
207 [&self.uri, PATH_V1]
208 .into_iter()
209 .chain(self.props.get("prefix").map(|s| &**s))
210 .chain(parts.iter().cloned())
211 .join("/")
212 }
213
214 fn config_endpoint(&self) -> String {
215 [&self.uri, PATH_V1, "config"].join("/")
216 }
217
218 pub(crate) fn get_token_endpoint(&self) -> String {
219 self.explicit_oauth2_server_uri()
220 .unwrap_or_else(|| default_token_endpoint(&self.uri))
221 }
222
223 pub(crate) fn explicit_oauth2_server_uri(&self) -> Option<String> {
225 self.props.get("oauth2-server-uri").cloned()
226 }
227
228 fn namespaces_endpoint(&self) -> String {
229 self.url_prefixed(&["namespaces"])
230 }
231
232 fn namespace_endpoint(&self, ns: &NamespaceIdent) -> String {
233 self.url_prefixed(&["namespaces", &ns.to_url_string()])
234 }
235
236 fn tables_endpoint(&self, ns: &NamespaceIdent) -> String {
237 self.url_prefixed(&["namespaces", &ns.to_url_string(), "tables"])
238 }
239
240 fn rename_table_endpoint(&self) -> String {
241 self.url_prefixed(&["tables", "rename"])
242 }
243
244 fn register_table_endpoint(&self, ns: &NamespaceIdent) -> String {
245 self.url_prefixed(&["namespaces", &ns.to_url_string(), "register"])
246 }
247
248 fn table_endpoint(&self, table: &TableIdent) -> String {
249 self.url_prefixed(&[
250 "namespaces",
251 &table.namespace.to_url_string(),
252 "tables",
253 &table.name,
254 ])
255 }
256
257 pub(crate) fn client(&self) -> Client {
261 self.client
262 .clone()
263 .unwrap_or_else(|| self.default_client.get_or_init(Client::default).clone())
264 }
265
266 pub(crate) fn token(&self) -> Option<String> {
270 self.props.get("token").cloned()
271 }
272
273 pub(crate) fn credential(&self) -> Option<(Option<String>, String)> {
276 credential_from_props(&self.props)
277 }
278
279 pub(crate) fn extra_headers(&self) -> Result<HeaderMap> {
281 extra_headers_from_props(&self.props)
282 }
283
284 pub(crate) fn extra_oauth_params(&self) -> HashMap<String, String> {
286 oauth_params_from_props(&self.props)
287 }
288
289 pub(crate) fn disable_header_redaction(&self) -> bool {
294 self.props
295 .get(REST_CATALOG_PROP_DISABLE_HEADER_REDACTION)
296 .map(|v| v.eq_ignore_ascii_case("true"))
297 .unwrap_or(false)
298 }
299
300 pub(crate) fn merge_with_config(mut self, mut config: CatalogConfig) -> Self {
302 if let Some(uri) = config.overrides.remove(REST_CATALOG_PROP_URI) {
303 self.uri = uri;
304 }
305
306 let mut props = config.defaults;
307 props.extend(self.props);
308 if let Some(warehouse) = &self.warehouse {
311 props.insert(REST_CATALOG_PROP_WAREHOUSE.to_string(), warehouse.clone());
312 }
313 props.extend(config.overrides);
314
315 self.props = props;
316 self
317 }
318}
319
320pub(crate) fn credential_from_props(
328 props: &HashMap<String, String>,
329) -> Option<(Option<String>, String)> {
330 let cred = props.get("credential")?;
331
332 match cred.split_once(':') {
333 Some((client_id, client_secret)) => {
334 Some((Some(client_id.to_string()), client_secret.to_string()))
335 }
336 None => Some((None, cred.to_string())),
337 }
338}
339
340pub(crate) fn extra_headers_from_props(props: &HashMap<String, String>) -> Result<HeaderMap> {
347 let mut headers = HeaderMap::from_iter([
348 (
349 header::CONTENT_TYPE,
350 HeaderValue::from_static("application/json"),
351 ),
352 (
353 HeaderName::from_static("x-client-version"),
354 HeaderValue::from_static(ICEBERG_REST_SPEC_VERSION),
355 ),
356 (
357 header::USER_AGENT,
358 HeaderValue::from_str(&format!("iceberg-rs/{CARGO_PKG_VERSION}")).unwrap(),
359 ),
360 ]);
361
362 headers.extend(explicit_headers_from_props(props)?);
363
364 Ok(headers)
365}
366
367pub(crate) fn default_token_endpoint(uri: &str) -> String {
369 [uri, PATH_V1, "oauth", "tokens"].join("/")
370}
371
372pub(crate) fn explicit_headers_from_props(props: &HashMap<String, String>) -> Result<HeaderMap> {
374 let mut headers = HeaderMap::new();
375 for (key, value) in props
376 .iter()
377 .filter_map(|(k, v)| k.strip_prefix("header.").map(|k| (k, v)))
378 {
379 headers.insert(
380 HeaderName::from_str(key).map_err(|e| {
381 Error::new(
382 ErrorKind::DataInvalid,
383 format!("Invalid header name: {key}"),
384 )
385 .with_source(e)
386 })?,
387 HeaderValue::from_str(value).map_err(|e| {
388 Error::new(
389 ErrorKind::DataInvalid,
390 format!("Invalid value for header: {key}"),
392 )
393 .with_source(e)
394 })?,
395 );
396 }
397
398 Ok(headers)
399}
400
401pub(crate) fn oauth_params_from_props(props: &HashMap<String, String>) -> HashMap<String, String> {
403 let mut params = HashMap::new();
404
405 if let Some(scope) = props.get("scope") {
406 params.insert("scope".to_string(), scope.to_string());
407 } else {
408 params.insert("scope".to_string(), "catalog".to_string());
409 }
410
411 let optional_params = ["audience", "resource"];
412 for param_name in optional_params {
413 if let Some(value) = props.get(param_name) {
414 params.insert(param_name.to_string(), value.to_string());
415 }
416 }
417
418 params
419}
420
421struct RestCatalogClient {
422 auth_manager: Arc<dyn AuthManager>,
425 catalog_session: Arc<dyn AuthSession>,
427 http_client: HttpClient,
430 config: RestCatalogConfig,
434 endpoints: HashSet<Endpoint>,
436}
437
438impl Debug for RestCatalogClient {
439 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
440 f.debug_struct("RestCatalogClient")
443 .field("http_client", &self.http_client)
444 .field("config", &self.config)
445 .field("endpoints", &self.endpoints)
446 .finish_non_exhaustive()
447 }
448}
449
450impl RestCatalogClient {
451 async fn init(
454 user_config: &RestCatalogConfig,
455 auth_manager: Arc<dyn AuthManager>,
456 ) -> Result<Self> {
457 let http_client = HttpClient::new(user_config)?;
458 let catalog_config = {
462 let init_session = auth_manager
463 .init_session(&http_client, &Self::auth_props(user_config))
464 .await?;
465 Self::load_config(&http_client, init_session.as_ref(), user_config).await?
466 };
467 let endpoints = match &catalog_config.endpoints {
470 Some(advertised) if !advertised.is_empty() => advertised.iter().cloned().collect(),
471 _ => crate::endpoint::DEFAULT_ENDPOINTS.clone(),
472 };
473 let config = user_config.clone().merge_with_config(catalog_config);
474 let http_client = http_client.update_with(&config)?;
475 let catalog_session = auth_manager
478 .catalog_session(&http_client, &Self::auth_props(&config))
479 .await?;
480
481 Ok(Self {
482 auth_manager,
483 catalog_session,
484 config,
485 http_client,
486 endpoints,
487 })
488 }
489
490 #[cfg(test)]
492 async fn token(&self) -> Option<String> {
493 self.http_client.token(self.catalog_session.as_ref()).await
494 }
495
496 async fn query_catalog(
498 &self,
499 context: &SessionContext,
500 request: HttpRequest,
501 ) -> Result<HttpResponse> {
502 let session = self
503 .auth_manager
504 .contextual_session(context, Arc::clone(&self.catalog_session))
505 .await?;
506 self.http_client
507 .query_catalog(session.as_ref(), request)
508 .await
509 }
510
511 fn auth_props(config: &RestCatalogConfig) -> HashMap<String, String> {
514 let mut props = config.props.clone();
519 props.insert(REST_CATALOG_PROP_URI.to_string(), config.uri.clone());
520 if let Some(warehouse) = &config.warehouse {
521 props
524 .entry(REST_CATALOG_PROP_WAREHOUSE.to_string())
525 .or_insert_with(|| warehouse.clone());
526 }
527 props
528 }
529
530 async fn load_config(
534 http_client: &HttpClient,
535 auth_session: &dyn AuthSession,
536 user_config: &RestCatalogConfig,
537 ) -> Result<CatalogConfig> {
538 let mut request_builder = http_client.request(Method::GET, user_config.config_endpoint());
539
540 if let Some(warehouse_location) = &user_config.warehouse {
541 request_builder = request_builder.query(&[("warehouse", warehouse_location)]);
542 }
543
544 let request = HttpRequest::build(request_builder)?;
545
546 let http_response = http_client.query_catalog(auth_session, request).await?;
547
548 match http_response.status() {
549 StatusCode::OK => deserialize_catalog_response(http_response),
550 _ => Err(deserialize_unexpected_catalog_error(
551 http_response,
552 http_client.disable_header_redaction(),
553 )),
554 }
555 }
556}
557
558#[derive(Debug)]
564pub struct RestCatalog {
565 session_context: SessionContext,
566 inner: Arc<RestSessionCatalog>,
567}
568
569impl RestCatalog {
570 #[cfg(test)]
572 fn new(
573 context: SessionContext,
574 config: RestCatalogConfig,
575 auth_manager: Option<Box<dyn AuthManager>>,
576 storage_factory: Option<Arc<dyn StorageFactory>>,
577 runtime: Runtime,
578 kms_client: Option<Arc<dyn KeyManagementClient>>,
579 ) -> Self {
580 let session_catalog = Arc::new(RestSessionCatalog::new(
581 config,
582 auth_manager,
583 storage_factory,
584 runtime,
585 kms_client,
586 ));
587
588 Self::from_session_catalog(context, session_catalog)
589 }
590
591 fn from_session_catalog(context: SessionContext, inner: Arc<RestSessionCatalog>) -> Self {
592 Self {
593 session_context: context,
594 inner,
595 }
596 }
597
598 #[cfg(test)]
599 async fn client(&self) -> Result<&RestCatalogClient> {
600 self.inner.client().await
601 }
602}
603
604#[async_trait]
608impl Catalog for RestCatalog {
609 async fn list_namespaces(
610 &self,
611 parent: Option<&NamespaceIdent>,
612 ) -> Result<Vec<NamespaceIdent>> {
613 self.inner
614 .list_namespaces(&self.session_context, parent)
615 .await
616 }
617
618 async fn create_namespace(
619 &self,
620 namespace: &NamespaceIdent,
621 properties: HashMap<String, String>,
622 ) -> Result<Namespace> {
623 self.inner
624 .create_namespace(&self.session_context, namespace, properties)
625 .await
626 }
627
628 async fn get_namespace(&self, namespace: &NamespaceIdent) -> Result<Namespace> {
629 self.inner
630 .get_namespace(&self.session_context, namespace)
631 .await
632 }
633
634 async fn namespace_exists(&self, ns: &NamespaceIdent) -> Result<bool> {
635 self.inner.namespace_exists(&self.session_context, ns).await
636 }
637
638 async fn update_namespace(
639 &self,
640 namespace: &NamespaceIdent,
641 properties: HashMap<String, String>,
642 ) -> Result<()> {
643 self.inner
644 .update_namespace(&self.session_context, namespace, properties)
645 .await
646 }
647
648 async fn drop_namespace(&self, namespace: &NamespaceIdent) -> Result<()> {
649 self.inner
650 .drop_namespace(&self.session_context, namespace)
651 .await
652 }
653
654 async fn list_tables(&self, namespace: &NamespaceIdent) -> Result<Vec<TableIdent>> {
655 self.inner
656 .list_tables(&self.session_context, namespace)
657 .await
658 }
659
660 async fn create_table(
661 &self,
662 namespace: &NamespaceIdent,
663 creation: TableCreation,
664 ) -> Result<Table> {
665 self.inner
666 .create_table(&self.session_context, namespace, creation)
667 .await
668 }
669
670 async fn load_table(&self, table_ident: &TableIdent) -> Result<Table> {
671 self.inner
672 .load_table(&self.session_context, table_ident)
673 .await
674 }
675
676 async fn drop_table(&self, table: &TableIdent) -> Result<()> {
677 self.inner.drop_table(&self.session_context, table).await
678 }
679
680 async fn purge_table(&self, table: &TableIdent) -> Result<()> {
681 self.inner.purge_table(&self.session_context, table).await
682 }
683
684 async fn table_exists(&self, table: &TableIdent) -> Result<bool> {
685 self.inner.table_exists(&self.session_context, table).await
686 }
687
688 async fn rename_table(&self, src: &TableIdent, dest: &TableIdent) -> Result<()> {
689 self.inner
690 .rename_table(&self.session_context, src, dest)
691 .await
692 }
693
694 async fn register_table(
695 &self,
696 table_ident: &TableIdent,
697 metadata_location: String,
698 ) -> Result<Table> {
699 self.inner
700 .register_table(&self.session_context, table_ident, metadata_location)
701 .await
702 }
703
704 async fn update_table(&self, commit: TableCommit) -> Result<Table> {
705 self.inner.update_table(&self.session_context, commit).await
706 }
707}
708
709#[derive(Debug)]
715pub struct RestSessionCatalog {
716 auth_manager_override: Option<Arc<dyn AuthManager>>,
720 user_config: RestCatalogConfig,
724 client: OnceCell<RestCatalogClient>,
725 storage_factory: Option<Arc<dyn StorageFactory>>,
727 runtime: Runtime,
728 kms_client: Option<Arc<dyn KeyManagementClient>>,
730}
731
732impl RestSessionCatalog {
733 fn new(
735 config: RestCatalogConfig,
736 auth_manager_override: Option<Box<dyn AuthManager>>,
737 storage_factory: Option<Arc<dyn StorageFactory>>,
738 runtime: Runtime,
739 kms_client: Option<Arc<dyn KeyManagementClient>>,
740 ) -> Self {
741 Self {
742 auth_manager_override: auth_manager_override.map(Arc::from),
743 user_config: config,
744 client: OnceCell::new(),
745 storage_factory,
746 runtime,
747 kms_client,
748 }
749 }
750
751 async fn delete_table(
753 &self,
754 context: &SessionContext,
755 table: &TableIdent,
756 purge: bool,
757 ) -> Result<()> {
758 let client = self.client().await?;
759
760 let mut request_builder = client
761 .http_client
762 .request(Method::DELETE, client.config.table_endpoint(table));
763
764 if purge {
765 request_builder = request_builder.query(&[("purgeRequested", "true")]);
766 }
767
768 let request = HttpRequest::build(request_builder)?;
769 let http_response = client.query_catalog(context, request).await?;
770
771 match http_response.status() {
772 StatusCode::NO_CONTENT | StatusCode::OK => Ok(()),
773 StatusCode::NOT_FOUND => Err(Error::new(
774 ErrorKind::TableNotFound,
775 "Tried to drop a table that does not exist",
776 )),
777 _ => Err(deserialize_unexpected_catalog_error(
778 http_response,
779 client.http_client.disable_header_redaction(),
780 )),
781 }
782 }
783
784 fn auth_type(config: &RestCatalogConfig) -> String {
789 config
790 .props
791 .get(REST_CATALOG_PROP_AUTH_TYPE)
792 .map(|auth_type| auth_type.to_ascii_lowercase())
794 .unwrap_or_else(|| {
795 if config.token().is_some()
796 || config.credential().is_some()
797 || config.explicit_oauth2_server_uri().is_some()
798 {
799 AUTH_TYPE_OAUTH2.to_string()
800 } else {
801 AUTH_TYPE_NONE.to_string()
802 }
803 })
804 }
805
806 fn resolve_auth_manager(&self) -> Result<Arc<dyn AuthManager>> {
809 if let Some(auth_manager) = &self.auth_manager_override {
810 return Ok(auth_manager.clone());
811 }
812 let config = &self.user_config;
813 let auth_type = Self::auth_type(config);
814 if auth_type == AUTH_TYPE_OAUTH2 && !config.props.contains_key(REST_CATALOG_PROP_AUTH_TYPE)
817 {
818 tracing::warn!(
819 "Inferring {REST_CATALOG_PROP_AUTH_TYPE}={AUTH_TYPE_OAUTH2} from the configured \
820 OAuth properties; set it explicitly to avoid this warning"
821 );
822 }
823 match auth_type.as_str() {
824 AUTH_TYPE_NONE => Ok(Arc::new(NoopAuthManager)),
825 AUTH_TYPE_OAUTH2 => Ok(Arc::new(OAuth2Manager::from_config(config)?)),
826 other => Err(Error::new(
827 ErrorKind::DataInvalid,
828 format!(
829 "unknown '{REST_CATALOG_PROP_AUTH_TYPE}': {other}; use \
830 `RestSessionCatalogBuilder::with_auth_manager` or \
831 `RestCatalogBuilder::with_auth_manager` to inject a custom auth manager"
832 ),
833 )),
834 }
835 }
836
837 async fn client(&self) -> Result<&RestCatalogClient> {
839 self.client
840 .get_or_try_init(|| async {
841 RestCatalogClient::init(&self.user_config, self.resolve_auth_manager()?).await
842 })
843 .await
844 }
845
846 pub(crate) async fn supports_endpoint(&self, endpoint: &Endpoint) -> Result<bool> {
850 Ok(self.client().await?.endpoints.contains(endpoint))
851 }
852
853 async fn check_exists_via_head(
856 &self,
857 context: &SessionContext,
858 client: &RestCatalogClient,
859 url: String,
860 ) -> Result<bool> {
861 let request = HttpRequest::build(client.http_client.request(Method::HEAD, url))?;
862 let http_response = client.query_catalog(context, request).await?;
863
864 match http_response.status() {
865 StatusCode::NO_CONTENT | StatusCode::OK => Ok(true),
866 StatusCode::NOT_FOUND => Ok(false),
867 _ => Err(deserialize_unexpected_catalog_error(
868 http_response,
869 client.http_client.disable_header_redaction(),
870 )),
871 }
872 }
873
874 async fn load_file_io(
875 &self,
876 metadata_location: Option<&str>,
877 extra_config: Option<HashMap<String, String>>,
878 ) -> Result<FileIO> {
879 let mut props = self.client().await?.config.props.clone();
880 if let Some(config) = extra_config {
881 props.extend(config);
882 }
883
884 let warehouse_path = match self.client().await?.config.warehouse.as_deref() {
887 Some(url) if Url::parse(url).is_ok() => Some(url),
888 Some(_) => None,
889 None => None,
890 };
891
892 if metadata_location.or(warehouse_path).is_none() {
893 return Err(Error::new(
894 ErrorKind::Unexpected,
895 "Unable to load file io, neither warehouse nor metadata location is set!",
896 ));
897 }
898
899 let factory = self
901 .storage_factory
902 .clone()
903 .ok_or_else(|| {
904 Error::new(
905 ErrorKind::Unexpected,
906 "StorageFactory must be provided for REST catalog table operations. Use `with_storage_factory` to configure it.",
907 )
908 })?;
909
910 let file_io = FileIOBuilder::new(factory).with_props(props).build();
911
912 Ok(file_io)
913 }
914}
915
916#[async_trait]
919impl SessionCatalog for RestSessionCatalog {
920 async fn list_namespaces(
921 &self,
922 context: &SessionContext,
923 parent: Option<&NamespaceIdent>,
924 ) -> Result<Vec<NamespaceIdent>> {
925 let client = self.client().await?;
926 let endpoint = client.config.namespaces_endpoint();
927 let mut namespaces = Vec::new();
928 let mut next_token = None;
929
930 loop {
931 let mut request = client.http_client.request(Method::GET, endpoint.clone());
932
933 if let Some(ns) = parent {
935 request = request.query(&[("parent", ns.to_url_string())]);
936 }
937
938 if let Some(token) = next_token {
939 request = request.query(&[("pageToken", token)]);
940 }
941
942 let http_response = client
943 .query_catalog(context, HttpRequest::build(request)?)
944 .await?;
945
946 match http_response.status() {
947 StatusCode::OK => {
948 let response =
949 deserialize_catalog_response::<ListNamespaceResponse>(http_response)?;
950
951 namespaces.extend(response.namespaces);
952
953 match response.next_page_token {
954 Some(token) => next_token = Some(token),
955 None => break,
956 }
957 }
958 StatusCode::NOT_FOUND => {
959 return Err(Error::new(
960 ErrorKind::NamespaceNotFound,
961 "The parent parameter of the namespace provided does not exist",
962 ));
963 }
964 _ => {
965 return Err(deserialize_unexpected_catalog_error(
966 http_response,
967 client.http_client.disable_header_redaction(),
968 ));
969 }
970 }
971 }
972
973 Ok(namespaces)
974 }
975
976 async fn create_namespace(
977 &self,
978 context: &SessionContext,
979 namespace: &NamespaceIdent,
980 properties: HashMap<String, String>,
981 ) -> Result<Namespace> {
982 let client = self.client().await?;
983
984 let request = HttpRequest::build(
985 client
986 .http_client
987 .request(Method::POST, client.config.namespaces_endpoint())
988 .json(&CreateNamespaceRequest {
989 namespace: namespace.clone(),
990 properties,
991 }),
992 )?;
993
994 let http_response = client.query_catalog(context, request).await?;
995
996 match http_response.status() {
997 StatusCode::OK => {
998 let response = deserialize_catalog_response::<NamespaceResponse>(http_response)?;
999 Ok(Namespace::from(response))
1000 }
1001 StatusCode::CONFLICT => Err(Error::new(
1002 ErrorKind::NamespaceAlreadyExists,
1003 "Tried to create a namespace that already exists",
1004 )),
1005 _ => Err(deserialize_unexpected_catalog_error(
1006 http_response,
1007 client.http_client.disable_header_redaction(),
1008 )),
1009 }
1010 }
1011
1012 async fn get_namespace(
1013 &self,
1014 context: &SessionContext,
1015 namespace: &NamespaceIdent,
1016 ) -> Result<Namespace> {
1017 let client = self.client().await?;
1018
1019 let request = HttpRequest::build(
1020 client
1021 .http_client
1022 .request(Method::GET, client.config.namespace_endpoint(namespace)),
1023 )?;
1024
1025 let http_response = client.query_catalog(context, request).await?;
1026
1027 match http_response.status() {
1028 StatusCode::OK => {
1029 let response = deserialize_catalog_response::<NamespaceResponse>(http_response)?;
1030 Ok(Namespace::from(response))
1031 }
1032 StatusCode::NOT_FOUND => Err(Error::new(
1033 ErrorKind::NamespaceNotFound,
1034 "Tried to get a namespace that does not exist",
1035 )),
1036 _ => Err(deserialize_unexpected_catalog_error(
1037 http_response,
1038 client.http_client.disable_header_redaction(),
1039 )),
1040 }
1041 }
1042
1043 async fn namespace_exists(
1044 &self,
1045 context: &SessionContext,
1046 ns: &NamespaceIdent,
1047 ) -> Result<bool> {
1048 if !self.supports_endpoint(&V1_NAMESPACE_EXISTS).await? {
1053 return match self.get_namespace(context, ns).await {
1054 Ok(_) => Ok(true),
1055 Err(e) if e.kind() == ErrorKind::NamespaceNotFound => Ok(false),
1056 Err(e) => Err(e),
1057 };
1058 }
1059
1060 let client = self.client().await?;
1061 self.check_exists_via_head(context, client, client.config.namespace_endpoint(ns))
1062 .await
1063 }
1064
1065 async fn update_namespace(
1066 &self,
1067 _context: &SessionContext,
1068 _namespace: &NamespaceIdent,
1069 _properties: HashMap<String, String>,
1070 ) -> Result<()> {
1071 Err(Error::new(
1072 ErrorKind::FeatureUnsupported,
1073 "Updating namespace not supported yet!",
1074 ))
1075 }
1076
1077 async fn drop_namespace(
1078 &self,
1079 context: &SessionContext,
1080 namespace: &NamespaceIdent,
1081 ) -> Result<()> {
1082 let client = self.client().await?;
1083
1084 let request = HttpRequest::build(
1085 client
1086 .http_client
1087 .request(Method::DELETE, client.config.namespace_endpoint(namespace)),
1088 )?;
1089
1090 let http_response = client.query_catalog(context, request).await?;
1091
1092 match http_response.status() {
1093 StatusCode::NO_CONTENT | StatusCode::OK => Ok(()),
1094 StatusCode::NOT_FOUND => Err(Error::new(
1095 ErrorKind::NamespaceNotFound,
1096 "Tried to drop a namespace that does not exist",
1097 )),
1098 _ => Err(deserialize_unexpected_catalog_error(
1099 http_response,
1100 client.http_client.disable_header_redaction(),
1101 )),
1102 }
1103 }
1104
1105 async fn list_tables(
1106 &self,
1107 context: &SessionContext,
1108 namespace: &NamespaceIdent,
1109 ) -> Result<Vec<TableIdent>> {
1110 let client = self.client().await?;
1111 let endpoint = client.config.tables_endpoint(namespace);
1112 let mut identifiers = Vec::new();
1113 let mut next_token = None;
1114
1115 loop {
1116 let mut request = client.http_client.request(Method::GET, endpoint.clone());
1117
1118 if let Some(token) = next_token {
1119 request = request.query(&[("pageToken", token)]);
1120 }
1121
1122 let http_response = client
1123 .query_catalog(context, HttpRequest::build(request)?)
1124 .await?;
1125
1126 match http_response.status() {
1127 StatusCode::OK => {
1128 let response =
1129 deserialize_catalog_response::<ListTablesResponse>(http_response)?;
1130
1131 identifiers.extend(response.identifiers);
1132
1133 match response.next_page_token {
1134 Some(token) => next_token = Some(token),
1135 None => break,
1136 }
1137 }
1138 StatusCode::NOT_FOUND => {
1139 return Err(Error::new(
1140 ErrorKind::NamespaceNotFound,
1141 "Tried to list tables of a namespace that does not exist",
1142 ));
1143 }
1144 _ => {
1145 return Err(deserialize_unexpected_catalog_error(
1146 http_response,
1147 client.http_client.disable_header_redaction(),
1148 ));
1149 }
1150 }
1151 }
1152
1153 Ok(identifiers)
1154 }
1155
1156 async fn create_table(
1163 &self,
1164 context: &SessionContext,
1165 namespace: &NamespaceIdent,
1166 creation: TableCreation,
1167 ) -> Result<Table> {
1168 let client = self.client().await?;
1169
1170 let table_ident = TableIdent::new(namespace.clone(), creation.name.clone());
1171
1172 let request = HttpRequest::build(
1173 client
1174 .http_client
1175 .request(Method::POST, client.config.tables_endpoint(namespace))
1176 .json(&CreateTableRequest {
1177 name: creation.name,
1178 location: creation.location,
1179 schema: creation.schema,
1180 partition_spec: creation.partition_spec,
1181 write_order: creation.sort_order,
1182 stage_create: Some(false),
1183 properties: creation.properties,
1184 }),
1185 )?;
1186
1187 let http_response = client.query_catalog(context, request).await?;
1188
1189 let response = match http_response.status() {
1190 StatusCode::OK => deserialize_catalog_response::<LoadTableResult>(http_response)?,
1191 StatusCode::NOT_FOUND => {
1192 return Err(Error::new(
1193 ErrorKind::NamespaceNotFound,
1194 "Tried to create a table under a namespace that does not exist",
1195 ));
1196 }
1197 StatusCode::CONFLICT => {
1198 return Err(Error::new(
1199 ErrorKind::TableAlreadyExists,
1200 "The table already exists",
1201 ));
1202 }
1203 _ => {
1204 return Err(deserialize_unexpected_catalog_error(
1205 http_response,
1206 client.http_client.disable_header_redaction(),
1207 ));
1208 }
1209 };
1210
1211 let metadata_location = response.metadata_location.as_ref().ok_or(Error::new(
1212 ErrorKind::DataInvalid,
1213 "Metadata location missing in `create_table` response!",
1214 ))?;
1215
1216 let config = response
1217 .config
1218 .into_iter()
1219 .chain(self.user_config.props.clone())
1220 .collect();
1221
1222 let file_io = self
1223 .load_file_io(Some(metadata_location), Some(config))
1224 .await?;
1225
1226 let mut table_builder = Table::builder()
1227 .identifier(table_ident.clone())
1228 .file_io(file_io)
1229 .metadata(response.metadata)
1230 .runtime(self.runtime.clone());
1231 if let Some(kms_client) = self.kms_client.clone() {
1232 table_builder = table_builder.kms_client(kms_client);
1233 }
1234
1235 if let Some(metadata_location) = response.metadata_location {
1236 table_builder.metadata_location(metadata_location).build()
1237 } else {
1238 table_builder.build()
1239 }
1240 }
1241
1242 async fn load_table(
1248 &self,
1249 context: &SessionContext,
1250 table_ident: &TableIdent,
1251 ) -> Result<Table> {
1252 let client = self.client().await?;
1253
1254 let request = HttpRequest::build(
1255 client
1256 .http_client
1257 .request(Method::GET, client.config.table_endpoint(table_ident)),
1258 )?;
1259
1260 let http_response = client.query_catalog(context, request).await?;
1261
1262 let response = match http_response.status() {
1263 StatusCode::OK | StatusCode::NOT_MODIFIED => {
1264 deserialize_catalog_response::<LoadTableResult>(http_response)?
1265 }
1266 StatusCode::NOT_FOUND => {
1267 return Err(Error::new(
1268 ErrorKind::TableNotFound,
1269 "Tried to load a table that does not exist",
1270 ));
1271 }
1272 _ => {
1273 return Err(deserialize_unexpected_catalog_error(
1274 http_response,
1275 client.http_client.disable_header_redaction(),
1276 ));
1277 }
1278 };
1279
1280 let config = response
1281 .config
1282 .into_iter()
1283 .chain(self.user_config.props.clone())
1284 .collect();
1285
1286 let file_io = self
1287 .load_file_io(response.metadata_location.as_deref(), Some(config))
1288 .await?;
1289
1290 let mut table_builder = Table::builder()
1291 .identifier(table_ident.clone())
1292 .file_io(file_io)
1293 .metadata(response.metadata)
1294 .runtime(self.runtime.clone());
1295 if let Some(kms_client) = self.kms_client.clone() {
1296 table_builder = table_builder.kms_client(kms_client);
1297 }
1298
1299 if let Some(metadata_location) = response.metadata_location {
1300 table_builder.metadata_location(metadata_location).build()
1301 } else {
1302 table_builder.build()
1303 }
1304 }
1305
1306 async fn drop_table(&self, context: &SessionContext, table: &TableIdent) -> Result<()> {
1308 self.delete_table(context, table, false).await
1309 }
1310
1311 async fn purge_table(&self, context: &SessionContext, table: &TableIdent) -> Result<()> {
1314 self.delete_table(context, table, true).await
1315 }
1316
1317 async fn table_exists(&self, context: &SessionContext, table: &TableIdent) -> Result<bool> {
1319 if !self.supports_endpoint(&V1_TABLE_EXISTS).await? {
1323 return match self.load_table(context, table).await {
1324 Ok(_) => Ok(true),
1325 Err(e) if e.kind() == ErrorKind::TableNotFound => Ok(false),
1326 Err(e) => Err(e),
1327 };
1328 }
1329
1330 let client = self.client().await?;
1331 self.check_exists_via_head(context, client, client.config.table_endpoint(table))
1332 .await
1333 }
1334
1335 async fn rename_table(
1337 &self,
1338 context: &SessionContext,
1339 src: &TableIdent,
1340 dest: &TableIdent,
1341 ) -> Result<()> {
1342 let client = self.client().await?;
1343
1344 let request = HttpRequest::build(
1345 client
1346 .http_client
1347 .request(Method::POST, client.config.rename_table_endpoint())
1348 .json(&RenameTableRequest {
1349 source: src.clone(),
1350 destination: dest.clone(),
1351 }),
1352 )?;
1353
1354 let http_response = client.query_catalog(context, request).await?;
1355
1356 match http_response.status() {
1357 StatusCode::NO_CONTENT | StatusCode::OK => Ok(()),
1358 StatusCode::NOT_FOUND => Err(Error::new(
1359 ErrorKind::TableNotFound,
1360 "Tried to rename a table that does not exist (is the namespace correct?)",
1361 )),
1362 StatusCode::CONFLICT => Err(Error::new(
1363 ErrorKind::TableAlreadyExists,
1364 "Tried to rename a table to a name that already exists",
1365 )),
1366 _ => Err(deserialize_unexpected_catalog_error(
1367 http_response,
1368 client.http_client.disable_header_redaction(),
1369 )),
1370 }
1371 }
1372
1373 async fn register_table(
1374 &self,
1375 context: &SessionContext,
1376 table_ident: &TableIdent,
1377 metadata_location: String,
1378 ) -> Result<Table> {
1379 let client = self.client().await?;
1380
1381 let request = HttpRequest::build(
1382 client
1383 .http_client
1384 .request(
1385 Method::POST,
1386 client
1387 .config
1388 .register_table_endpoint(table_ident.namespace()),
1389 )
1390 .json(&RegisterTableRequest {
1391 name: table_ident.name.clone(),
1392 metadata_location: metadata_location.clone(),
1393 overwrite: Some(false),
1394 }),
1395 )?;
1396
1397 let http_response = client.query_catalog(context, request).await?;
1398
1399 let response: LoadTableResult = match http_response.status() {
1400 StatusCode::OK => deserialize_catalog_response::<LoadTableResult>(http_response)?,
1401 StatusCode::NOT_FOUND => {
1402 return Err(Error::new(
1403 ErrorKind::NamespaceNotFound,
1404 "The namespace specified does not exist.",
1405 ));
1406 }
1407 StatusCode::CONFLICT => {
1408 return Err(Error::new(
1409 ErrorKind::TableAlreadyExists,
1410 "The given table already exists.",
1411 ));
1412 }
1413 _ => {
1414 return Err(deserialize_unexpected_catalog_error(
1415 http_response,
1416 client.http_client.disable_header_redaction(),
1417 ));
1418 }
1419 };
1420
1421 let metadata_location = response.metadata_location.as_ref().ok_or(Error::new(
1422 ErrorKind::DataInvalid,
1423 "Metadata location missing in `register_table` response!",
1424 ))?;
1425
1426 let file_io = self.load_file_io(Some(metadata_location), None).await?;
1427
1428 let mut table_builder = Table::builder()
1429 .identifier(table_ident.clone())
1430 .file_io(file_io)
1431 .metadata(response.metadata)
1432 .metadata_location(metadata_location.clone())
1433 .runtime(self.runtime.clone());
1434 if let Some(kms_client) = self.kms_client.clone() {
1435 table_builder = table_builder.kms_client(kms_client);
1436 }
1437 table_builder.build()
1438 }
1439
1440 async fn update_table(
1441 &self,
1442 context: &SessionContext,
1443 mut commit: TableCommit,
1444 ) -> Result<Table> {
1445 let client = self.client().await?;
1446
1447 let request = HttpRequest::build(
1448 client
1449 .http_client
1450 .request(
1451 Method::POST,
1452 client.config.table_endpoint(commit.identifier()),
1453 )
1454 .json(&CommitTableRequest {
1455 identifier: Some(commit.identifier().clone()),
1456 requirements: commit.take_requirements(),
1457 updates: commit.take_updates(),
1458 }),
1459 )?;
1460
1461 let http_response = client.query_catalog(context, request).await?;
1462
1463 let response: CommitTableResponse = match http_response.status() {
1464 StatusCode::OK => deserialize_catalog_response(http_response)?,
1465 StatusCode::NOT_FOUND => {
1466 return Err(Error::new(
1467 ErrorKind::TableNotFound,
1468 "Tried to update a table that does not exist",
1469 ));
1470 }
1471 StatusCode::CONFLICT => {
1472 return Err(Error::new(
1473 ErrorKind::CatalogCommitConflicts,
1474 "CatalogCommitConflicts, one or more requirements failed. The client may retry.",
1475 )
1476 .with_retryable(true));
1477 }
1478 StatusCode::INTERNAL_SERVER_ERROR => {
1479 return Err(Error::new(
1480 ErrorKind::Unexpected,
1481 "An unknown server-side problem occurred; the commit state is unknown.",
1482 ));
1483 }
1484 StatusCode::BAD_GATEWAY => {
1485 return Err(Error::new(
1486 ErrorKind::Unexpected,
1487 "A gateway or proxy received an invalid response from the upstream server; the commit state is unknown.",
1488 ));
1489 }
1490 StatusCode::GATEWAY_TIMEOUT => {
1491 return Err(Error::new(
1492 ErrorKind::Unexpected,
1493 "A server-side gateway timeout occurred; the commit state is unknown.",
1494 ));
1495 }
1496 _ => {
1497 return Err(deserialize_unexpected_catalog_error(
1498 http_response,
1499 client.http_client.disable_header_redaction(),
1500 ));
1501 }
1502 };
1503
1504 let file_io = self
1505 .load_file_io(Some(&response.metadata_location), None)
1506 .await?;
1507
1508 let mut table_builder = Table::builder()
1509 .identifier(commit.identifier().clone())
1510 .file_io(file_io)
1511 .metadata(response.metadata)
1512 .metadata_location(response.metadata_location)
1513 .runtime(self.runtime.clone());
1514 if let Some(kms_client) = self.kms_client.clone() {
1515 table_builder = table_builder.kms_client(kms_client);
1516 }
1517 table_builder.build()
1518 }
1519}
1520
1521#[derive(Debug)]
1526pub struct RestSessionCatalogBuilder {
1527 config: RestCatalogConfig,
1528 auth_manager: Option<Box<dyn AuthManager>>,
1529 storage_factory: Option<Arc<dyn StorageFactory>>,
1530 kms_client_factory: Option<Arc<dyn KmsClientFactory>>,
1531 runtime: Option<Runtime>,
1532}
1533
1534impl Default for RestSessionCatalogBuilder {
1535 fn default() -> Self {
1536 Self {
1537 config: RestCatalogConfig {
1538 name: None,
1539 uri: "".to_string(),
1540 warehouse: None,
1541 props: HashMap::new(),
1542 client: None,
1543 default_client: Arc::new(OnceLock::new()),
1544 },
1545 auth_manager: None,
1546 storage_factory: None,
1547 kms_client_factory: None,
1548 runtime: None,
1549 }
1550 }
1551}
1552
1553impl RestSessionCatalogBuilder {
1554 pub fn with_client(mut self, client: Client) -> Self {
1556 self.config.client = Some(client);
1557 self
1558 }
1559
1560 pub fn with_auth_manager<M>(mut self, auth_manager: M) -> Self
1565 where M: AuthManager + 'static {
1566 self.auth_manager = Some(Box::new(auth_manager));
1567 self
1568 }
1569
1570 pub fn with_storage_factory(mut self, storage_factory: Arc<dyn StorageFactory>) -> Self {
1595 self.storage_factory = Some(storage_factory);
1596 self
1597 }
1598
1599 pub fn with_kms_client_factory(
1620 mut self,
1621 kms_client_factory: Arc<dyn KmsClientFactory>,
1622 ) -> Self {
1623 self.kms_client_factory = Some(kms_client_factory);
1624 self
1625 }
1626
1627 pub fn with_runtime(mut self, runtime: Runtime) -> Self {
1633 self.runtime = Some(runtime);
1634 self
1635 }
1636
1637 pub fn load(
1642 mut self,
1643 name: impl Into<String>,
1644 props: HashMap<String, String>,
1645 ) -> impl Future<Output = Result<RestSessionCatalog>> + Send {
1646 self.config.name = Some(name.into());
1647
1648 if props.contains_key(REST_CATALOG_PROP_URI) {
1649 self.config.uri = props
1650 .get(REST_CATALOG_PROP_URI)
1651 .cloned()
1652 .unwrap_or_default();
1653 }
1654
1655 if props.contains_key(REST_CATALOG_PROP_WAREHOUSE) {
1656 self.config.warehouse = props.get(REST_CATALOG_PROP_WAREHOUSE).cloned()
1657 }
1658
1659 self.config.props = props
1661 .into_iter()
1662 .filter(|(k, _)| k != REST_CATALOG_PROP_URI && k != REST_CATALOG_PROP_WAREHOUSE)
1663 .collect();
1664
1665 async move {
1666 if self.config.name.is_none() {
1667 Err(Error::new(
1668 ErrorKind::DataInvalid,
1669 "Catalog name is required",
1670 ))
1671 } else if self.config.uri.is_empty() {
1672 Err(Error::new(
1673 ErrorKind::DataInvalid,
1674 "Catalog uri is required",
1675 ))
1676 } else {
1677 let runtime = self.runtime.unwrap_or_else(Runtime::current);
1678 let kms_client = match self.kms_client_factory {
1679 Some(factory) => Some(factory.create_kms_client(&self.config.props).await?),
1680 None => None,
1681 };
1682
1683 Ok(RestSessionCatalog::new(
1684 self.config,
1685 self.auth_manager,
1686 self.storage_factory,
1687 runtime,
1688 kms_client,
1689 ))
1690 }
1691 }
1692 }
1693}
1694
1695#[cfg(test)]
1696mod tests {
1697 use std::fs::File;
1698 use std::io::BufReader;
1699 use std::sync::{Arc, Mutex};
1700
1701 use chrono::{TimeZone, Utc};
1702 use iceberg::io::LocalFsStorageFactory;
1703 use iceberg::spec::{
1704 FormatVersion, NestedField, NullOrder, Operation, PrimitiveType, Schema, Snapshot,
1705 SnapshotLog, SortDirection, SortField, SortOrder, Summary, Transform, Type,
1706 UnboundPartitionField, UnboundPartitionSpec,
1707 };
1708 use iceberg::transaction::{ApplyTransactionAction, Transaction};
1709 use iceberg_test_utils::test_runtime;
1710 use mockito::{Mock, Server, ServerGuard};
1711 use serde_json::json;
1712 use uuid::uuid;
1713
1714 use super::*;
1715 use crate::auth::AuthSession;
1716 use crate::request::HttpRequest;
1717
1718 #[derive(Debug)]
1719 struct PlainContextSession;
1720
1721 #[async_trait]
1722 impl AuthSession for PlainContextSession {
1723 async fn authenticate(&self, _request: &mut HttpRequest) -> Result<()> {
1724 Ok(())
1725 }
1726 }
1727
1728 #[derive(Debug)]
1729 struct ContextSession(String);
1730
1731 #[async_trait]
1732 impl AuthSession for ContextSession {
1733 async fn authenticate(&self, request: &mut HttpRequest) -> Result<()> {
1734 request.headers_mut().insert(
1735 "x-session-id",
1736 HeaderValue::from_str(&self.0).expect("valid test session ID"),
1737 );
1738 Ok(())
1739 }
1740 }
1741
1742 #[derive(Debug)]
1743 struct ContextManager {
1744 catalog_session: Arc<dyn AuthSession>,
1745 seen_session_ids: Arc<Mutex<Vec<String>>>,
1746 fail_contextual_session: bool,
1747 }
1748
1749 #[async_trait]
1750 impl AuthManager for ContextManager {
1751 async fn init_session(
1752 &self,
1753 _client: &HttpClient,
1754 _props: &HashMap<String, String>,
1755 ) -> Result<Box<dyn AuthSession>> {
1756 Ok(Box::new(PlainContextSession))
1757 }
1758
1759 async fn catalog_session(
1760 &self,
1761 _client: &HttpClient,
1762 _props: &HashMap<String, String>,
1763 ) -> Result<Arc<dyn AuthSession>> {
1764 Ok(self.catalog_session.clone())
1765 }
1766
1767 async fn contextual_session(
1768 &self,
1769 context: &SessionContext,
1770 catalog_session: Arc<dyn AuthSession>,
1771 ) -> Result<Arc<dyn AuthSession>> {
1772 assert!(Arc::ptr_eq(&catalog_session, &self.catalog_session));
1773 self.seen_session_ids
1774 .lock()
1775 .unwrap()
1776 .push(context.session_id().to_string());
1777 if self.fail_contextual_session {
1778 return Err(Error::new(
1779 ErrorKind::Unexpected,
1780 "contextual session failure",
1781 ));
1782 }
1783 Ok(Arc::new(ContextSession(context.session_id().to_string())))
1784 }
1785 }
1786
1787 fn test_catalog(config: RestCatalogConfig) -> RestSessionCatalog {
1788 RestSessionCatalog::new(config, None, None, Runtime::current(), None)
1789 }
1790
1791 fn test_catalog_with<M>(config: RestCatalogConfig, auth_manager: M) -> RestSessionCatalog
1792 where M: AuthManager + 'static {
1793 RestSessionCatalog::new(
1794 config,
1795 Some(Box::new(auth_manager)),
1796 None,
1797 Runtime::current(),
1798 None,
1799 )
1800 }
1801
1802 fn context_catalog(
1803 config: RestCatalogConfig,
1804 fail_contextual_session: bool,
1805 ) -> (RestSessionCatalog, Arc<Mutex<Vec<String>>>) {
1806 let seen_session_ids = Arc::new(Mutex::new(Vec::new()));
1807 let catalog_session: Arc<dyn AuthSession> = Arc::new(PlainContextSession);
1808 let catalog = test_catalog_with(config, ContextManager {
1809 catalog_session,
1810 seen_session_ids: seen_session_ids.clone(),
1811 fail_contextual_session,
1812 });
1813 (catalog, seen_session_ids)
1814 }
1815
1816 fn test_client() -> HttpClient {
1817 HttpClient::new(
1818 &RestCatalogConfig::builder()
1819 .uri("http://localhost".to_string())
1820 .build(),
1821 )
1822 .unwrap()
1823 }
1824
1825 fn session_catalog(config: RestCatalogConfig) -> RestSessionCatalog {
1827 RestSessionCatalog::new(
1828 config,
1829 None,
1830 Some(Arc::new(LocalFsStorageFactory)),
1831 Runtime::current(),
1832 None,
1833 )
1834 }
1835
1836 #[tokio::test]
1837 async fn test_update_config() {
1838 let mut server = Server::new_async().await;
1839
1840 let config_mock = server
1841 .mock("GET", "/v1/config")
1842 .with_status(200)
1843 .with_body(
1844 r#"{
1845 "overrides": {
1846 "warehouse": "s3://iceberg-catalog"
1847 },
1848 "defaults": {}
1849 }"#,
1850 )
1851 .create_async()
1852 .await;
1853
1854 let catalog = session_catalog(RestCatalogConfig::builder().uri(server.url()).build());
1855
1856 assert_eq!(
1857 catalog
1858 .client()
1859 .await
1860 .unwrap()
1861 .config
1862 .props
1863 .get("warehouse"),
1864 Some(&"s3://iceberg-catalog".to_string())
1865 );
1866
1867 config_mock.assert_async().await;
1868 }
1869
1870 async fn create_config_mock(server: &mut ServerGuard) -> Mock {
1871 server
1872 .mock("GET", "/v1/config")
1873 .with_status(200)
1874 .with_body(
1875 r#"{
1876 "overrides": {
1877 "warehouse": "s3://iceberg-catalog"
1878 },
1879 "defaults": {}
1880 }"#,
1881 )
1882 .create_async()
1883 .await
1884 }
1885
1886 async fn create_config_mock_with_exists_endpoints(server: &mut ServerGuard) -> Mock {
1889 server
1890 .mock("GET", "/v1/config")
1891 .with_status(200)
1892 .with_body(
1893 r#"{
1894 "overrides": { "warehouse": "s3://iceberg-catalog" },
1895 "defaults": {},
1896 "endpoints": [
1897 "HEAD /v1/{prefix}/namespaces/{namespace}",
1898 "HEAD /v1/{prefix}/namespaces/{namespace}/tables/{table}"
1899 ]
1900 }"#,
1901 )
1902 .create_async()
1903 .await
1904 }
1905
1906 #[tokio::test]
1907 async fn test_config_advertised_endpoints() {
1908 let mut server = Server::new_async().await;
1909
1910 let config_mock = server
1911 .mock("GET", "/v1/config")
1912 .with_status(200)
1913 .with_body(
1914 r#"{
1915 "overrides": {},
1916 "defaults": {},
1917 "endpoints": [
1918 "GET /v1/{prefix}/namespaces",
1919 "POST /v1/{prefix}/namespaces/{namespace}/tables/{table}/plan"
1920 ]
1921 }"#,
1922 )
1923 .create_async()
1924 .await;
1925
1926 let catalog = RestCatalog::new(
1927 SessionContext::empty(),
1928 RestCatalogConfig::builder().uri(server.url()).build(),
1929 None,
1930 Some(Arc::new(LocalFsStorageFactory)),
1931 Runtime::current(),
1932 None,
1933 );
1934
1935 let plan = "POST /v1/{prefix}/namespaces/{namespace}/tables/{table}/plan"
1936 .parse::<Endpoint>()
1937 .unwrap();
1938 assert!(catalog.inner.supports_endpoint(&plan).await.unwrap());
1939 let delete_ns = "DELETE /v1/{prefix}/namespaces/{namespace}"
1941 .parse::<Endpoint>()
1942 .unwrap();
1943 assert!(!catalog.inner.supports_endpoint(&delete_ns).await.unwrap());
1944
1945 config_mock.assert_async().await;
1946 }
1947
1948 #[tokio::test]
1949 async fn test_config_without_endpoints_falls_back_to_default_set() {
1950 let mut server = Server::new_async().await;
1951
1952 let config_mock = server
1953 .mock("GET", "/v1/config")
1954 .with_status(200)
1955 .with_body(r#"{ "overrides": {}, "defaults": {} }"#)
1956 .create_async()
1957 .await;
1958
1959 let catalog = RestCatalog::new(
1960 SessionContext::empty(),
1961 RestCatalogConfig::builder().uri(server.url()).build(),
1962 None,
1963 Some(Arc::new(LocalFsStorageFactory)),
1964 Runtime::current(),
1965 None,
1966 );
1967
1968 let load_table = "GET /v1/{prefix}/namespaces/{namespace}/tables/{table}"
1971 .parse::<Endpoint>()
1972 .unwrap();
1973 assert!(catalog.inner.supports_endpoint(&load_table).await.unwrap());
1974 let plan = "POST /v1/{prefix}/namespaces/{namespace}/tables/{table}/plan"
1976 .parse::<Endpoint>()
1977 .unwrap();
1978 assert!(!catalog.inner.supports_endpoint(&plan).await.unwrap());
1979
1980 config_mock.assert_async().await;
1981 }
1982
1983 #[tokio::test]
1984 async fn test_config_with_empty_endpoints_falls_back_to_default_set() {
1985 let mut server = Server::new_async().await;
1986
1987 let config_mock = server
1990 .mock("GET", "/v1/config")
1991 .with_status(200)
1992 .with_body(r#"{ "overrides": {}, "defaults": {}, "endpoints": [] }"#)
1993 .create_async()
1994 .await;
1995
1996 let catalog = RestCatalog::new(
1997 SessionContext::empty(),
1998 RestCatalogConfig::builder().uri(server.url()).build(),
1999 None,
2000 Some(Arc::new(LocalFsStorageFactory)),
2001 Runtime::current(),
2002 None,
2003 );
2004
2005 let load_table = "GET /v1/{prefix}/namespaces/{namespace}/tables/{table}"
2006 .parse::<Endpoint>()
2007 .unwrap();
2008 assert!(catalog.inner.supports_endpoint(&load_table).await.unwrap());
2009
2010 config_mock.assert_async().await;
2011 }
2012
2013 async fn create_oauth_mock(server: &mut ServerGuard) -> Mock {
2014 create_oauth_mock_with_path(server, "/v1/oauth/tokens", "ey000000000000", 200).await
2015 }
2016
2017 async fn create_oauth_mock_with_path(
2018 server: &mut ServerGuard,
2019 path: &str,
2020 token: &str,
2021 status: usize,
2022 ) -> Mock {
2023 let body = format!(
2024 r#"{{
2025 "access_token": "{token}",
2026 "token_type": "Bearer",
2027 "issued_token_type": "urn:ietf:params:oauth:token-type:access_token",
2028 "expires_in": 86400
2029 }}"#
2030 );
2031 server
2032 .mock("POST", path)
2033 .with_status(status)
2034 .with_body(body)
2035 .expect(1)
2036 .create_async()
2037 .await
2038 }
2039
2040 #[tokio::test]
2041 async fn test_oauth() {
2042 let mut server = Server::new_async().await;
2043 let oauth_mock = create_oauth_mock(&mut server).await;
2044 let config_mock = create_config_mock(&mut server).await;
2045
2046 let mut props = HashMap::new();
2047 props.insert("credential".to_string(), "client1:secret1".to_string());
2048
2049 let catalog = session_catalog(
2050 RestCatalogConfig::builder()
2051 .uri(server.url())
2052 .props(props)
2053 .build(),
2054 );
2055
2056 let token = catalog.client().await.unwrap().token().await;
2057 oauth_mock.assert_async().await;
2058 config_mock.assert_async().await;
2059 assert_eq!(token, Some("ey000000000000".to_string()));
2060 }
2061
2062 #[tokio::test]
2063 async fn test_oauth_with_optional_param() {
2064 let mut props = HashMap::new();
2065 props.insert("credential".to_string(), "client1:secret1".to_string());
2066 props.insert("scope".to_string(), "custom_scope".to_string());
2067 props.insert("audience".to_string(), "custom_audience".to_string());
2068 props.insert("resource".to_string(), "custom_resource".to_string());
2069
2070 let mut server = Server::new_async().await;
2071 let oauth_mock = server
2072 .mock("POST", "/v1/oauth/tokens")
2073 .match_body(mockito::Matcher::Regex("scope=custom_scope".to_string()))
2074 .match_body(mockito::Matcher::Regex(
2075 "audience=custom_audience".to_string(),
2076 ))
2077 .match_body(mockito::Matcher::Regex(
2078 "resource=custom_resource".to_string(),
2079 ))
2080 .with_status(200)
2081 .with_body(
2082 r#"{
2083 "access_token": "ey000000000000",
2084 "token_type": "Bearer",
2085 "issued_token_type": "urn:ietf:params:oauth:token-type:access_token",
2086 "expires_in": 86400
2087 }"#,
2088 )
2089 .expect(1)
2090 .create_async()
2091 .await;
2092
2093 let config_mock = create_config_mock(&mut server).await;
2094
2095 let catalog = session_catalog(
2096 RestCatalogConfig::builder()
2097 .uri(server.url())
2098 .props(props)
2099 .build(),
2100 );
2101
2102 let token = catalog.client().await.unwrap().token().await;
2103
2104 oauth_mock.assert_async().await;
2105 config_mock.assert_async().await;
2106 assert_eq!(token, Some("ey000000000000".to_string()));
2107 }
2108
2109 #[tokio::test]
2110 async fn test_http_headers() {
2111 let server = Server::new_async().await;
2112 let mut props = HashMap::new();
2113 props.insert("credential".to_string(), "client1:secret1".to_string());
2114
2115 let config = RestCatalogConfig::builder()
2116 .uri(server.url())
2117 .props(props)
2118 .build();
2119 let headers: HeaderMap = config.extra_headers().unwrap();
2120
2121 let expected_headers = HeaderMap::from_iter([
2122 (
2123 header::CONTENT_TYPE,
2124 HeaderValue::from_static("application/json"),
2125 ),
2126 (
2127 HeaderName::from_static("x-client-version"),
2128 HeaderValue::from_static(ICEBERG_REST_SPEC_VERSION),
2129 ),
2130 (
2131 header::USER_AGENT,
2132 HeaderValue::from_str(&format!("iceberg-rs/{CARGO_PKG_VERSION}")).unwrap(),
2133 ),
2134 ]);
2135 assert_eq!(headers, expected_headers);
2136 }
2137
2138 #[tokio::test]
2139 async fn test_http_headers_with_custom_headers() {
2140 let server = Server::new_async().await;
2141 let mut props = HashMap::new();
2142 props.insert("credential".to_string(), "client1:secret1".to_string());
2143 props.insert(
2144 "header.content-type".to_string(),
2145 "application/yaml".to_string(),
2146 );
2147 props.insert(
2148 "header.customized-header".to_string(),
2149 "some/value".to_string(),
2150 );
2151
2152 let config = RestCatalogConfig::builder()
2153 .uri(server.url())
2154 .props(props)
2155 .build();
2156 let headers: HeaderMap = config.extra_headers().unwrap();
2157
2158 let expected_headers = HeaderMap::from_iter([
2159 (
2160 header::CONTENT_TYPE,
2161 HeaderValue::from_static("application/yaml"),
2162 ),
2163 (
2164 HeaderName::from_static("x-client-version"),
2165 HeaderValue::from_static(ICEBERG_REST_SPEC_VERSION),
2166 ),
2167 (
2168 header::USER_AGENT,
2169 HeaderValue::from_str(&format!("iceberg-rs/{CARGO_PKG_VERSION}")).unwrap(),
2170 ),
2171 (
2172 HeaderName::from_static("customized-header"),
2173 HeaderValue::from_static("some/value"),
2174 ),
2175 ]);
2176 assert_eq!(headers, expected_headers);
2177 }
2178
2179 #[tokio::test]
2180 async fn test_oauth_with_oauth2_server_uri() {
2181 let mut server = Server::new_async().await;
2182 let config_mock = create_config_mock(&mut server).await;
2183
2184 let mut auth_server = Server::new_async().await;
2185 let auth_server_path = "/some/path";
2186 let oauth_mock =
2187 create_oauth_mock_with_path(&mut auth_server, auth_server_path, "ey000000000000", 200)
2188 .await;
2189
2190 let mut props = HashMap::new();
2191 props.insert("credential".to_string(), "client1:secret1".to_string());
2192 props.insert(
2193 "oauth2-server-uri".to_string(),
2194 format!("{}{}", auth_server.url(), auth_server_path).to_string(),
2195 );
2196
2197 let catalog = session_catalog(
2198 RestCatalogConfig::builder()
2199 .uri(server.url())
2200 .props(props)
2201 .build(),
2202 );
2203
2204 let token = catalog.client().await.unwrap().token().await;
2205
2206 oauth_mock.assert_async().await;
2207 config_mock.assert_async().await;
2208 assert_eq!(token, Some("ey000000000000".to_string()));
2209 }
2210
2211 #[tokio::test]
2212 async fn test_config_override() {
2213 let mut server = Server::new_async().await;
2214 let mut redirect_server = Server::new_async().await;
2215 let new_uri = redirect_server.url();
2216
2217 let config_mock = server
2218 .mock("GET", "/v1/config")
2219 .with_status(200)
2220 .with_body(
2221 json!(
2222 {
2223 "overrides": {
2224 "uri": new_uri,
2225 "warehouse": "s3://iceberg-catalog",
2226 "prefix": "ice/warehouses/my"
2227 },
2228 "defaults": {},
2229 }
2230 )
2231 .to_string(),
2232 )
2233 .create_async()
2234 .await;
2235
2236 let list_ns_mock = redirect_server
2237 .mock("GET", "/v1/ice/warehouses/my/namespaces")
2238 .with_body(
2239 r#"{
2240 "namespaces": []
2241 }"#,
2242 )
2243 .create_async()
2244 .await;
2245
2246 let catalog = session_catalog(RestCatalogConfig::builder().uri(server.url()).build());
2247
2248 let _namespaces = catalog
2249 .list_namespaces(&SessionContext::empty(), None)
2250 .await
2251 .unwrap();
2252
2253 config_mock.assert_async().await;
2254 list_ns_mock.assert_async().await;
2255 }
2256
2257 #[tokio::test]
2258 async fn test_list_namespace() {
2259 let mut server = Server::new_async().await;
2260
2261 let config_mock = create_config_mock(&mut server).await;
2262
2263 let list_ns_mock = server
2264 .mock("GET", "/v1/namespaces")
2265 .with_body(
2266 r#"{
2267 "namespaces": [
2268 ["ns1", "ns11"],
2269 ["ns2"]
2270 ]
2271 }"#,
2272 )
2273 .create_async()
2274 .await;
2275
2276 let catalog = session_catalog(RestCatalogConfig::builder().uri(server.url()).build());
2277
2278 let namespaces = catalog
2279 .list_namespaces(&SessionContext::empty(), None)
2280 .await
2281 .unwrap();
2282
2283 let expected_ns = vec![
2284 NamespaceIdent::from_vec(vec!["ns1".to_string(), "ns11".to_string()]).unwrap(),
2285 NamespaceIdent::from_vec(vec!["ns2".to_string()]).unwrap(),
2286 ];
2287
2288 assert_eq!(expected_ns, namespaces);
2289
2290 config_mock.assert_async().await;
2291 list_ns_mock.assert_async().await;
2292 }
2293
2294 #[tokio::test]
2295 async fn test_auth_type_none_disables_auth() {
2296 let props = HashMap::from([
2298 (REST_CATALOG_PROP_AUTH_TYPE.to_string(), "none".to_string()),
2299 ("token".to_string(), "some-oauth-token".to_string()),
2300 ]);
2301 let config = RestCatalogConfig::builder()
2302 .uri("http://localhost".to_string())
2303 .props(props)
2304 .build();
2305
2306 let session = test_catalog(config)
2307 .resolve_auth_manager()
2308 .unwrap()
2309 .init_session(&test_client(), &HashMap::new())
2310 .await
2311 .unwrap();
2312 let mut req = HttpRequest::new(
2313 Client::new()
2314 .get("https://rest.example.com/v1/config")
2315 .build()
2316 .unwrap(),
2317 );
2318 session.authenticate(&mut req).await.unwrap();
2319 assert!(req.headers().get("authorization").is_none());
2320 }
2321
2322 #[tokio::test]
2323 async fn test_header_prop_overrides_token_on_the_wire() {
2324 let mut server = Server::new_async().await;
2328 let config_mock = create_config_mock(&mut server).await;
2329 let list_ns_mock = server
2330 .mock("GET", "/v1/namespaces")
2331 .match_header("authorization", "Basic xyz")
2332 .with_body(r#"{"namespaces": []}"#)
2333 .create_async()
2334 .await;
2335
2336 let props = HashMap::from([
2337 ("token".to_string(), "some-oauth-token".to_string()),
2338 ("header.authorization".to_string(), "Basic xyz".to_string()),
2339 ]);
2340 let catalog = RestCatalog::new(
2341 SessionContext::empty(),
2342 RestCatalogConfig::builder()
2343 .uri(server.url())
2344 .props(props)
2345 .build(),
2346 None,
2347 Some(Arc::new(LocalFsStorageFactory)),
2348 Runtime::current(),
2349 None,
2350 );
2351
2352 catalog.list_namespaces(None).await.unwrap();
2353 config_mock.assert_async().await;
2354 list_ns_mock.assert_async().await;
2355 }
2356
2357 #[tokio::test]
2358 async fn test_builtin_oauth_endpoint_follows_uri_override() {
2359 let mut bootstrap = Server::new_async().await;
2363 let overridden = Server::new_async().await;
2364
2365 let config_mock = bootstrap
2366 .mock("GET", "/v1/config")
2367 .with_status(200)
2368 .with_body(format!(
2369 r#"{{"overrides": {{"uri": "{}"}}, "defaults": {{}}}}"#,
2370 overridden.url()
2371 ))
2372 .create_async()
2373 .await;
2374 let bootstrap_oauth_mock =
2376 create_oauth_mock_with_path(&mut bootstrap, "/v1/oauth/tokens", "tok-boot", 200).await;
2377
2378 let props = HashMap::from([("credential".to_string(), "client1:secret1".to_string())]);
2379 let catalog = RestCatalog::new(
2380 SessionContext::empty(),
2381 RestCatalogConfig::builder()
2382 .uri(bootstrap.url())
2383 .props(props)
2384 .build(),
2385 None,
2386 Some(Arc::new(LocalFsStorageFactory)),
2387 Runtime::current(),
2388 None,
2389 );
2390
2391 let client = catalog.client().await.unwrap();
2392 config_mock.assert_async().await;
2393 bootstrap_oauth_mock.assert_async().await;
2394 let session_debug = format!("{:?}", client.catalog_session);
2397 assert!(session_debug.contains(&format!("{}/v1/oauth/tokens", overridden.url())));
2398 }
2399
2400 #[tokio::test]
2401 async fn test_concurrent_authenticate_single_token_exchange() {
2402 let mut server = Server::new_async().await;
2405 let oauth_mock =
2407 create_oauth_mock_with_path(&mut server, "/v1/oauth/tokens", "tok-once", 200).await;
2408
2409 let manager = OAuth2Manager::new(format!("{}/v1/oauth/tokens", server.url()))
2410 .with_credential(Some("client1".to_string()), "secret1".to_string());
2411 let session: Arc<dyn AuthSession> = Arc::from(
2412 manager
2413 .init_session(&test_client(), &HashMap::new())
2414 .await
2415 .unwrap(),
2416 );
2417
2418 let client = Client::new();
2419 let attempts = (0..8).map(|_| {
2420 let session = session.clone();
2421 let client = client.clone();
2422 async move {
2423 let mut req = HttpRequest::new(
2424 client
2425 .get("https://rest.example.com/v1/config")
2426 .build()
2427 .unwrap(),
2428 );
2429 session.authenticate(&mut req).await.unwrap();
2430 req.headers()
2431 .get("authorization")
2432 .unwrap()
2433 .to_str()
2434 .unwrap()
2435 .to_string()
2436 }
2437 });
2438 let bearers = futures::future::join_all(attempts).await;
2439
2440 oauth_mock.assert_async().await;
2441 assert!(bearers.iter().all(|b| b == "Bearer tok-once"));
2442 }
2443
2444 #[tokio::test]
2445 async fn test_seeded_token_takes_precedence_over_credential() {
2446 let mut server = Server::new_async().await;
2449 let oauth_mock = server
2450 .mock("POST", "/v1/oauth/tokens")
2451 .expect(0)
2452 .create_async()
2453 .await;
2454
2455 let manager = OAuth2Manager::new(format!("{}/v1/oauth/tokens", server.url()))
2456 .with_token("tok-seed")
2457 .with_credential(Some("client1".to_string()), "secret1".to_string());
2458 let session = manager
2459 .init_session(&test_client(), &HashMap::new())
2460 .await
2461 .unwrap();
2462
2463 let mut req = HttpRequest::new(
2464 Client::new()
2465 .get("https://rest.example.com/v1/config")
2466 .build()
2467 .unwrap(),
2468 );
2469 session.authenticate(&mut req).await.unwrap();
2470 assert_eq!(
2471 req.headers().get("authorization").unwrap(),
2472 "Bearer tok-seed"
2473 );
2474
2475 oauth_mock.assert_async().await;
2476 }
2477
2478 #[tokio::test]
2479 async fn test_injected_oauth_manager_keeps_endpoint_and_options() {
2480 let mut server = Server::new_async().await;
2484 let config_mock = server
2487 .mock("GET", "/v1/config")
2488 .with_status(200)
2489 .with_body(r#"{"defaults": {"credential": "client1:secret1"}, "overrides": {}}"#)
2490 .create_async()
2491 .await;
2492
2493 let default_endpoint_mock = server
2495 .mock("POST", "/v1/oauth/tokens")
2496 .expect(0)
2497 .create_async()
2498 .await;
2499 let custom_endpoint_mock = server
2502 .mock("POST", "/custom/oauth/tokens")
2503 .match_header("x-tenant", "t1")
2504 .match_body(mockito::Matcher::AllOf(vec![
2507 mockito::Matcher::Regex("scope=catalog".to_string()),
2508 mockito::Matcher::Regex("audience=aud-1".to_string()),
2509 ]))
2510 .with_status(200)
2511 .with_body(
2512 r#"{
2513 "access_token": "ey000000000000",
2514 "token_type": "Bearer",
2515 "issued_token_type": "urn:ietf:params:oauth:token-type:access_token",
2516 "expires_in": 86400
2517 }"#,
2518 )
2519 .create_async()
2520 .await;
2521 let ns_mock = server
2522 .mock("GET", "/v1/namespaces")
2523 .match_header("authorization", "Bearer ey000000000000")
2524 .with_body(r#"{"namespaces": []}"#)
2525 .create_async()
2526 .await;
2527
2528 let manager = OAuth2Manager::new(format!("{}/custom/oauth/tokens", server.url()))
2529 .with_extra_headers(HeaderMap::from_iter([(
2530 HeaderName::from_static("x-tenant"),
2531 HeaderValue::from_static("t1"),
2532 )]))
2533 .with_extra_oauth_params(HashMap::from([(
2534 "audience".to_string(),
2535 "aud-1".to_string(),
2536 )]));
2537 let catalog = RestCatalog::new(
2538 SessionContext::empty(),
2539 RestCatalogConfig::builder().uri(server.url()).build(),
2540 Some(Box::new(manager)),
2541 Some(Arc::new(LocalFsStorageFactory)),
2542 Runtime::current(),
2543 None,
2544 );
2545
2546 catalog.list_namespaces(None).await.unwrap();
2547
2548 config_mock.assert_async().await;
2549 custom_endpoint_mock.assert_async().await;
2550 default_endpoint_mock.assert_async().await;
2551 ns_mock.assert_async().await;
2552 }
2553
2554 #[tokio::test]
2555 async fn test_props_token_takes_precedence_over_props_credential() {
2556 let mut server = Server::new_async().await;
2558 let oauth_mock = server
2559 .mock("POST", "/v1/oauth/tokens")
2560 .expect(0)
2561 .create_async()
2562 .await;
2563
2564 let manager = OAuth2Manager::new(format!("{}/v1/oauth/tokens", server.url()));
2565 let session = manager
2566 .init_session(
2567 &test_client(),
2568 &HashMap::from([
2569 ("token".to_string(), "tok-props".to_string()),
2570 ("credential".to_string(), "client1:secret1".to_string()),
2571 ]),
2572 )
2573 .await
2574 .unwrap();
2575
2576 let mut req = HttpRequest::new(
2577 Client::new()
2578 .get("https://rest.example.com/v1/config")
2579 .build()
2580 .unwrap(),
2581 );
2582 session.authenticate(&mut req).await.unwrap();
2583 assert_eq!(
2584 req.headers().get("authorization").unwrap(),
2585 "Bearer tok-props"
2586 );
2587 oauth_mock.assert_async().await;
2588 }
2589
2590 #[tokio::test]
2591 async fn test_manager_exchanges_over_the_catalog_client() {
2592 let mut server = Server::new_async().await;
2596 let token_mock = server
2597 .mock("POST", "/v1/oauth/tokens")
2598 .match_header("x-from", "manager")
2599 .match_header("x-catalog-only", mockito::Matcher::Missing)
2600 .with_status(200)
2601 .with_body(r#"{"access_token": "tok", "token_type": "Bearer"}"#)
2602 .create_async()
2603 .await;
2604
2605 let catalog_client = HttpClient::new(
2606 &RestCatalogConfig::builder()
2607 .uri(server.url())
2608 .props(HashMap::from([(
2609 "header.x-catalog-only".to_string(),
2610 "not-on-token-requests".to_string(),
2611 )]))
2612 .build(),
2613 )
2614 .unwrap();
2615
2616 let manager = OAuth2Manager::new(format!("{}/v1/oauth/tokens", server.url()))
2617 .with_credential(Some("client1".to_string()), "secret1".to_string())
2618 .with_extra_headers(HeaderMap::from_iter([(
2619 HeaderName::from_static("x-from"),
2620 HeaderValue::from_static("manager"),
2621 )]));
2622 let session = manager
2623 .init_session(&catalog_client, &HashMap::new())
2624 .await
2625 .unwrap();
2626
2627 let mut req = HttpRequest::new(
2628 Client::new()
2629 .get("https://rest.example.com/v1/namespaces")
2630 .build()
2631 .unwrap(),
2632 );
2633 session.authenticate(&mut req).await.unwrap();
2634 token_mock.assert_async().await;
2635 }
2636
2637 #[tokio::test]
2638 async fn test_handshake_is_authenticated_by_the_init_session() {
2639 let mut server = Server::new_async().await;
2642 let config_mock = server
2643 .mock("GET", "/v1/config")
2644 .match_header("authorization", "Bearer tok-init")
2645 .with_status(200)
2646 .with_body(r#"{"defaults": {}, "overrides": {}}"#)
2647 .create_async()
2648 .await;
2649
2650 let catalog = RestCatalog::new(
2651 SessionContext::empty(),
2652 RestCatalogConfig::builder()
2653 .uri(server.url())
2654 .props(HashMap::from([(
2655 "token".to_string(),
2656 "tok-init".to_string(),
2657 )]))
2658 .build(),
2659 None,
2660 Some(Arc::new(LocalFsStorageFactory)),
2661 Runtime::current(),
2662 None,
2663 );
2664
2665 catalog.client().await.unwrap();
2666 config_mock.assert_async().await;
2667 }
2668
2669 #[tokio::test]
2670 async fn test_init_session_receives_user_props() {
2671 use tokio::sync::Mutex as AsyncMutex;
2672
2673 #[derive(Debug)]
2676 struct PlainSession;
2677 #[async_trait]
2678 impl AuthSession for PlainSession {
2679 async fn authenticate(&self, _request: &mut HttpRequest) -> Result<()> {
2680 Ok(())
2681 }
2682 }
2683
2684 #[derive(Debug)]
2685 struct CapturingManager(Arc<AsyncMutex<Option<HashMap<String, String>>>>);
2686 #[async_trait]
2687 impl AuthManager for CapturingManager {
2688 async fn init_session(
2689 &self,
2690 _client: &HttpClient,
2691 props: &HashMap<String, String>,
2692 ) -> Result<Box<dyn AuthSession>> {
2693 *self.0.lock().await = Some(props.clone());
2694 Ok(Box::new(PlainSession))
2695 }
2696 async fn catalog_session(
2697 &self,
2698 _client: &HttpClient,
2699 _props: &HashMap<String, String>,
2700 ) -> Result<Arc<dyn AuthSession>> {
2701 Ok(Arc::new(PlainSession))
2702 }
2703 }
2704
2705 let mut server = Server::new_async().await;
2706 let config_mock = create_config_mock(&mut server).await;
2707 let captured = Arc::new(AsyncMutex::new(None));
2708 let catalog = RestCatalog::new(
2709 SessionContext::empty(),
2710 RestCatalogConfig::builder()
2711 .uri(server.url())
2712 .props(HashMap::from([(
2713 "token".to_string(),
2714 "tok-user".to_string(),
2715 )]))
2716 .build(),
2717 Some(Box::new(CapturingManager(captured.clone()))),
2718 Some(Arc::new(LocalFsStorageFactory)),
2719 Runtime::current(),
2720 None,
2721 );
2722
2723 catalog.client().await.unwrap();
2724 config_mock.assert_async().await;
2725 let props = captured.lock().await.clone().unwrap();
2726 assert_eq!(props.get("token").map(String::as_str), Some("tok-user"));
2727 assert_eq!(
2728 props.get(REST_CATALOG_PROP_URI).map(String::as_str),
2729 Some(server.url().as_str())
2730 );
2731 }
2732
2733 #[tokio::test]
2734 async fn test_catalog_session_receives_resolved_warehouse() {
2735 use tokio::sync::Mutex as AsyncMutex;
2736
2737 #[derive(Debug)]
2741 struct PlainSession;
2742 #[async_trait]
2743 impl AuthSession for PlainSession {
2744 async fn authenticate(&self, _request: &mut HttpRequest) -> Result<()> {
2745 Ok(())
2746 }
2747 }
2748
2749 #[derive(Debug)]
2750 struct CapturingManager(Arc<AsyncMutex<Option<HashMap<String, String>>>>);
2751 #[async_trait]
2752 impl AuthManager for CapturingManager {
2753 async fn init_session(
2754 &self,
2755 _client: &HttpClient,
2756 _props: &HashMap<String, String>,
2757 ) -> Result<Box<dyn AuthSession>> {
2758 Ok(Box::new(PlainSession))
2759 }
2760 async fn catalog_session(
2761 &self,
2762 _client: &HttpClient,
2763 props: &HashMap<String, String>,
2764 ) -> Result<Arc<dyn AuthSession>> {
2765 *self.0.lock().await = Some(props.clone());
2766 Ok(Arc::new(PlainSession))
2767 }
2768 }
2769
2770 let mut server = Server::new_async().await;
2772 let config_mock = server
2773 .mock("GET", "/v1/config")
2774 .match_query(mockito::Matcher::UrlEncoded(
2775 "warehouse".to_string(),
2776 "client-wh".to_string(),
2777 ))
2778 .with_status(200)
2779 .with_body(r#"{"defaults": {"warehouse": "default-wh"}, "overrides": {}}"#)
2780 .create_async()
2781 .await;
2782 let captured = Arc::new(AsyncMutex::new(None));
2783 let catalog = RestCatalog::new(
2784 SessionContext::empty(),
2785 RestCatalogConfig::builder()
2786 .uri(server.url())
2787 .warehouse("client-wh".to_string())
2788 .build(),
2789 Some(Box::new(CapturingManager(captured.clone()))),
2790 Some(Arc::new(LocalFsStorageFactory)),
2791 Runtime::current(),
2792 None,
2793 );
2794 catalog.client().await.unwrap();
2795 config_mock.assert_async().await;
2796 let props = captured.lock().await.clone().unwrap();
2797 assert_eq!(
2798 props.get("warehouse").map(String::as_str),
2799 Some("client-wh")
2800 );
2801
2802 let mut server = Server::new_async().await;
2804 let config_mock = server
2805 .mock("GET", "/v1/config")
2806 .match_query(mockito::Matcher::UrlEncoded(
2807 "warehouse".to_string(),
2808 "client-wh".to_string(),
2809 ))
2810 .with_status(200)
2811 .with_body(r#"{"defaults": {}, "overrides": {"warehouse": "override-wh"}}"#)
2812 .create_async()
2813 .await;
2814 let captured = Arc::new(AsyncMutex::new(None));
2815 let catalog = RestCatalog::new(
2816 SessionContext::empty(),
2817 RestCatalogConfig::builder()
2818 .uri(server.url())
2819 .warehouse("client-wh".to_string())
2820 .build(),
2821 Some(Box::new(CapturingManager(captured.clone()))),
2822 Some(Arc::new(LocalFsStorageFactory)),
2823 Runtime::current(),
2824 None,
2825 );
2826 catalog.client().await.unwrap();
2827 config_mock.assert_async().await;
2828 let props = captured.lock().await.clone().unwrap();
2829 assert_eq!(
2830 props.get("warehouse").map(String::as_str),
2831 Some("override-wh")
2832 );
2833 }
2834
2835 #[tokio::test]
2836 async fn test_init_session_dropped_before_catalog_session() {
2837 use std::sync::atomic::{AtomicBool, Ordering};
2838
2839 #[derive(Debug)]
2842 struct GuardSession(Arc<AtomicBool>);
2843 impl Drop for GuardSession {
2844 fn drop(&mut self) {
2845 self.0.store(true, Ordering::SeqCst);
2846 }
2847 }
2848 #[async_trait]
2849 impl AuthSession for GuardSession {
2850 async fn authenticate(&self, _request: &mut HttpRequest) -> Result<()> {
2851 Ok(())
2852 }
2853 }
2854
2855 #[derive(Debug)]
2856 struct PlainSession;
2857 #[async_trait]
2858 impl AuthSession for PlainSession {
2859 async fn authenticate(&self, _request: &mut HttpRequest) -> Result<()> {
2860 Ok(())
2861 }
2862 }
2863
2864 #[derive(Debug)]
2865 struct GuardManager(Arc<AtomicBool>);
2866 #[async_trait]
2867 impl AuthManager for GuardManager {
2868 async fn init_session(
2869 &self,
2870 _client: &HttpClient,
2871 _props: &HashMap<String, String>,
2872 ) -> Result<Box<dyn AuthSession>> {
2873 Ok(Box::new(GuardSession(self.0.clone())))
2874 }
2875 async fn catalog_session(
2876 &self,
2877 _client: &HttpClient,
2878 _props: &HashMap<String, String>,
2879 ) -> Result<Arc<dyn AuthSession>> {
2880 if !self.0.load(Ordering::SeqCst) {
2881 return Err(Error::new(
2882 ErrorKind::Unexpected,
2883 "init session must be dropped before catalog_session",
2884 ));
2885 }
2886 Ok(Arc::new(PlainSession))
2887 }
2888 }
2889
2890 let mut server = Server::new_async().await;
2891 let config_mock = create_config_mock(&mut server).await;
2892
2893 let dropped = Arc::new(AtomicBool::new(false));
2894 let catalog = RestCatalog::new(
2895 SessionContext::empty(),
2896 RestCatalogConfig::builder().uri(server.url()).build(),
2897 Some(Box::new(GuardManager(dropped.clone()))),
2898 Some(Arc::new(LocalFsStorageFactory)),
2899 Runtime::current(),
2900 None,
2901 );
2902
2903 catalog.client().await.unwrap();
2904 config_mock.assert_async().await;
2905 assert!(dropped.load(Ordering::SeqCst));
2906 }
2907
2908 #[tokio::test]
2909 async fn test_contextual_session_authenticates_each_catalog_request() {
2910 let mut server = Server::new_async().await;
2911 let config_mock = create_config_mock(&mut server).await;
2912 let first_page = server
2913 .mock("GET", "/v1/namespaces")
2914 .match_header("x-session-id", "session-123")
2915 .with_body(r#"{"namespaces": [["ns1"]], "next-page-token": "next"}"#)
2916 .create_async()
2917 .await;
2918 let second_page = server
2919 .mock("GET", "/v1/namespaces?pageToken=next")
2920 .match_header("x-session-id", "session-123")
2921 .with_body(r#"{"namespaces": [["ns2"]]}"#)
2922 .create_async()
2923 .await;
2924
2925 let (catalog, seen_session_ids) = context_catalog(
2926 RestCatalogConfig::builder().uri(server.url()).build(),
2927 false,
2928 );
2929 let context = SessionContext::builder()
2930 .session_id("session-123".to_string())
2931 .build();
2932
2933 let namespaces = catalog.list_namespaces(&context, None).await.unwrap();
2934
2935 assert_eq!(namespaces, vec![
2936 NamespaceIdent::new("ns1".to_string()),
2937 NamespaceIdent::new("ns2".to_string()),
2938 ]);
2939 assert_eq!(*seen_session_ids.lock().unwrap(), vec![
2940 "session-123",
2941 "session-123"
2942 ]);
2943 config_mock.assert_async().await;
2944 first_page.assert_async().await;
2945 second_page.assert_async().await;
2946 }
2947
2948 #[tokio::test]
2949 async fn test_contextual_session_authenticates_namespace_exists_head() {
2950 let mut server = Server::new_async().await;
2951 let config_mock = create_config_mock_with_exists_endpoints(&mut server).await;
2952 let exists_mock = server
2953 .mock("HEAD", "/v1/namespaces/ns1")
2954 .match_header("x-session-id", "namespace-session")
2955 .with_status(204)
2956 .create_async()
2957 .await;
2958 let (catalog, seen_session_ids) = context_catalog(
2959 RestCatalogConfig::builder().uri(server.url()).build(),
2960 false,
2961 );
2962 let context = SessionContext::builder()
2963 .session_id("namespace-session".to_string())
2964 .build();
2965
2966 assert!(
2967 catalog
2968 .namespace_exists(&context, &NamespaceIdent::new("ns1".to_string()))
2969 .await
2970 .unwrap()
2971 );
2972 assert_eq!(*seen_session_ids.lock().unwrap(), vec!["namespace-session"]);
2973 config_mock.assert_async().await;
2974 exists_mock.assert_async().await;
2975 }
2976
2977 #[tokio::test]
2978 async fn test_contextual_session_authenticates_table_exists_head() {
2979 let mut server = Server::new_async().await;
2980 let config_mock = create_config_mock_with_exists_endpoints(&mut server).await;
2981 let exists_mock = server
2982 .mock("HEAD", "/v1/namespaces/ns1/tables/table1")
2983 .match_header("x-session-id", "table-session")
2984 .with_status(204)
2985 .create_async()
2986 .await;
2987 let (catalog, seen_session_ids) = context_catalog(
2988 RestCatalogConfig::builder().uri(server.url()).build(),
2989 false,
2990 );
2991 let context = SessionContext::builder()
2992 .session_id("table-session".to_string())
2993 .build();
2994
2995 assert!(
2996 catalog
2997 .table_exists(
2998 &context,
2999 &TableIdent::new(NamespaceIdent::new("ns1".to_string()), "table1".to_string(),),
3000 )
3001 .await
3002 .unwrap()
3003 );
3004 assert_eq!(*seen_session_ids.lock().unwrap(), vec!["table-session"]);
3005 config_mock.assert_async().await;
3006 exists_mock.assert_async().await;
3007 }
3008
3009 #[tokio::test]
3010 async fn test_contextual_session_authenticates_write_operation() {
3011 let mut server = Server::new_async().await;
3012 let config_mock = create_config_mock(&mut server).await;
3013 let drop_mock = server
3014 .mock("DELETE", "/v1/namespaces/ns1")
3015 .match_header("x-session-id", "write-session")
3016 .with_status(204)
3017 .create_async()
3018 .await;
3019 let (catalog, seen_session_ids) = context_catalog(
3020 RestCatalogConfig::builder().uri(server.url()).build(),
3021 false,
3022 );
3023 let context = SessionContext::builder()
3024 .session_id("write-session".to_string())
3025 .build();
3026
3027 catalog
3028 .drop_namespace(&context, &NamespaceIdent::new("ns1".to_string()))
3029 .await
3030 .unwrap();
3031 assert_eq!(*seen_session_ids.lock().unwrap(), vec!["write-session"]);
3032 config_mock.assert_async().await;
3033 drop_mock.assert_async().await;
3034 }
3035
3036 #[tokio::test]
3037 async fn test_contextual_session_error_prevents_operation_request() {
3038 let mut server = Server::new_async().await;
3039 let config_mock = create_config_mock(&mut server).await;
3040 let list_mock = server
3041 .mock("GET", "/v1/namespaces")
3042 .expect(0)
3043 .create_async()
3044 .await;
3045 let (catalog, seen_session_ids) =
3046 context_catalog(RestCatalogConfig::builder().uri(server.url()).build(), true);
3047 let context = SessionContext::builder()
3048 .session_id("failing-session".to_string())
3049 .build();
3050
3051 let error = catalog.list_namespaces(&context, None).await.unwrap_err();
3052 assert_eq!(error.kind(), ErrorKind::Unexpected);
3053 assert_eq!(error.message(), "contextual session failure");
3054 assert_eq!(*seen_session_ids.lock().unwrap(), vec!["failing-session"]);
3055 config_mock.assert_async().await;
3056 list_mock.assert_async().await;
3057 }
3058
3059 #[test]
3060 fn test_config_debug_redacts_secrets() {
3061 let config = RestCatalogConfig::builder()
3062 .uri("http://localhost".to_string())
3063 .props(HashMap::from([
3064 ("token".to_string(), "tok-secret".to_string()),
3065 ("credential".to_string(), "id:cred-secret".to_string()),
3066 ("header.authorization".to_string(), "Basic xyz".to_string()),
3067 ("adls.account-key".to_string(), "adls-secret".to_string()),
3068 ("s3.sse.key".to_string(), "sse-secret".to_string()),
3069 (
3070 "adls.connection-string".to_string(),
3071 "cs-secret".to_string(),
3072 ),
3073 ("warehouse".to_string(), "wh1".to_string()),
3074 ]))
3075 .build();
3076
3077 let out = format!("{config:?}");
3078 assert!(!out.contains("tok-secret"));
3079 assert!(!out.contains("cred-secret"));
3080 assert!(!out.contains("Basic xyz"));
3081 assert!(!out.contains("adls-secret"));
3082 assert!(!out.contains("sse-secret"));
3083 assert!(!out.contains("cs-secret"));
3084 assert!(out.contains("[REDACTED]"));
3085 assert!(out.contains("wh1"));
3086 }
3087
3088 #[tokio::test]
3089 async fn test_auth_type_defaults() {
3090 let bare = RestCatalogConfig::builder()
3093 .uri("http://localhost".to_string())
3094 .build();
3095 assert!(
3096 format!("{:?}", test_catalog(bare).resolve_auth_manager().unwrap())
3097 .contains("NoopAuthManager")
3098 );
3099
3100 let with_token = RestCatalogConfig::builder()
3101 .uri("http://localhost".to_string())
3102 .props(HashMap::from([("token".to_string(), "tok".to_string())]))
3103 .build();
3104 assert!(
3105 format!(
3106 "{:?}",
3107 test_catalog(with_token).resolve_auth_manager().unwrap()
3108 )
3109 .contains("OAuth2Manager")
3110 );
3111
3112 let mixed_case = RestCatalogConfig::builder()
3114 .uri("http://localhost".to_string())
3115 .props(HashMap::from([(
3116 REST_CATALOG_PROP_AUTH_TYPE.to_string(),
3117 "OAuth2".to_string(),
3118 )]))
3119 .build();
3120 assert!(
3121 format!(
3122 "{:?}",
3123 test_catalog(mixed_case).resolve_auth_manager().unwrap()
3124 )
3125 .contains("OAuth2Manager")
3126 );
3127
3128 let with_endpoint = RestCatalogConfig::builder()
3131 .uri("http://localhost".to_string())
3132 .props(HashMap::from([(
3133 "oauth2-server-uri".to_string(),
3134 "http://auth.example.com/tokens".to_string(),
3135 )]))
3136 .build();
3137 assert!(
3138 format!(
3139 "{:?}",
3140 test_catalog(with_endpoint).resolve_auth_manager().unwrap()
3141 )
3142 .contains("OAuth2Manager")
3143 );
3144 }
3145
3146 #[tokio::test]
3147 async fn test_unknown_auth_type_is_rejected() {
3148 let props = HashMap::from([(
3149 REST_CATALOG_PROP_AUTH_TYPE.to_string(),
3150 "kerberos".to_string(),
3151 )]);
3152 let config = RestCatalogConfig::builder()
3153 .uri("http://localhost".to_string())
3154 .props(props)
3155 .build();
3156
3157 let err = test_catalog(config).resolve_auth_manager().unwrap_err();
3158 assert!(err.message().contains(REST_CATALOG_PROP_AUTH_TYPE));
3159 }
3160
3161 #[tokio::test]
3162 async fn test_with_auth_manager_overrides_config() {
3163 #[derive(Debug)]
3165 struct StubAuthManager;
3166 #[async_trait]
3167 impl AuthManager for StubAuthManager {
3168 async fn init_session(
3169 &self,
3170 _client: &HttpClient,
3171 _props: &HashMap<String, String>,
3172 ) -> Result<Box<dyn AuthSession>> {
3173 unimplemented!()
3174 }
3175 async fn catalog_session(
3176 &self,
3177 _client: &HttpClient,
3178 _props: &HashMap<String, String>,
3179 ) -> Result<Arc<dyn AuthSession>> {
3180 unimplemented!()
3181 }
3182 }
3183
3184 let config = RestCatalogConfig::builder()
3185 .uri("http://localhost".to_string())
3186 .props(HashMap::from([(
3187 REST_CATALOG_PROP_AUTH_TYPE.to_string(),
3188 "kerberos".to_string(),
3189 )]))
3190 .build();
3191
3192 let catalog = test_catalog_with(config, StubAuthManager);
3194 assert!(catalog.resolve_auth_manager().is_ok());
3195 }
3196
3197 #[tokio::test]
3198 async fn test_list_namespace_with_pagination() {
3199 let mut server = Server::new_async().await;
3200
3201 let config_mock = create_config_mock(&mut server).await;
3202
3203 let list_ns_mock_page1 = server
3204 .mock("GET", "/v1/namespaces")
3205 .with_body(
3206 r#"{
3207 "namespaces": [
3208 ["ns1", "ns11"],
3209 ["ns2"]
3210 ],
3211 "next-page-token": "token123"
3212 }"#,
3213 )
3214 .create_async()
3215 .await;
3216
3217 let list_ns_mock_page2 = server
3218 .mock("GET", "/v1/namespaces?pageToken=token123")
3219 .with_body(
3220 r#"{
3221 "namespaces": [
3222 ["ns3"],
3223 ["ns4", "ns41"]
3224 ]
3225 }"#,
3226 )
3227 .create_async()
3228 .await;
3229
3230 let catalog = session_catalog(RestCatalogConfig::builder().uri(server.url()).build());
3231
3232 let namespaces = catalog
3233 .list_namespaces(&SessionContext::empty(), None)
3234 .await
3235 .unwrap();
3236
3237 let expected_ns = vec![
3238 NamespaceIdent::from_vec(vec!["ns1".to_string(), "ns11".to_string()]).unwrap(),
3239 NamespaceIdent::from_vec(vec!["ns2".to_string()]).unwrap(),
3240 NamespaceIdent::from_vec(vec!["ns3".to_string()]).unwrap(),
3241 NamespaceIdent::from_vec(vec!["ns4".to_string(), "ns41".to_string()]).unwrap(),
3242 ];
3243
3244 assert_eq!(expected_ns, namespaces);
3245
3246 config_mock.assert_async().await;
3247 list_ns_mock_page1.assert_async().await;
3248 list_ns_mock_page2.assert_async().await;
3249 }
3250
3251 #[tokio::test]
3252 async fn test_list_namespace_with_multiple_pages() {
3253 let mut server = Server::new_async().await;
3254
3255 let config_mock = create_config_mock(&mut server).await;
3256
3257 let list_ns_mock_page1 = server
3259 .mock("GET", "/v1/namespaces")
3260 .with_body(
3261 r#"{
3262 "namespaces": [
3263 ["ns1", "ns11"],
3264 ["ns2"]
3265 ],
3266 "next-page-token": "page2"
3267 }"#,
3268 )
3269 .create_async()
3270 .await;
3271
3272 let list_ns_mock_page2 = server
3274 .mock("GET", "/v1/namespaces?pageToken=page2")
3275 .with_body(
3276 r#"{
3277 "namespaces": [
3278 ["ns3"],
3279 ["ns4", "ns41"]
3280 ],
3281 "next-page-token": "page3"
3282 }"#,
3283 )
3284 .create_async()
3285 .await;
3286
3287 let list_ns_mock_page3 = server
3289 .mock("GET", "/v1/namespaces?pageToken=page3")
3290 .with_body(
3291 r#"{
3292 "namespaces": [
3293 ["ns5", "ns51", "ns511"]
3294 ],
3295 "next-page-token": "page4"
3296 }"#,
3297 )
3298 .create_async()
3299 .await;
3300
3301 let list_ns_mock_page4 = server
3303 .mock("GET", "/v1/namespaces?pageToken=page4")
3304 .with_body(
3305 r#"{
3306 "namespaces": [
3307 ["ns6"],
3308 ["ns7"]
3309 ],
3310 "next-page-token": "page5"
3311 }"#,
3312 )
3313 .create_async()
3314 .await;
3315
3316 let list_ns_mock_page5 = server
3318 .mock("GET", "/v1/namespaces?pageToken=page5")
3319 .with_body(
3320 r#"{
3321 "namespaces": [
3322 ["ns8", "ns81"]
3323 ]
3324 }"#,
3325 )
3326 .create_async()
3327 .await;
3328
3329 let catalog = session_catalog(RestCatalogConfig::builder().uri(server.url()).build());
3330
3331 let namespaces = catalog
3332 .list_namespaces(&SessionContext::empty(), None)
3333 .await
3334 .unwrap();
3335
3336 let expected_ns = vec![
3337 NamespaceIdent::from_vec(vec!["ns1".to_string(), "ns11".to_string()]).unwrap(),
3338 NamespaceIdent::from_vec(vec!["ns2".to_string()]).unwrap(),
3339 NamespaceIdent::from_vec(vec!["ns3".to_string()]).unwrap(),
3340 NamespaceIdent::from_vec(vec!["ns4".to_string(), "ns41".to_string()]).unwrap(),
3341 NamespaceIdent::from_vec(vec![
3342 "ns5".to_string(),
3343 "ns51".to_string(),
3344 "ns511".to_string(),
3345 ])
3346 .unwrap(),
3347 NamespaceIdent::from_vec(vec!["ns6".to_string()]).unwrap(),
3348 NamespaceIdent::from_vec(vec!["ns7".to_string()]).unwrap(),
3349 NamespaceIdent::from_vec(vec!["ns8".to_string(), "ns81".to_string()]).unwrap(),
3350 ];
3351
3352 assert_eq!(expected_ns, namespaces);
3353
3354 config_mock.assert_async().await;
3356 list_ns_mock_page1.assert_async().await;
3357 list_ns_mock_page2.assert_async().await;
3358 list_ns_mock_page3.assert_async().await;
3359 list_ns_mock_page4.assert_async().await;
3360 list_ns_mock_page5.assert_async().await;
3361 }
3362
3363 #[tokio::test]
3364 async fn test_create_namespace() {
3365 let mut server = Server::new_async().await;
3366
3367 let config_mock = create_config_mock(&mut server).await;
3368
3369 let create_ns_mock = server
3370 .mock("POST", "/v1/namespaces")
3371 .with_body(
3372 r#"{
3373 "namespace": [ "ns1", "ns11"],
3374 "properties" : {
3375 "key1": "value1"
3376 }
3377 }"#,
3378 )
3379 .create_async()
3380 .await;
3381
3382 let catalog = session_catalog(RestCatalogConfig::builder().uri(server.url()).build());
3383
3384 let namespaces = catalog
3385 .create_namespace(
3386 &SessionContext::empty(),
3387 &NamespaceIdent::from_vec(vec!["ns1".to_string(), "ns11".to_string()]).unwrap(),
3388 HashMap::from([("key1".to_string(), "value1".to_string())]),
3389 )
3390 .await
3391 .unwrap();
3392
3393 let expected_ns = Namespace::with_properties(
3394 NamespaceIdent::from_vec(vec!["ns1".to_string(), "ns11".to_string()]).unwrap(),
3395 HashMap::from([("key1".to_string(), "value1".to_string())]),
3396 );
3397
3398 assert_eq!(expected_ns, namespaces);
3399
3400 config_mock.assert_async().await;
3401 create_ns_mock.assert_async().await;
3402 }
3403
3404 #[tokio::test]
3405 async fn test_get_namespace() {
3406 let mut server = Server::new_async().await;
3407
3408 let config_mock = create_config_mock(&mut server).await;
3409
3410 let get_ns_mock = server
3411 .mock("GET", "/v1/namespaces/ns1")
3412 .with_body(
3413 r#"{
3414 "namespace": [ "ns1"],
3415 "properties" : {
3416 "key1": "value1"
3417 }
3418 }"#,
3419 )
3420 .create_async()
3421 .await;
3422
3423 let catalog = session_catalog(RestCatalogConfig::builder().uri(server.url()).build());
3424
3425 let namespaces = catalog
3426 .get_namespace(
3427 &SessionContext::empty(),
3428 &NamespaceIdent::new("ns1".to_string()),
3429 )
3430 .await
3431 .unwrap();
3432
3433 let expected_ns = Namespace::with_properties(
3434 NamespaceIdent::new("ns1".to_string()),
3435 HashMap::from([("key1".to_string(), "value1".to_string())]),
3436 );
3437
3438 assert_eq!(expected_ns, namespaces);
3439
3440 config_mock.assert_async().await;
3441 get_ns_mock.assert_async().await;
3442 }
3443
3444 #[tokio::test]
3445 async fn check_namespace_exists() {
3446 let mut server = Server::new_async().await;
3447
3448 let config_mock = create_config_mock_with_exists_endpoints(&mut server).await;
3449
3450 let get_ns_mock = server
3451 .mock("HEAD", "/v1/namespaces/ns1")
3452 .with_status(204)
3453 .create_async()
3454 .await;
3455
3456 let catalog = session_catalog(RestCatalogConfig::builder().uri(server.url()).build());
3457
3458 assert!(
3459 catalog
3460 .namespace_exists(
3461 &SessionContext::empty(),
3462 &NamespaceIdent::new("ns1".to_string())
3463 )
3464 .await
3465 .unwrap()
3466 );
3467
3468 config_mock.assert_async().await;
3469 get_ns_mock.assert_async().await;
3470 }
3471
3472 #[tokio::test]
3473 async fn test_namespace_exists_falls_back_to_get_when_head_not_advertised() {
3474 let mut server = Server::new_async().await;
3475
3476 let config_mock = create_config_mock(&mut server).await;
3479 let get_ns_mock = server
3480 .mock("GET", "/v1/namespaces/ns1")
3481 .with_status(200)
3482 .with_body(
3483 r#"{
3484 "namespace": ["ns1"],
3485 "properties": {}
3486 }"#,
3487 )
3488 .create_async()
3489 .await;
3490
3491 let catalog = RestCatalog::new(
3492 SessionContext::empty(),
3493 RestCatalogConfig::builder().uri(server.url()).build(),
3494 None,
3495 Some(Arc::new(LocalFsStorageFactory)),
3496 Runtime::current(),
3497 None,
3498 );
3499
3500 assert!(
3501 catalog
3502 .namespace_exists(&NamespaceIdent::new("ns1".to_string()))
3503 .await
3504 .unwrap()
3505 );
3506
3507 config_mock.assert_async().await;
3508 get_ns_mock.assert_async().await;
3509 }
3510
3511 #[tokio::test]
3512 async fn test_drop_namespace() {
3513 let mut server = Server::new_async().await;
3514
3515 let config_mock = create_config_mock(&mut server).await;
3516
3517 let drop_ns_mock = server
3518 .mock("DELETE", "/v1/namespaces/ns1")
3519 .with_status(204)
3520 .create_async()
3521 .await;
3522
3523 let catalog = session_catalog(RestCatalogConfig::builder().uri(server.url()).build());
3524
3525 catalog
3526 .drop_namespace(
3527 &SessionContext::empty(),
3528 &NamespaceIdent::new("ns1".to_string()),
3529 )
3530 .await
3531 .unwrap();
3532
3533 config_mock.assert_async().await;
3534 drop_ns_mock.assert_async().await;
3535 }
3536
3537 #[tokio::test]
3538 async fn test_list_tables() {
3539 let mut server = Server::new_async().await;
3540
3541 let config_mock = create_config_mock(&mut server).await;
3542
3543 let list_tables_mock = server
3544 .mock("GET", "/v1/namespaces/ns1/tables")
3545 .with_status(200)
3546 .with_body(
3547 r#"{
3548 "identifiers": [
3549 {
3550 "namespace": ["ns1"],
3551 "name": "table1"
3552 },
3553 {
3554 "namespace": ["ns1"],
3555 "name": "table2"
3556 }
3557 ]
3558 }"#,
3559 )
3560 .create_async()
3561 .await;
3562
3563 let catalog = session_catalog(RestCatalogConfig::builder().uri(server.url()).build());
3564
3565 let tables = catalog
3566 .list_tables(
3567 &SessionContext::empty(),
3568 &NamespaceIdent::new("ns1".to_string()),
3569 )
3570 .await
3571 .unwrap();
3572
3573 let expected_tables = vec![
3574 TableIdent::new(NamespaceIdent::new("ns1".to_string()), "table1".to_string()),
3575 TableIdent::new(NamespaceIdent::new("ns1".to_string()), "table2".to_string()),
3576 ];
3577
3578 assert_eq!(tables, expected_tables);
3579
3580 config_mock.assert_async().await;
3581 list_tables_mock.assert_async().await;
3582 }
3583
3584 #[tokio::test]
3585 async fn test_list_tables_with_pagination() {
3586 let mut server = Server::new_async().await;
3587
3588 let config_mock = create_config_mock(&mut server).await;
3589
3590 let list_tables_mock_page1 = server
3591 .mock("GET", "/v1/namespaces/ns1/tables")
3592 .with_status(200)
3593 .with_body(
3594 r#"{
3595 "identifiers": [
3596 {
3597 "namespace": ["ns1"],
3598 "name": "table1"
3599 },
3600 {
3601 "namespace": ["ns1"],
3602 "name": "table2"
3603 }
3604 ],
3605 "next-page-token": "token456"
3606 }"#,
3607 )
3608 .create_async()
3609 .await;
3610
3611 let list_tables_mock_page2 = server
3612 .mock("GET", "/v1/namespaces/ns1/tables?pageToken=token456")
3613 .with_status(200)
3614 .with_body(
3615 r#"{
3616 "identifiers": [
3617 {
3618 "namespace": ["ns1"],
3619 "name": "table3"
3620 },
3621 {
3622 "namespace": ["ns1"],
3623 "name": "table4"
3624 }
3625 ]
3626 }"#,
3627 )
3628 .create_async()
3629 .await;
3630
3631 let catalog = session_catalog(RestCatalogConfig::builder().uri(server.url()).build());
3632
3633 let tables = catalog
3634 .list_tables(
3635 &SessionContext::empty(),
3636 &NamespaceIdent::new("ns1".to_string()),
3637 )
3638 .await
3639 .unwrap();
3640
3641 let expected_tables = vec![
3642 TableIdent::new(NamespaceIdent::new("ns1".to_string()), "table1".to_string()),
3643 TableIdent::new(NamespaceIdent::new("ns1".to_string()), "table2".to_string()),
3644 TableIdent::new(NamespaceIdent::new("ns1".to_string()), "table3".to_string()),
3645 TableIdent::new(NamespaceIdent::new("ns1".to_string()), "table4".to_string()),
3646 ];
3647
3648 assert_eq!(tables, expected_tables);
3649
3650 config_mock.assert_async().await;
3651 list_tables_mock_page1.assert_async().await;
3652 list_tables_mock_page2.assert_async().await;
3653 }
3654
3655 #[tokio::test]
3656 async fn test_list_tables_with_multiple_pages() {
3657 let mut server = Server::new_async().await;
3658
3659 let config_mock = create_config_mock(&mut server).await;
3660
3661 let list_tables_mock_page1 = server
3663 .mock("GET", "/v1/namespaces/ns1/tables")
3664 .with_status(200)
3665 .with_body(
3666 r#"{
3667 "identifiers": [
3668 {
3669 "namespace": ["ns1"],
3670 "name": "table1"
3671 },
3672 {
3673 "namespace": ["ns1"],
3674 "name": "table2"
3675 }
3676 ],
3677 "next-page-token": "page2"
3678 }"#,
3679 )
3680 .create_async()
3681 .await;
3682
3683 let list_tables_mock_page2 = server
3685 .mock("GET", "/v1/namespaces/ns1/tables?pageToken=page2")
3686 .with_status(200)
3687 .with_body(
3688 r#"{
3689 "identifiers": [
3690 {
3691 "namespace": ["ns1"],
3692 "name": "table3"
3693 },
3694 {
3695 "namespace": ["ns1"],
3696 "name": "table4"
3697 }
3698 ],
3699 "next-page-token": "page3"
3700 }"#,
3701 )
3702 .create_async()
3703 .await;
3704
3705 let list_tables_mock_page3 = server
3707 .mock("GET", "/v1/namespaces/ns1/tables?pageToken=page3")
3708 .with_status(200)
3709 .with_body(
3710 r#"{
3711 "identifiers": [
3712 {
3713 "namespace": ["ns1"],
3714 "name": "table5"
3715 }
3716 ],
3717 "next-page-token": "page4"
3718 }"#,
3719 )
3720 .create_async()
3721 .await;
3722
3723 let list_tables_mock_page4 = server
3725 .mock("GET", "/v1/namespaces/ns1/tables?pageToken=page4")
3726 .with_status(200)
3727 .with_body(
3728 r#"{
3729 "identifiers": [
3730 {
3731 "namespace": ["ns1"],
3732 "name": "table6"
3733 },
3734 {
3735 "namespace": ["ns1"],
3736 "name": "table7"
3737 }
3738 ],
3739 "next-page-token": "page5"
3740 }"#,
3741 )
3742 .create_async()
3743 .await;
3744
3745 let list_tables_mock_page5 = server
3747 .mock("GET", "/v1/namespaces/ns1/tables?pageToken=page5")
3748 .with_status(200)
3749 .with_body(
3750 r#"{
3751 "identifiers": [
3752 {
3753 "namespace": ["ns1"],
3754 "name": "table8"
3755 }
3756 ]
3757 }"#,
3758 )
3759 .create_async()
3760 .await;
3761
3762 let catalog = session_catalog(RestCatalogConfig::builder().uri(server.url()).build());
3763
3764 let tables = catalog
3765 .list_tables(
3766 &SessionContext::empty(),
3767 &NamespaceIdent::new("ns1".to_string()),
3768 )
3769 .await
3770 .unwrap();
3771
3772 let expected_tables = vec![
3773 TableIdent::new(NamespaceIdent::new("ns1".to_string()), "table1".to_string()),
3774 TableIdent::new(NamespaceIdent::new("ns1".to_string()), "table2".to_string()),
3775 TableIdent::new(NamespaceIdent::new("ns1".to_string()), "table3".to_string()),
3776 TableIdent::new(NamespaceIdent::new("ns1".to_string()), "table4".to_string()),
3777 TableIdent::new(NamespaceIdent::new("ns1".to_string()), "table5".to_string()),
3778 TableIdent::new(NamespaceIdent::new("ns1".to_string()), "table6".to_string()),
3779 TableIdent::new(NamespaceIdent::new("ns1".to_string()), "table7".to_string()),
3780 TableIdent::new(NamespaceIdent::new("ns1".to_string()), "table8".to_string()),
3781 ];
3782
3783 assert_eq!(tables, expected_tables);
3784
3785 config_mock.assert_async().await;
3787 list_tables_mock_page1.assert_async().await;
3788 list_tables_mock_page2.assert_async().await;
3789 list_tables_mock_page3.assert_async().await;
3790 list_tables_mock_page4.assert_async().await;
3791 list_tables_mock_page5.assert_async().await;
3792 }
3793
3794 #[tokio::test]
3795 async fn test_drop_tables() {
3796 let mut server = Server::new_async().await;
3797
3798 let config_mock = create_config_mock(&mut server).await;
3799
3800 let delete_table_mock = server
3801 .mock("DELETE", "/v1/namespaces/ns1/tables/table1")
3802 .with_status(204)
3803 .create_async()
3804 .await;
3805
3806 let catalog = session_catalog(RestCatalogConfig::builder().uri(server.url()).build());
3807
3808 catalog
3809 .drop_table(
3810 &SessionContext::empty(),
3811 &TableIdent::new(NamespaceIdent::new("ns1".to_string()), "table1".to_string()),
3812 )
3813 .await
3814 .unwrap();
3815
3816 config_mock.assert_async().await;
3817 delete_table_mock.assert_async().await;
3818 }
3819
3820 #[tokio::test]
3821 async fn test_check_table_exists() {
3822 let mut server = Server::new_async().await;
3823
3824 let config_mock = create_config_mock_with_exists_endpoints(&mut server).await;
3825
3826 let check_table_exists_mock = server
3827 .mock("HEAD", "/v1/namespaces/ns1/tables/table1")
3828 .with_status(204)
3829 .create_async()
3830 .await;
3831
3832 let catalog = session_catalog(RestCatalogConfig::builder().uri(server.url()).build());
3833
3834 assert!(
3835 catalog
3836 .table_exists(
3837 &SessionContext::empty(),
3838 &TableIdent::new(NamespaceIdent::new("ns1".to_string()), "table1".to_string(),),
3839 )
3840 .await
3841 .unwrap()
3842 );
3843
3844 config_mock.assert_async().await;
3845 check_table_exists_mock.assert_async().await;
3846 }
3847
3848 #[tokio::test]
3849 async fn test_table_exists_falls_back_to_load_when_head_not_advertised() {
3850 let mut server = Server::new_async().await;
3851
3852 let config_mock = create_config_mock(&mut server).await;
3855 let load_table_mock = server
3856 .mock("GET", "/v1/namespaces/ns1/tables/table1")
3857 .with_status(200)
3858 .with_body_from_file(format!(
3859 "{}/testdata/{}",
3860 env!("CARGO_MANIFEST_DIR"),
3861 "load_table_response.json"
3862 ))
3863 .create_async()
3864 .await;
3865
3866 let catalog = RestCatalog::new(
3867 SessionContext::empty(),
3868 RestCatalogConfig::builder().uri(server.url()).build(),
3869 None,
3870 Some(Arc::new(LocalFsStorageFactory)),
3871 Runtime::current(),
3872 None,
3873 );
3874
3875 assert!(
3876 catalog
3877 .table_exists(&TableIdent::new(
3878 NamespaceIdent::new("ns1".to_string()),
3879 "table1".to_string(),
3880 ))
3881 .await
3882 .unwrap()
3883 );
3884
3885 config_mock.assert_async().await;
3886 load_table_mock.assert_async().await;
3887 }
3888
3889 #[tokio::test]
3890 async fn test_rename_table() {
3891 let mut server = Server::new_async().await;
3892
3893 let config_mock = create_config_mock(&mut server).await;
3894
3895 let rename_table_mock = server
3896 .mock("POST", "/v1/tables/rename")
3897 .with_status(204)
3898 .create_async()
3899 .await;
3900
3901 let catalog = session_catalog(RestCatalogConfig::builder().uri(server.url()).build());
3902
3903 catalog
3904 .rename_table(
3905 &SessionContext::empty(),
3906 &TableIdent::new(NamespaceIdent::new("ns1".to_string()), "table1".to_string()),
3907 &TableIdent::new(NamespaceIdent::new("ns1".to_string()), "table2".to_string()),
3908 )
3909 .await
3910 .unwrap();
3911
3912 config_mock.assert_async().await;
3913 rename_table_mock.assert_async().await;
3914 }
3915
3916 #[tokio::test]
3917 async fn test_load_table() {
3918 let mut server = Server::new_async().await;
3919
3920 let config_mock = create_config_mock(&mut server).await;
3921
3922 let rename_table_mock = server
3923 .mock("GET", "/v1/namespaces/ns1/tables/test1")
3924 .with_status(200)
3925 .with_body_from_file(format!(
3926 "{}/testdata/{}",
3927 env!("CARGO_MANIFEST_DIR"),
3928 "load_table_response.json"
3929 ))
3930 .create_async()
3931 .await;
3932
3933 let catalog = session_catalog(RestCatalogConfig::builder().uri(server.url()).build());
3934
3935 let table = catalog
3936 .load_table(
3937 &SessionContext::empty(),
3938 &TableIdent::new(NamespaceIdent::new("ns1".to_string()), "test1".to_string()),
3939 )
3940 .await
3941 .unwrap();
3942
3943 assert_eq!(
3944 &TableIdent::from_strs(vec!["ns1", "test1"]).unwrap(),
3945 table.identifier()
3946 );
3947 assert_eq!(
3948 "s3://warehouse/database/table/metadata/00001-5f2f8166-244c-4eae-ac36-384ecdec81fc.gz.metadata.json",
3949 table.metadata_location().unwrap()
3950 );
3951 assert_eq!(FormatVersion::V1, table.metadata().format_version());
3952 assert_eq!("s3://warehouse/database/table", table.metadata().location());
3953 assert_eq!(
3954 uuid!("b55d9dda-6561-423a-8bfc-787980ce421f"),
3955 table.metadata().uuid()
3956 );
3957 assert_eq!(
3958 Utc.timestamp_millis_opt(1646787054459).unwrap(),
3959 table.metadata().last_updated_timestamp().unwrap()
3960 );
3961 assert_eq!(
3962 vec![&Arc::new(
3963 Schema::builder()
3964 .with_fields(vec![
3965 NestedField::optional(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
3966 NestedField::optional(2, "data", Type::Primitive(PrimitiveType::String))
3967 .into(),
3968 ])
3969 .build()
3970 .unwrap()
3971 )],
3972 table.metadata().schemas_iter().collect::<Vec<_>>()
3973 );
3974 assert_eq!(
3975 &HashMap::from([
3976 ("owner".to_string(), "bryan".to_string()),
3977 (
3978 "write.metadata.compression-codec".to_string(),
3979 "gzip".to_string()
3980 )
3981 ]),
3982 table.metadata().properties()
3983 );
3984 assert_eq!(vec![&Arc::new(Snapshot::builder()
3985 .with_snapshot_id(3497810964824022504)
3986 .with_timestamp_ms(1646787054459)
3987 .with_manifest_list("s3://warehouse/database/table/metadata/snap-3497810964824022504-1-c4f68204-666b-4e50-a9df-b10c34bf6b82.avro")
3988 .with_sequence_number(0)
3989 .with_schema_id(0)
3990 .with_summary(Summary {
3991 operation: Operation::Append,
3992 additional_properties: HashMap::from_iter([
3993 ("spark.app.id", "local-1646787004168"),
3994 ("added-data-files", "1"),
3995 ("added-records", "1"),
3996 ("added-files-size", "697"),
3997 ("changed-partition-count", "1"),
3998 ("total-records", "1"),
3999 ("total-files-size", "697"),
4000 ("total-data-files", "1"),
4001 ("total-delete-files", "0"),
4002 ("total-position-deletes", "0"),
4003 ("total-equality-deletes", "0")
4004 ].iter().map(|p| (p.0.to_string(), p.1.to_string()))),
4005 }).build()
4006 )], table.metadata().snapshots().collect::<Vec<_>>());
4007 assert_eq!(
4008 &[SnapshotLog {
4009 timestamp_ms: 1646787054459,
4010 snapshot_id: 3497810964824022504,
4011 }],
4012 table.metadata().history()
4013 );
4014 assert_eq!(
4015 vec![&Arc::new(SortOrder {
4016 order_id: 0,
4017 fields: vec![],
4018 })],
4019 table.metadata().sort_orders_iter().collect::<Vec<_>>()
4020 );
4021
4022 config_mock.assert_async().await;
4023 rename_table_mock.assert_async().await;
4024 }
4025
4026 #[tokio::test]
4027 async fn test_load_table_404() {
4028 let mut server = Server::new_async().await;
4029
4030 let config_mock = create_config_mock(&mut server).await;
4031
4032 let rename_table_mock = server
4033 .mock("GET", "/v1/namespaces/ns1/tables/test1")
4034 .with_status(404)
4035 .with_body(r#"
4036{
4037 "error": {
4038 "message": "Table does not exist: ns1.test1 in warehouse 8bcb0838-50fc-472d-9ddb-8feb89ef5f1e",
4039 "type": "NoSuchNamespaceErrorException",
4040 "code": 404
4041 }
4042}
4043 "#)
4044 .create_async()
4045 .await;
4046
4047 let catalog = session_catalog(RestCatalogConfig::builder().uri(server.url()).build());
4048
4049 let table = catalog
4050 .load_table(
4051 &SessionContext::empty(),
4052 &TableIdent::new(NamespaceIdent::new("ns1".to_string()), "test1".to_string()),
4053 )
4054 .await;
4055
4056 assert!(table.is_err());
4057 assert!(table.err().unwrap().message().contains("does not exist"));
4058
4059 config_mock.assert_async().await;
4060 rename_table_mock.assert_async().await;
4061 }
4062
4063 #[tokio::test]
4064 async fn test_create_table() {
4065 let mut server = Server::new_async().await;
4066
4067 let config_mock = create_config_mock(&mut server).await;
4068
4069 let create_table_mock = server
4070 .mock("POST", "/v1/namespaces/ns1/tables")
4071 .with_status(200)
4072 .with_body_from_file(format!(
4073 "{}/testdata/{}",
4074 env!("CARGO_MANIFEST_DIR"),
4075 "create_table_response.json"
4076 ))
4077 .create_async()
4078 .await;
4079
4080 let catalog = session_catalog(RestCatalogConfig::builder().uri(server.url()).build());
4081
4082 let table_creation = TableCreation::builder()
4083 .name("test1".to_string())
4084 .schema(
4085 Schema::builder()
4086 .with_fields(vec![
4087 NestedField::optional(1, "foo", Type::Primitive(PrimitiveType::String))
4088 .into(),
4089 NestedField::required(2, "bar", Type::Primitive(PrimitiveType::Int)).into(),
4090 NestedField::optional(3, "baz", Type::Primitive(PrimitiveType::Boolean))
4091 .into(),
4092 ])
4093 .with_schema_id(1)
4094 .with_identifier_field_ids(vec![2])
4095 .build()
4096 .unwrap(),
4097 )
4098 .properties(HashMap::from([("owner".to_string(), "testx".to_string())]))
4099 .partition_spec(
4100 UnboundPartitionSpec::builder()
4101 .add_partition_field(
4102 UnboundPartitionField::builder()
4103 .source_ids(vec![1])
4104 .name("id")
4105 .transform(Transform::Truncate(3))
4106 .build()
4107 .unwrap(),
4108 )
4109 .unwrap()
4110 .build(),
4111 )
4112 .sort_order(
4113 SortOrder::builder()
4114 .with_sort_field(
4115 SortField::builder()
4116 .source_id(2)
4117 .transform(Transform::Identity)
4118 .direction(SortDirection::Ascending)
4119 .null_order(NullOrder::First)
4120 .build(),
4121 )
4122 .build_unbound()
4123 .unwrap(),
4124 )
4125 .build();
4126
4127 let table = catalog
4128 .create_table(
4129 &SessionContext::empty(),
4130 &NamespaceIdent::from_strs(["ns1"]).unwrap(),
4131 table_creation,
4132 )
4133 .await
4134 .unwrap();
4135
4136 assert_eq!(
4137 &TableIdent::from_strs(vec!["ns1", "test1"]).unwrap(),
4138 table.identifier()
4139 );
4140 assert_eq!(
4141 "s3://warehouse/database/table/metadata.json",
4142 table.metadata_location().unwrap()
4143 );
4144 assert_eq!(FormatVersion::V1, table.metadata().format_version());
4145 assert_eq!("s3://warehouse/database/table", table.metadata().location());
4146 assert_eq!(
4147 uuid!("bf289591-dcc0-4234-ad4f-5c3eed811a29"),
4148 table.metadata().uuid()
4149 );
4150 assert_eq!(
4151 1657810967051,
4152 table
4153 .metadata()
4154 .last_updated_timestamp()
4155 .unwrap()
4156 .timestamp_millis()
4157 );
4158 assert_eq!(
4159 vec![&Arc::new(
4160 Schema::builder()
4161 .with_fields(vec![
4162 NestedField::optional(1, "foo", Type::Primitive(PrimitiveType::String))
4163 .into(),
4164 NestedField::required(2, "bar", Type::Primitive(PrimitiveType::Int)).into(),
4165 NestedField::optional(3, "baz", Type::Primitive(PrimitiveType::Boolean))
4166 .into(),
4167 ])
4168 .with_schema_id(0)
4169 .with_identifier_field_ids(vec![2])
4170 .build()
4171 .unwrap()
4172 )],
4173 table.metadata().schemas_iter().collect::<Vec<_>>()
4174 );
4175 assert_eq!(
4176 &HashMap::from([
4177 (
4178 "write.delete.parquet.compression-codec".to_string(),
4179 "zstd".to_string()
4180 ),
4181 (
4182 "write.metadata.compression-codec".to_string(),
4183 "gzip".to_string()
4184 ),
4185 (
4186 "write.summary.partition-limit".to_string(),
4187 "100".to_string()
4188 ),
4189 (
4190 "write.parquet.compression-codec".to_string(),
4191 "zstd".to_string()
4192 ),
4193 ]),
4194 table.metadata().properties()
4195 );
4196 assert!(table.metadata().current_snapshot().is_none());
4197 assert!(table.metadata().history().is_empty());
4198 assert_eq!(
4199 vec![&Arc::new(SortOrder {
4200 order_id: 0,
4201 fields: vec![],
4202 })],
4203 table.metadata().sort_orders_iter().collect::<Vec<_>>()
4204 );
4205
4206 config_mock.assert_async().await;
4207 create_table_mock.assert_async().await;
4208 }
4209
4210 #[tokio::test]
4211 async fn test_create_table_409() {
4212 let mut server = Server::new_async().await;
4213
4214 let config_mock = create_config_mock(&mut server).await;
4215
4216 let create_table_mock = server
4217 .mock("POST", "/v1/namespaces/ns1/tables")
4218 .with_status(409)
4219 .with_body(r#"
4220{
4221 "error": {
4222 "message": "Table already exists: ns1.test1 in warehouse 8bcb0838-50fc-472d-9ddb-8feb89ef5f1e",
4223 "type": "AlreadyExistsException",
4224 "code": 409
4225 }
4226}
4227 "#)
4228 .create_async()
4229 .await;
4230
4231 let catalog = session_catalog(RestCatalogConfig::builder().uri(server.url()).build());
4232
4233 let table_creation = TableCreation::builder()
4234 .name("test1".to_string())
4235 .schema(
4236 Schema::builder()
4237 .with_fields(vec![
4238 NestedField::optional(1, "foo", Type::Primitive(PrimitiveType::String))
4239 .into(),
4240 NestedField::required(2, "bar", Type::Primitive(PrimitiveType::Int)).into(),
4241 NestedField::optional(3, "baz", Type::Primitive(PrimitiveType::Boolean))
4242 .into(),
4243 ])
4244 .with_schema_id(1)
4245 .with_identifier_field_ids(vec![2])
4246 .build()
4247 .unwrap(),
4248 )
4249 .properties(HashMap::from([("owner".to_string(), "testx".to_string())]))
4250 .build();
4251
4252 let table_result = catalog
4253 .create_table(
4254 &SessionContext::empty(),
4255 &NamespaceIdent::from_strs(["ns1"]).unwrap(),
4256 table_creation,
4257 )
4258 .await;
4259
4260 assert!(table_result.is_err());
4261 assert!(
4262 table_result
4263 .err()
4264 .unwrap()
4265 .message()
4266 .contains("already exists")
4267 );
4268
4269 config_mock.assert_async().await;
4270 create_table_mock.assert_async().await;
4271 }
4272
4273 #[tokio::test]
4274 async fn test_update_table() {
4275 let mut server = Server::new_async().await;
4276
4277 let config_mock = create_config_mock(&mut server).await;
4278
4279 let load_table_mock = server
4280 .mock("GET", "/v1/namespaces/ns1/tables/test1")
4281 .with_status(200)
4282 .with_body_from_file(format!(
4283 "{}/testdata/{}",
4284 env!("CARGO_MANIFEST_DIR"),
4285 "load_table_response.json"
4286 ))
4287 .create_async()
4288 .await;
4289
4290 let update_table_mock = server
4291 .mock("POST", "/v1/namespaces/ns1/tables/test1")
4292 .with_status(200)
4293 .with_body_from_file(format!(
4294 "{}/testdata/{}",
4295 env!("CARGO_MANIFEST_DIR"),
4296 "update_table_response.json"
4297 ))
4298 .create_async()
4299 .await;
4300
4301 let catalog = RestCatalog::new(
4302 SessionContext::empty(),
4303 RestCatalogConfig::builder().uri(server.url()).build(),
4304 None,
4305 Some(Arc::new(LocalFsStorageFactory)),
4306 Runtime::current(),
4307 None,
4308 );
4309
4310 let table1 = {
4311 let file = File::open(format!(
4312 "{}/testdata/{}",
4313 env!("CARGO_MANIFEST_DIR"),
4314 "create_table_response.json"
4315 ))
4316 .unwrap();
4317 let reader = BufReader::new(file);
4318 let resp = serde_json::from_reader::<_, LoadTableResult>(reader).unwrap();
4319
4320 Table::builder()
4321 .metadata(resp.metadata)
4322 .metadata_location(resp.metadata_location.unwrap())
4323 .identifier(TableIdent::from_strs(["ns1", "test1"]).unwrap())
4324 .file_io(FileIO::new_with_fs())
4325 .runtime(test_runtime())
4326 .build()
4327 .unwrap()
4328 };
4329
4330 let tx = Transaction::new(&table1);
4331 let table = tx
4332 .upgrade_table_version()
4333 .set_format_version(FormatVersion::V2)
4334 .apply(tx)
4335 .unwrap()
4336 .commit(&catalog)
4337 .await
4338 .unwrap();
4339
4340 assert_eq!(
4341 &TableIdent::from_strs(vec!["ns1", "test1"]).unwrap(),
4342 table.identifier()
4343 );
4344 assert_eq!(
4345 "s3://warehouse/database/table/metadata.json",
4346 table.metadata_location().unwrap()
4347 );
4348 assert_eq!(FormatVersion::V2, table.metadata().format_version());
4349 assert_eq!("s3://warehouse/database/table", table.metadata().location());
4350 assert_eq!(
4351 uuid!("bf289591-dcc0-4234-ad4f-5c3eed811a29"),
4352 table.metadata().uuid()
4353 );
4354 assert_eq!(
4355 1657810967051,
4356 table
4357 .metadata()
4358 .last_updated_timestamp()
4359 .unwrap()
4360 .timestamp_millis()
4361 );
4362 assert_eq!(
4363 vec![&Arc::new(
4364 Schema::builder()
4365 .with_fields(vec![
4366 NestedField::optional(1, "foo", Type::Primitive(PrimitiveType::String))
4367 .into(),
4368 NestedField::required(2, "bar", Type::Primitive(PrimitiveType::Int)).into(),
4369 NestedField::optional(3, "baz", Type::Primitive(PrimitiveType::Boolean))
4370 .into(),
4371 ])
4372 .with_schema_id(0)
4373 .with_identifier_field_ids(vec![2])
4374 .build()
4375 .unwrap()
4376 )],
4377 table.metadata().schemas_iter().collect::<Vec<_>>()
4378 );
4379 assert_eq!(
4380 &HashMap::from([
4381 (
4382 "write.delete.parquet.compression-codec".to_string(),
4383 "zstd".to_string()
4384 ),
4385 (
4386 "write.metadata.compression-codec".to_string(),
4387 "gzip".to_string()
4388 ),
4389 (
4390 "write.summary.partition-limit".to_string(),
4391 "100".to_string()
4392 ),
4393 (
4394 "write.parquet.compression-codec".to_string(),
4395 "zstd".to_string()
4396 ),
4397 ]),
4398 table.metadata().properties()
4399 );
4400 assert!(table.metadata().current_snapshot().is_none());
4401 assert!(table.metadata().history().is_empty());
4402 assert_eq!(
4403 vec![&Arc::new(SortOrder {
4404 order_id: 0,
4405 fields: vec![],
4406 })],
4407 table.metadata().sort_orders_iter().collect::<Vec<_>>()
4408 );
4409
4410 config_mock.assert_async().await;
4411 update_table_mock.assert_async().await;
4412 load_table_mock.assert_async().await
4413 }
4414
4415 #[tokio::test]
4416 async fn test_update_table_404() {
4417 let mut server = Server::new_async().await;
4418
4419 let config_mock = create_config_mock(&mut server).await;
4420
4421 let load_table_mock = server
4422 .mock("GET", "/v1/namespaces/ns1/tables/test1")
4423 .with_status(200)
4424 .with_body_from_file(format!(
4425 "{}/testdata/{}",
4426 env!("CARGO_MANIFEST_DIR"),
4427 "load_table_response.json"
4428 ))
4429 .create_async()
4430 .await;
4431
4432 let update_table_mock = server
4433 .mock("POST", "/v1/namespaces/ns1/tables/test1")
4434 .with_status(404)
4435 .with_body(
4436 r#"
4437{
4438 "error": {
4439 "message": "The given table does not exist",
4440 "type": "NoSuchTableException",
4441 "code": 404
4442 }
4443}
4444 "#,
4445 )
4446 .create_async()
4447 .await;
4448
4449 let catalog = RestCatalog::new(
4450 SessionContext::empty(),
4451 RestCatalogConfig::builder().uri(server.url()).build(),
4452 None,
4453 Some(Arc::new(LocalFsStorageFactory)),
4454 Runtime::current(),
4455 None,
4456 );
4457
4458 let table1 = {
4459 let file = File::open(format!(
4460 "{}/testdata/{}",
4461 env!("CARGO_MANIFEST_DIR"),
4462 "create_table_response.json"
4463 ))
4464 .unwrap();
4465 let reader = BufReader::new(file);
4466 let resp = serde_json::from_reader::<_, LoadTableResult>(reader).unwrap();
4467
4468 Table::builder()
4469 .metadata(resp.metadata)
4470 .metadata_location(resp.metadata_location.unwrap())
4471 .identifier(TableIdent::from_strs(["ns1", "test1"]).unwrap())
4472 .file_io(FileIO::new_with_fs())
4473 .runtime(test_runtime())
4474 .build()
4475 .unwrap()
4476 };
4477
4478 let tx = Transaction::new(&table1);
4479 let table_result = tx
4480 .upgrade_table_version()
4481 .set_format_version(FormatVersion::V2)
4482 .apply(tx)
4483 .unwrap()
4484 .commit(&catalog)
4485 .await;
4486
4487 assert!(table_result.is_err());
4488 assert!(
4489 table_result
4490 .err()
4491 .unwrap()
4492 .message()
4493 .contains("does not exist")
4494 );
4495
4496 config_mock.assert_async().await;
4497 update_table_mock.assert_async().await;
4498 load_table_mock.assert_async().await;
4499 }
4500
4501 #[tokio::test]
4502 async fn test_register_table() {
4503 let mut server = Server::new_async().await;
4504
4505 let config_mock = create_config_mock(&mut server).await;
4506
4507 let register_table_mock = server
4508 .mock("POST", "/v1/namespaces/ns1/register")
4509 .with_status(200)
4510 .with_body_from_file(format!(
4511 "{}/testdata/{}",
4512 env!("CARGO_MANIFEST_DIR"),
4513 "load_table_response.json"
4514 ))
4515 .create_async()
4516 .await;
4517
4518 let catalog = session_catalog(RestCatalogConfig::builder().uri(server.url()).build());
4519 let table_ident =
4520 TableIdent::new(NamespaceIdent::new("ns1".to_string()), "test1".to_string());
4521 let metadata_location = String::from(
4522 "s3://warehouse/database/table/metadata/00001-5f2f8166-244c-4eae-ac36-384ecdec81fc.gz.metadata.json",
4523 );
4524
4525 let table = catalog
4526 .register_table(&SessionContext::empty(), &table_ident, metadata_location)
4527 .await
4528 .unwrap();
4529
4530 assert_eq!(
4531 &TableIdent::from_strs(vec!["ns1", "test1"]).unwrap(),
4532 table.identifier()
4533 );
4534 assert_eq!(
4535 "s3://warehouse/database/table/metadata/00001-5f2f8166-244c-4eae-ac36-384ecdec81fc.gz.metadata.json",
4536 table.metadata_location().unwrap()
4537 );
4538
4539 config_mock.assert_async().await;
4540 register_table_mock.assert_async().await;
4541 }
4542
4543 #[tokio::test]
4544 async fn test_register_table_404() {
4545 let mut server = Server::new_async().await;
4546
4547 let config_mock = create_config_mock(&mut server).await;
4548
4549 let register_table_mock = server
4550 .mock("POST", "/v1/namespaces/ns1/register")
4551 .with_status(404)
4552 .with_body(
4553 r#"
4554{
4555 "error": {
4556 "message": "The namespace specified does not exist",
4557 "type": "NoSuchNamespaceErrorException",
4558 "code": 404
4559 }
4560}
4561 "#,
4562 )
4563 .create_async()
4564 .await;
4565
4566 let catalog = session_catalog(RestCatalogConfig::builder().uri(server.url()).build());
4567
4568 let table_ident =
4569 TableIdent::new(NamespaceIdent::new("ns1".to_string()), "test1".to_string());
4570 let metadata_location = String::from(
4571 "s3://warehouse/database/table/metadata/00001-5f2f8166-244c-4eae-ac36-384ecdec81fc.gz.metadata.json",
4572 );
4573 let table = catalog
4574 .register_table(&SessionContext::empty(), &table_ident, metadata_location)
4575 .await;
4576
4577 assert!(table.is_err());
4578 assert!(table.err().unwrap().message().contains("does not exist"));
4579
4580 config_mock.assert_async().await;
4581 register_table_mock.assert_async().await;
4582 }
4583
4584 #[tokio::test]
4585 async fn test_create_rest_catalog() {
4586 let builder = RestCatalogBuilder::default().with_client(Client::new());
4587
4588 let catalog = builder
4589 .load(
4590 "test",
4591 HashMap::from([
4592 (
4593 REST_CATALOG_PROP_URI.to_string(),
4594 "http://localhost:8080".to_string(),
4595 ),
4596 ("a".to_string(), "b".to_string()),
4597 ]),
4598 )
4599 .await;
4600
4601 assert!(catalog.is_ok());
4602
4603 let catalog = catalog.unwrap();
4604 let catalog_config = &catalog.inner.user_config;
4605 assert_eq!(catalog_config.name.as_deref(), Some("test"));
4606 assert_eq!(catalog_config.uri, "http://localhost:8080");
4607 assert_eq!(catalog_config.warehouse, None);
4608 assert!(catalog_config.client.is_some());
4609
4610 assert_eq!(catalog_config.props.get("a"), Some(&"b".to_string()));
4611 assert!(!catalog_config.props.contains_key(REST_CATALOG_PROP_URI));
4612 }
4613
4614 #[tokio::test]
4615 async fn test_create_rest_catalog_no_uri() {
4616 let builder = RestCatalogBuilder::default();
4617
4618 let catalog = builder
4619 .load(
4620 "test",
4621 HashMap::from([(
4622 REST_CATALOG_PROP_WAREHOUSE.to_string(),
4623 "s3://warehouse".to_string(),
4624 )]),
4625 )
4626 .await;
4627
4628 assert!(catalog.is_err());
4629 if let Err(err) = catalog {
4630 assert_eq!(err.kind(), ErrorKind::DataInvalid);
4631 assert_eq!(err.message(), "Catalog uri is required");
4632 }
4633 }
4634
4635 #[tokio::test]
4636 async fn test_create_session_catalog() {
4637 let builder = RestSessionCatalogBuilder::default();
4638
4639 let result = builder
4640 .load(
4641 "test",
4642 HashMap::from([
4643 (
4644 REST_CATALOG_PROP_URI.to_string(),
4645 "http://localhost:8080".to_string(),
4646 ),
4647 ("a".to_string(), "b".to_string()),
4648 ]),
4649 )
4650 .await;
4651
4652 assert!(result.is_ok());
4653
4654 let catalog = result.unwrap();
4655
4656 let catalog_config = catalog.user_config;
4657 assert_eq!(catalog_config.name.as_deref(), Some("test"));
4658 assert_eq!(catalog_config.uri, "http://localhost:8080");
4659 assert_eq!(catalog_config.warehouse, None);
4660 assert!(catalog_config.client.is_none());
4662
4663 assert_eq!(catalog_config.props.get("a"), Some(&"b".to_string()));
4665 assert!(!catalog_config.props.contains_key(REST_CATALOG_PROP_URI));
4666 }
4667
4668 #[tokio::test]
4669 async fn test_create_rest_catalog_with_session() {
4670 let context = SessionContext::builder()
4671 .session_id("test-id".to_string())
4672 .build();
4673
4674 let result = RestCatalogBuilder::default()
4675 .with_session_context(context)
4676 .load(
4677 "test",
4678 HashMap::from([(
4679 REST_CATALOG_PROP_URI.to_string(),
4680 "http://localhost:8080".to_string(),
4681 )]),
4682 )
4683 .await;
4684
4685 assert!(result.is_ok());
4686
4687 let catalog = result.unwrap();
4689 assert_eq!(catalog.session_context.session_id(), "test-id");
4690 }
4691
4692 #[tokio::test]
4693 async fn test_create_rest_catalog_default_session() {
4694 let result = RestCatalogBuilder::default()
4695 .load(
4696 "test",
4697 HashMap::from([(
4698 REST_CATALOG_PROP_URI.to_string(),
4699 "http://localhost:8080".to_string(),
4700 )]),
4701 )
4702 .await;
4703
4704 assert!(result.is_ok());
4705
4706 let catalog = result.unwrap();
4709 assert!(uuid::Uuid::parse_str(catalog.session_context.session_id()).is_ok());
4710 }
4711
4712 #[tokio::test]
4714 async fn test_rest_catalog_delegates() {
4715 let mut server = Server::new_async().await;
4716 let config_mock = create_config_mock(&mut server).await;
4717 let list_ns_mock = server
4718 .mock("GET", "/v1/namespaces")
4719 .with_body(
4720 r#"{
4721 "namespaces": [["ns1"]]
4722 }"#,
4723 )
4724 .create_async()
4725 .await;
4726
4727 let catalog = RestCatalog::new(
4728 SessionContext::empty(),
4729 RestCatalogConfig::builder().uri(server.url()).build(),
4730 None,
4731 Some(Arc::new(LocalFsStorageFactory)),
4732 Runtime::current(),
4733 None,
4734 );
4735
4736 let namespaces = catalog.list_namespaces(None).await.unwrap();
4737
4738 assert_eq!(namespaces, vec![
4739 NamespaceIdent::from_vec(vec!["ns1".to_string()]).unwrap()
4740 ]);
4741
4742 config_mock.assert_async().await;
4743 list_ns_mock.assert_async().await;
4744 }
4745}