use std::collections::BTreeMap;
use std::sync::{Arc, PoisonError, RwLock};
use std::time::Duration;
use async_trait::async_trait;
use datafusion::catalog::{CatalogProvider, SchemaProvider, TableProvider};
use datafusion::error::Result as DfResult;
use super::PLANNING_REUSE;
use super::provider::TimelakeProvider;
use crate::error::{Error, Result};
use crate::lake::Lake;
use crate::name::TableName;
pub(super) const INFORMATION_SCHEMA: &str = "information_schema";
type Listing = RwLock<BTreeMap<String, Vec<String>>>;
type Interval = Arc<RwLock<Option<Duration>>>;
#[derive(Debug)]
pub(super) struct LakeCatalog {
lake: Lake,
listing: Arc<Listing>,
interval: Interval,
}
impl LakeCatalog {
pub(super) fn new(lake: Lake) -> Self {
Self { lake, listing: Arc::default(), interval: Interval::default() }
}
pub(super) fn set_interval(&self, interval: Option<Duration>) {
*self.interval.write().unwrap_or_else(PoisonError::into_inner) = interval;
}
pub(super) async fn refresh_listing(&self) -> Result<()> {
let namespaces = self.lake.namespaces().await?;
let tables = futures::future::try_join_all(namespaces.iter().map(|ns| self.lake.tables(ns))).await?;
let listing = namespaces
.into_iter()
.zip(tables)
.map(|(ns, names)| (ns, names.iter().map(|n| n.name().to_string()).collect()))
.collect();
*self.listing.write().unwrap_or_else(PoisonError::into_inner) = listing;
Ok(())
}
}
impl CatalogProvider for LakeCatalog {
fn schema_names(&self) -> Vec<String> {
self.listing.read().unwrap_or_else(PoisonError::into_inner).keys().cloned().collect()
}
fn schema(&self, name: &str) -> Option<Arc<dyn SchemaProvider>> {
let (listing, interval) = (self.listing.clone(), self.interval.clone());
Some(Arc::new(NamespaceSchema { lake: self.lake.clone(), namespace: name.to_string(), listing, interval }))
}
}
#[derive(Debug)]
struct NamespaceSchema {
lake: Lake,
namespace: String,
listing: Arc<Listing>,
interval: Interval,
}
#[async_trait]
impl SchemaProvider for NamespaceSchema {
fn table_names(&self) -> Vec<String> {
let listing = self.listing.read().unwrap_or_else(PoisonError::into_inner);
listing.get(&self.namespace).cloned().unwrap_or_default()
}
async fn table(&self, name: &str) -> DfResult<Option<Arc<dyn TableProvider>>> {
let Ok(name) = TableName::new(&self.namespace, name) else { return Ok(None) };
let interval = *self.interval.read().unwrap_or_else(PoisonError::into_inner);
match self.lake.table_recent(&name, interval, PLANNING_REUSE).await {
Ok(t) => Ok(Some(Arc::new(TimelakeProvider::data(t)))),
Err(Error::NotFound(_)) => Ok(None),
Err(e) => Err(e.into()),
}
}
fn table_exist(&self, name: &str) -> bool {
self.table_names().iter().any(|t| t == name)
}
}