Skip to main content

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}