Skip to main content

iceberg_catalog_loader/
lib.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18use std::collections::HashMap;
19use std::sync::Arc;
20
21use async_trait::async_trait;
22use iceberg::encryption::kms::KmsClientFactory;
23use iceberg::io::StorageFactory;
24use iceberg::{Catalog, CatalogBuilder, Error, ErrorKind, Result};
25use iceberg_catalog_glue::GlueCatalogBuilder;
26use iceberg_catalog_hms::HmsCatalogBuilder;
27use iceberg_catalog_rest::RestCatalogBuilder;
28use iceberg_catalog_s3tables::S3TablesCatalogBuilder;
29use iceberg_catalog_sql::SqlCatalogBuilder;
30
31/// A CatalogBuilderFactory creating a new catalog builder.
32type CatalogBuilderFactory = fn() -> Box<dyn BoxedCatalogBuilder>;
33
34/// A registry of catalog builders.
35static CATALOG_REGISTRY: &[(&str, CatalogBuilderFactory)] = &[
36    ("rest", || Box::new(RestCatalogBuilder::default())),
37    ("glue", || Box::new(GlueCatalogBuilder::default())),
38    ("s3tables", || Box::new(S3TablesCatalogBuilder::default())),
39    ("hms", || Box::new(HmsCatalogBuilder::default())),
40    ("sql", || Box::new(SqlCatalogBuilder::default())),
41];
42
43/// Return the list of supported catalog types.
44pub fn supported_types() -> Vec<&'static str> {
45    CATALOG_REGISTRY.iter().map(|(k, _)| *k).collect()
46}
47
48#[async_trait]
49pub trait BoxedCatalogBuilder: Send {
50    fn with_storage_factory(
51        self: Box<Self>,
52        storage_factory: Arc<dyn StorageFactory>,
53    ) -> Box<dyn BoxedCatalogBuilder>;
54
55    fn with_kms_client_factory(
56        self: Box<Self>,
57        kms_client_factory: Arc<dyn KmsClientFactory>,
58    ) -> Box<dyn BoxedCatalogBuilder>;
59
60    async fn load(
61        self: Box<Self>,
62        name: String,
63        props: HashMap<String, String>,
64    ) -> Result<Arc<dyn Catalog>>;
65}
66
67#[async_trait]
68impl<T: CatalogBuilder + 'static> BoxedCatalogBuilder for T {
69    fn with_storage_factory(
70        self: Box<Self>,
71        storage_factory: Arc<dyn StorageFactory>,
72    ) -> Box<dyn BoxedCatalogBuilder> {
73        Box::new(CatalogBuilder::with_storage_factory(*self, storage_factory))
74    }
75
76    fn with_kms_client_factory(
77        self: Box<Self>,
78        kms_client_factory: Arc<dyn KmsClientFactory>,
79    ) -> Box<dyn BoxedCatalogBuilder> {
80        Box::new(CatalogBuilder::with_kms_client_factory(
81            *self,
82            kms_client_factory,
83        ))
84    }
85
86    async fn load(
87        self: Box<Self>,
88        name: String,
89        props: HashMap<String, String>,
90    ) -> Result<Arc<dyn Catalog>> {
91        let builder = *self;
92        Ok(Arc::new(builder.load(name, props).await?) as Arc<dyn Catalog>)
93    }
94}
95
96/// Load a catalog from a string.
97pub fn load(r#type: &str) -> Result<Box<dyn BoxedCatalogBuilder>> {
98    let key = r#type.trim();
99    if let Some((_, factory)) = CATALOG_REGISTRY
100        .iter()
101        .find(|(k, _)| k.eq_ignore_ascii_case(key))
102    {
103        Ok(factory())
104    } else {
105        Err(Error::new(
106            ErrorKind::FeatureUnsupported,
107            format!(
108                "Unsupported catalog type: {}. Supported types: {}",
109                r#type,
110                supported_types().join(", ")
111            ),
112        ))
113    }
114}
115
116/// Ergonomic catalog loader builder pattern.
117pub struct CatalogLoader<'a> {
118    catalog_type: &'a str,
119}
120
121impl<'a> From<&'a str> for CatalogLoader<'a> {
122    fn from(s: &'a str) -> Self {
123        Self { catalog_type: s }
124    }
125}
126
127impl CatalogLoader<'_> {
128    pub async fn load(
129        self,
130        name: String,
131        props: HashMap<String, String>,
132    ) -> Result<Arc<dyn Catalog>> {
133        let builder = load(self.catalog_type)?;
134        builder.load(name, props).await
135    }
136}
137
138#[cfg(test)]
139mod tests {
140    use std::collections::HashMap;
141    use std::sync::Arc;
142
143    use iceberg::io::LocalFsStorageFactory;
144    use sqlx::migrate::MigrateDatabase;
145    use tempfile::TempDir;
146
147    use crate::{CatalogLoader, load};
148
149    #[tokio::test]
150    async fn test_load_unsupported_catalog() {
151        let result = load("unsupported");
152        assert!(result.is_err());
153    }
154
155    #[tokio::test]
156    async fn test_catalog_loader_pattern() {
157        use iceberg_catalog_rest::REST_CATALOG_PROP_URI;
158
159        let catalog = CatalogLoader::from("rest")
160            .load(
161                "rest".to_string(),
162                HashMap::from([
163                    (
164                        REST_CATALOG_PROP_URI.to_string(),
165                        "http://localhost:8080".to_string(),
166                    ),
167                    ("key".to_string(), "value".to_string()),
168                ]),
169            )
170            .await;
171
172        assert!(catalog.is_ok());
173    }
174
175    #[tokio::test]
176    async fn test_catalog_loader_pattern_rest_catalog() {
177        use iceberg_catalog_rest::REST_CATALOG_PROP_URI;
178
179        let catalog_loader = load("rest").unwrap();
180        let catalog = catalog_loader
181            .load(
182                "rest".to_string(),
183                HashMap::from([
184                    (
185                        REST_CATALOG_PROP_URI.to_string(),
186                        "http://localhost:8080".to_string(),
187                    ),
188                    ("key".to_string(), "value".to_string()),
189                ]),
190            )
191            .await;
192
193        assert!(catalog.is_ok());
194    }
195
196    #[tokio::test]
197    async fn test_catalog_loader_pattern_glue_catalog() {
198        use iceberg_catalog_glue::GLUE_CATALOG_PROP_WAREHOUSE;
199
200        let catalog_loader = load("glue").unwrap();
201        let catalog = catalog_loader
202            .load(
203                "glue".to_string(),
204                HashMap::from([
205                    (
206                        GLUE_CATALOG_PROP_WAREHOUSE.to_string(),
207                        "s3://test".to_string(),
208                    ),
209                    ("key".to_string(), "value".to_string()),
210                ]),
211            )
212            .await;
213
214        assert!(catalog.is_ok());
215    }
216
217    #[tokio::test]
218    async fn test_catalog_loader_pattern_s3tables() {
219        use iceberg_catalog_s3tables::S3TABLES_CATALOG_PROP_TABLE_BUCKET_ARN;
220
221        let catalog = CatalogLoader::from("s3tables")
222            .load(
223                "s3tables".to_string(),
224                HashMap::from([
225                    (
226                        S3TABLES_CATALOG_PROP_TABLE_BUCKET_ARN.to_string(),
227                        "arn:aws:s3tables:us-east-1:123456789012:bucket/test".to_string(),
228                    ),
229                    ("key".to_string(), "value".to_string()),
230                ]),
231            )
232            .await;
233
234        assert!(catalog.is_ok());
235    }
236
237    #[tokio::test]
238    async fn test_catalog_loader_pattern_hms_catalog() {
239        use iceberg_catalog_hms::{HMS_CATALOG_PROP_URI, HMS_CATALOG_PROP_WAREHOUSE};
240
241        let catalog_loader = load("hms").unwrap();
242        let catalog = catalog_loader
243            .with_storage_factory(Arc::new(LocalFsStorageFactory))
244            .load(
245                "hms".to_string(),
246                HashMap::from([
247                    (HMS_CATALOG_PROP_URI.to_string(), "127.0.0.1:1".to_string()),
248                    (
249                        HMS_CATALOG_PROP_WAREHOUSE.to_string(),
250                        "s3://warehouse".to_string(),
251                    ),
252                    ("key".to_string(), "value".to_string()),
253                ]),
254            )
255            .await;
256
257        assert!(catalog.is_ok());
258    }
259
260    fn temp_path() -> String {
261        let temp_dir = TempDir::new().unwrap();
262        temp_dir.path().to_str().unwrap().to_string()
263    }
264
265    #[tokio::test]
266    async fn test_catalog_loader_pattern_sql_catalog() {
267        use iceberg_catalog_sql::{SQL_CATALOG_PROP_URI, SQL_CATALOG_PROP_WAREHOUSE};
268
269        let uri = format!("sqlite:{}", temp_path());
270        sqlx::Sqlite::create_database(&uri).await.unwrap();
271
272        let catalog_loader = load("sql").unwrap();
273        let catalog = catalog_loader
274            .with_storage_factory(Arc::new(LocalFsStorageFactory))
275            .load(
276                "sql".to_string(),
277                HashMap::from([
278                    (SQL_CATALOG_PROP_URI.to_string(), uri),
279                    (
280                        SQL_CATALOG_PROP_WAREHOUSE.to_string(),
281                        "s3://warehouse".to_string(),
282                    ),
283                ]),
284            )
285            .await;
286
287        assert!(catalog.is_ok());
288    }
289
290    #[tokio::test]
291    async fn test_error_message_includes_supported_types() {
292        let err = match load("does-not-exist") {
293            Ok(_) => panic!("expected error for unsupported type"),
294            Err(e) => e,
295        };
296        let msg = err.message().to_string();
297        assert!(msg.contains("Supported types:"));
298        // Should include at least the built-in type
299        assert!(msg.contains("rest"));
300    }
301}