use std::sync::Arc;
use datafusion::prelude::SessionContext;
pub use datafusion_federation;
pub trait FederatedContext {
fn federate(self) -> SessionContext;
fn is_federated(&self) -> bool;
}
impl FederatedContext for SessionContext {
fn federate(self) -> SessionContext {
use datafusion_federation::{FederatedQueryPlanner, default_optimizer_rules};
let state = self.state();
if state.optimizer().rules.iter().any(|rule| rule.name() == "federation_optimizer_rule") {
self
} else {
SessionContext::new_with_state(
self.into_state_builder()
.with_optimizer_rules(default_optimizer_rules())
.with_query_planner(Arc::new(FederatedQueryPlanner::new()))
.build(),
)
}
}
fn is_federated(&self) -> bool {
self.state()
.optimizer()
.rules
.iter()
.any(|rule| rule.name() == "federation_optimizer_rule")
}
}