diff --git a/datafusion/core/src/execution/session_state.rs b/datafusion/core/src/execution/session_state.rs index f656fae4016f..a46e5c36156c 100644 --- a/datafusion/core/src/execution/session_state.rs +++ b/datafusion/core/src/execution/session_state.rs @@ -1193,6 +1193,18 @@ impl SessionStateBuilder { self } + /// Add a [`TableProviderFactory`] to the map of factories + pub fn with_table_factory( + mut self, + key: String, + table_factory: Arc, + ) -> Self { + let mut table_factories = self.table_factories.unwrap_or_default(); + table_factories.insert(key, table_factory); + self.table_factories = Some(table_factories); + self + } + /// Set the map of [`TableProviderFactory`]s pub fn with_table_factories( mut self, @@ -1929,4 +1941,22 @@ mod tests { Optimizer::default().rules.len() + 1 ); } + + #[test] + fn test_with_table_factories() -> Result<()> { + use crate::test_util::TestTableFactory; + + let state = SessionStateBuilder::new().build(); + let table_factories = state.table_factories(); + assert!(table_factories.is_empty()); + + let table_factory = Arc::new(TestTableFactory {}); + let state = SessionStateBuilder::new() + .with_table_factory("employee".to_string(), table_factory) + .build(); + let table_factories = state.table_factories(); + assert_eq!(table_factories.len(), 1); + assert!(table_factories.contains_key("employee")); + Ok(()) + } }