Support partitioning of information schema content per catalog

datafusion/catalog/src/information_schema.rs

@@ -44,6 +44,7 @@ use datafusion_expr::{TableType, Volatility};
 use datafusion_physical_plan::SendableRecordBatchStream;
 use datafusion_physical_plan::stream::RecordBatchStreamAdapter;
 use datafusion_physical_plan::streaming::PartitionStream;
+use datafusion_session::CatalogProvider;
 use itertools::Itertools;
 use std::collections::{BTreeSet, HashMap, HashSet};
 use std::fmt::Debug;
@@ -81,10 +82,55 @@ pub struct InformationSchemaProvider {
 }
 
 impl InformationSchemaProvider {
+    /// Creates a new [`InformationSchemaProvider`] for the provided `catalog_list`
+    pub fn for_catalog(catalog_name: String, catalog: Arc<dyn CatalogProvider>) -> Self {
+        #[derive(Debug)]
+        struct SingleCatalogList {
+            catalog_name: String,
+            catalog: Arc<dyn CatalogProvider>,
+        }
+
+        impl CatalogProviderList for SingleCatalogList {
+            fn register_catalog(
+                &self,
+                _: String,
+                _: Arc<dyn CatalogProvider>,
+            ) -> Option<Arc<dyn CatalogProvider>> {
+                None
+            }
+
+            fn catalog_names(&self) -> Vec<String> {
+                vec![self.catalog_name.clone()]
+            }
+
+            fn catalog(&self, name: &str) -> Option<Arc<dyn CatalogProvider>> {
+                if name == self.catalog_name {
+                    Some(Arc::clone(&self.catalog))
+                } else {
+                    None
+                }
+            }
+        }
+
+        let catalog_list = Arc::new(SingleCatalogList {
+            catalog_name,
+            catalog,
+        });
+
+        Self {
+            config: InformationSchemaConfig {
+                system_catalog_name: None,
+                catalog_list,
+                table_functions: HashMap::new(),
+            },
+        }
+    }
+
     /// Creates a new [`InformationSchemaProvider`] for the provided `catalog_list`
     pub fn new(catalog_list: Arc<dyn CatalogProviderList>) -> Self {
         Self {
             config: InformationSchemaConfig {
+                system_catalog_name: None,
                 catalog_list,
                 table_functions: HashMap::new(),
             },
@@ -100,10 +146,16 @@ impl InformationSchemaProvider {
         self.config.table_functions = table_functions;
         self
     }
+
+    pub fn with_system_catalog(mut self, catalog_name: String) -> Self {
+        self.config.system_catalog_name = Some(catalog_name);
+        self
+    }
 }
 
 #[derive(Clone, Debug)]
 struct InformationSchemaConfig {
+    system_catalog_name: Option<String>,
     catalog_list: Arc<dyn CatalogProviderList>,
     table_functions: HashMap<String, Arc<TableFunction>>,
 }
@@ -150,6 +202,17 @@ impl InformationSchemaConfig {
             }
         }
 
+        if let Some(system_catalog_name) = &self.system_catalog_name {
+            for table_name in INFORMATION_SCHEMA_TABLES {
+                builder.add_table(
+                    system_catalog_name,
+                    INFORMATION_SCHEMA,
+                    table_name,
+                    TableType::View,
+                );
+            }
+        }
+
         Ok(())
     }
 
@@ -179,16 +242,13 @@ impl InformationSchemaConfig {
                 if schema_name != INFORMATION_SCHEMA {
                     // schema name may not exist in the catalog, so we need to check
                     if let Some(schema) = catalog.schema(&schema_name) {
-                        for table_name in schema.table_names() {
-                            if let Some(table) = schema.table(&table_name).await? {
-                                builder.add_view(
-                                    &catalog_name,
-                                    &schema_name,
-                                    &table_name,
-                                    table.get_table_definition(),
-                                )
-                            }
-                        }
+                        Self::add_views(
+                            builder,
+                            &catalog_name,
+                            &schema_name,
+                            schema.as_ref(),
+                        )
+                        .await?;
                     }
                 }
             }
@@ -197,6 +257,25 @@ impl InformationSchemaConfig {
         Ok(())
     }
 
