use std::sync::Arc; use api::query::{function::FuncMetaManagerRef, session::SessionCtx}; use async_trait::async_trait; use datafusion::arrow::datatypes::DataType; use datafusion::common::Result as DFResult; use datafusion::datasource::listing::ListingTable; use datafusion::logical_expr::var_provider::is_system_variables; use datafusion::logical_expr::{AggregateUDF, ScalarUDF, TableSource, WindowUDF}; use datafusion::variable::VarType; use datafusion::{ config::ConfigOptions, sql::{planner::ContextProvider, TableReference}, }; use crate::data_source::table_source::{TableHandle, TableSourceAdapter}; pub mod base_table; #[async_trait] pub trait ContextProviderExtension: ContextProvider { fn get_table_source_(&self, name: TableReference) -> datafusion::common::Result>; } pub type TableHandleProviderRef = Arc; pub trait TableHandleProvider { fn build_table_handle(&self, provider: Arc) -> DFResult; } pub struct MetadataProvider { provider: Arc, session: SessionCtx, config_options: ConfigOptions, func_manager: FuncMetaManagerRef, current_session_table_provider: TableHandleProviderRef, } impl MetadataProvider { #[allow(clippy::too_many_arguments)] pub fn new( provider: Arc, current_session_table_provider: TableHandleProviderRef, func_manager: FuncMetaManagerRef, session: SessionCtx, ) -> Self { Self { provider, current_session_table_provider, config_options: session.inner().config_options().clone(), session, func_manager, } } fn build_table_handle(&self) -> datafusion::common::Result { self.current_session_table_provider.build_table_handle(self.provider.clone()) } } impl ContextProviderExtension for MetadataProvider { fn get_table_source_(&self, table_ref: TableReference) -> datafusion::common::Result> { let name = table_ref.clone().resolve("", ""); let table_name = &*name.table; let table_handle = self.build_table_handle()?; Ok(Arc::new(TableSourceAdapter::try_new(table_ref.clone(), table_name, table_handle)?)) } } impl ContextProvider for MetadataProvider { fn get_function_meta(&self, name: &str) -> Option> { self.func_manager .udf(name) .ok() .or(self.session.inner().scalar_functions().get(name).cloned()) } fn get_aggregate_meta(&self, name: &str) -> Option> { self.func_manager.udaf(name).ok() } fn get_variable_type(&self, variable_names: &[String]) -> Option { if variable_names.is_empty() { return None; } let var_type = if is_system_variables(variable_names) { VarType::System } else { VarType::UserDefined }; self.session .inner() .execution_props() .get_var_provider(var_type) .and_then(|p| p.get_type(variable_names)) } fn options(&self) -> &ConfigOptions { // TODO refactor &self.config_options } fn get_window_meta(&self, name: &str) -> Option> { self.func_manager.udwf(name).ok() } fn get_table_source(&self, name: TableReference) -> DFResult> { Ok(self.get_table_source_(name)?) } fn udf_names(&self) -> Vec { todo!() } fn udaf_names(&self) -> Vec { todo!() } fn udwf_names(&self) -> Vec { todo!() } }