use crate::datasource::TableProvider;
use crate::error::{DataFusionError, Result};
use std::any::Any;
use std::collections::HashMap;
use std::sync::{Arc, RwLock};
pub trait SchemaProvider: Sync + Send {
fn as_any(&self) -> &dyn Any;
fn table_names(&self) -> Vec<String>;
fn table(&self, name: &str) -> Option<Arc<dyn TableProvider>>;
#[allow(unused_variables)]
fn register_table(
&self,
name: String,
table: Arc<dyn TableProvider>,
) -> Result<Option<Arc<dyn TableProvider>>> {
Err(DataFusionError::Execution(
"schema provider does not support registering tables".to_owned(),
))
}
#[allow(unused_variables)]
fn deregister_table(&self, name: &str) -> Result<Option<Arc<dyn TableProvider>>> {
Err(DataFusionError::Execution(
"schema provider does not support deregistering tables".to_owned(),
))
}
}
pub struct MemorySchemaProvider {
tables: RwLock<HashMap<String, Arc<dyn TableProvider>>>,
}
impl MemorySchemaProvider {
pub fn new() -> Self {
Self {
tables: RwLock::new(HashMap::new()),
}
}
}
impl SchemaProvider for MemorySchemaProvider {
fn as_any(&self) -> &dyn Any {
self
}
fn table_names(&self) -> Vec<String> {
let tables = self.tables.read().unwrap();
tables.keys().cloned().collect()
}
fn table(&self, name: &str) -> Option<Arc<dyn TableProvider>> {
let tables = self.tables.read().unwrap();
tables.get(name).cloned()
}
fn register_table(
&self,
name: String,
table: Arc<dyn TableProvider>,
) -> Result<Option<Arc<dyn TableProvider>>> {
let mut tables = self.tables.write().unwrap();
Ok(tables.insert(name, table))
}
fn deregister_table(&self, name: &str) -> Result<Option<Arc<dyn TableProvider>>> {
let mut tables = self.tables.write().unwrap();
Ok(tables.remove(name))
}
}