+    async fn add_views(
+        builder: &mut InformationSchemaViewBuilder,
+        catalog_name: &str,
+        schema_name: &str,
+        schema: &dyn SchemaProvider,
+    ) -> Result<(), DataFusionError> {
+        for table_name in schema.table_names() {
+            if let Some(table) = schema.table(&table_name).await? {
+                builder.add_view(
+                    catalog_name,
+                    schema_name,
+                    &table_name,
+                    table.get_table_definition(),
+                )
+            }
+        }
+        Ok(())
+    }
+
     /// Construct the `information_schema.columns` virtual table
     async fn make_columns(
         &self,
@@ -209,21 +288,13 @@ impl InformationSchemaConfig {
                 if schema_name != INFORMATION_SCHEMA {
                     // schema name may not exist in the catalog, so we need to check
                     if let Some(schema) = catalog.schema(&schema_name) {
-                        for table_name in schema.table_names() {
-                            if let Some(table) = schema.table(&table_name).await? {
-                                for (field_position, field) in
-                                    table.schema().fields().iter().enumerate()
-                                {
-                                    builder.add_column(
-                                        &catalog_name,
-                                        &schema_name,
-                                        &table_name,
-                                        field_position,
-                                        field,
-                                    )
-                                }
-                            }
-                        }
+                        Self::add_columns(
+                            builder,
+                            &catalog_name,
+                            &schema_name,
+                            schema.as_ref(),
+                        )
+                        .await?;
                     }
                 }
             }
@@ -232,6 +303,29 @@ impl InformationSchemaConfig {
         Ok(())
     }
 
+    async fn add_columns(
+        builder: &mut InformationSchemaColumnsBuilder,
+        catalog_name: &str,
+        schema_name: &str,
+        schema: &dyn SchemaProvider,
+    ) -> Result<(), DataFusionError> {
+        for table_name in schema.table_names() {
+            if let Some(table) = schema.table(&table_name).await? {
+                for (field_position, field) in table.schema().fields().iter().enumerate()
+                {
+                    builder.add_column(
+                        catalog_name,
+                        schema_name,
+                        &table_name,
+                        field_position,
+                        field,
+                    )
+                }
+            }
+        }
+        Ok(())
+    }
+
     /// Construct the `information_schema.df_settings` virtual table
     fn make_df_settings(
         &self,
@@ -561,6 +655,24 @@ fn remove_native_type_prefix(native_type: &NativeType) -> String {
 
 #[async_trait]
 impl SchemaProvider for InformationSchemaProvider {
+    fn table_names(&self) -> Vec<String> {
+        self.config.table_names()
+    }
+
+    async fn table(
+        &self,
+        name: &str,
+    ) -> Result<Option<Arc<dyn TableProvider>>, DataFusionError> {
+        self.config.get_table(name)
+    }
+
+    fn table_exist(&self, name: &str) -> bool {
+        self.config.table_exist(name)
+    }
+}
+
+#[async_trait]
+impl SchemaProvider for InformationSchemaConfig {
     fn table_names(&self) -> Vec<String> {
         INFORMATION_SCHEMA_TABLES
             .iter()
@@ -572,7 +684,20 @@ impl SchemaProvider for InformationSchemaProvider {
         &self,
         name: &str,
     ) -> Result<Option<Arc<dyn TableProvider>>, DataFusionError> {
-        let config = self.config.clone();
+        self.get_table(name)
+    }
+
+    fn table_exist(&self, name: &str) -> bool {
+        INFORMATION_SCHEMA_TABLES.contains(&name.to_ascii_lowercase().as_str())
+    }
+}
+
+impl InformationSchemaConfig {
+    fn get_table(
+        &self,
+        name: &str,
+    ) -> Result<Option<Arc<dyn TableProvider>>, DataFusionError> {
+        let config = self.clone();
         let table: Arc<dyn PartitionStream> = match name.to_ascii_lowercase().as_str() {
             TABLES => Arc::new(InformationSchemaTables::new(config)),
             COLUMNS => Arc::new(InformationSchemaColumns::new(config)),
@@ -588,10 +713,6 @@ impl SchemaProvider for InformationSchemaProvider {
             StreamingTable::try_new(Arc::clone(table.schema()), vec![table]).unwrap(),
         )))
     }
-
-    fn table_exist(&self, name: &str) -> bool {
-        INFORMATION_SCHEMA_TABLES.contains(&name.to_ascii_lowercase().as_str())
-    }
 }
 
 #[derive(Debug)]
@@ -1553,6 +1674,7 @@ mod tests {
     #[tokio::test]
     async fn make_tables_uses_table_type() {
         let config = InformationSchemaConfig {
+            system_catalog_name: None,
             catalog_list: Arc::new(Fixture),
             table_functions: HashMap::new(),
         };

datafusion/common/src/config.rs

@@ -242,6 +242,12 @@ config_namespace! {
         /// virtual tables for displaying schema information
         pub information_schema: bool, default = false
 
+        /// The system catalog name - when set together with `information_schema`, the
+        /// `information_schema` for each catalog will only describe objects contained in the
+        /// catalog itself. The global, cross-catalog `information_schema` is available under
+        /// the virtual catalog named by this option.
+        pub system_catalog: Option<String>, default = None
+
         /// Location scanned to load tables for `default` schema
         pub location: Option<String>, default = None
 

datafusion/core/src/execution/context/mod.rs

@@ -2859,7 +2859,9 @@ mod tests {
     async fn sql_create_catalog() -> Result<()> {
         // the information schema used to introduce cyclic Arcs
         let ctx = SessionContext::new_with_config(
-            SessionConfig::new().with_information_schema(true),
+            SessionConfig::new()
+                .with_information_schema(true)
+                .with_system_catalog(Some(String::from("system"))),
         );
 
         // Create catalog
@@ -2875,7 +2877,7 @@ mod tests {
             .await?;
 
         // Check table exists in schema
-        let results = ctx.sql("SELECT * FROM information_schema.tables WHERE table_catalog='test' AND table_schema='abc' AND table_name = 'y'").await.unwrap().collect().await.unwrap();
+        let results = ctx.sql("SELECT * FROM system.information_schema.tables WHERE table_catalog='test' AND table_schema='abc' AND table_name = 'y'").await.unwrap().collect().await.unwrap();
 
         assert_eq!(results[0].num_rows(), 1);
         Ok(())