iceberg_datafusion/catalog.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 datafusion::catalog::{CatalogProvider, SchemaProvider};
22use futures::future::try_join_all;
23use iceberg::{Catalog, NamespaceIdent, Result};
24
25use crate::schema::IcebergSchemaProvider;
26
27/// Provides an interface to manage and access multiple schemas
28/// within an Iceberg [`Catalog`].
29///
30/// Acts as a centralized catalog provider that aggregates
31/// multiple [`SchemaProvider`], each associated with distinct namespaces.
32#[derive(Debug)]
33pub struct IcebergCatalogProvider {
34 /// A `HashMap` where keys are namespace names
35 /// and values are dynamic references to objects implementing the
36 /// [`SchemaProvider`] trait.
37 schemas: HashMap<String, Arc<dyn SchemaProvider>>,
38}
39
40impl IcebergCatalogProvider {
41 /// Asynchronously tries to construct a new [`IcebergCatalogProvider`]
42 /// using the given client to fetch and initialize schema providers for
43 /// each namespace in the Iceberg [`Catalog`].
44 ///
45 /// This method retrieves the list of namespace names
46 /// attempts to create a schema provider for each namespace, and
47 /// collects these providers into a `HashMap`.
48 pub async fn try_new(client: Arc<dyn Catalog>) -> Result<Self> {
49 // TODO:
50 // Schemas and providers should be cached and evicted based on time
51 // As of right now; schemas might become stale.
52 let schema_names: Vec<_> = client
53 .list_namespaces(None)
54 .await?
55 .iter()
56 .flat_map(|ns| ns.as_ref().clone())
57 .collect();
58
59 let providers = try_join_all(
60 schema_names
61 .iter()
62 .map(|name| {
63 IcebergSchemaProvider::try_new(
64 client.clone(),
65 NamespaceIdent::new(name.clone()),
66 )
67 })
68 .collect::<Vec<_>>(),
69 )
70 .await?;
71
72 let schemas: HashMap<String, Arc<dyn SchemaProvider>> = schema_names
73 .into_iter()
74 .zip(providers)
75 .map(|(name, provider)| {
76 let provider = Arc::new(provider) as Arc<dyn SchemaProvider>;
77 (name, provider)
78 })
79 .collect();
80
81 Ok(IcebergCatalogProvider { schemas })
82 }
83}
84
85impl CatalogProvider for IcebergCatalogProvider {
86 fn schema_names(&self) -> Vec<String> {
87 self.schemas.keys().cloned().collect()
88 }
89
90 fn schema(&self, name: &str) -> Option<Arc<dyn SchemaProvider>> {
91 self.schemas.get(name).cloned()
92 }
93}