feat: add pg_views/pg_matviews/pg_tables views

Author: sunng87

datafusion-postgres/src/pg_catalog.rs

@@ -26,6 +26,7 @@ mod pg_database;
 mod pg_get_expr_udf;
 mod pg_namespace;
 mod pg_settings;
+mod pg_tables;
 mod pg_views;
 
 const PG_CATALOG_TABLE_PG_AGGREGATE: &str = "pg_aggregate";
@@ -90,6 +91,9 @@ const PG_CATALOG_TABLE_PG_TABLESPACE: &str = "pg_tablespace";
 const PG_CATALOG_TABLE_PG_TRIGGER: &str = "pg_trigger";
 const PG_CATALOG_TABLE_PG_USER_MAPPING: &str = "pg_user_mapping";
 const PG_CATALOG_VIEW_PG_SETTINGS: &str = "pg_settings";
+const PG_CATALOG_VIEW_PG_VIEWS: &str = "pg_views";
+const PG_CATALOG_VIEW_PG_MATVIEWS: &str = "pg_matviews";
+const PG_CATALOG_VIEW_PG_TABLES: &str = "pg_tables";
 
 /// Determine PostgreSQL table type (relkind) from DataFusion TableProvider
 fn get_table_type(table: &Arc<dyn TableProvider>) -> &'static str {
@@ -185,6 +189,8 @@ pub const PG_CATALOG_TABLES: &[&str] = &[
     PG_CATALOG_TABLE_PG_TRIGGER,
     PG_CATALOG_TABLE_PG_USER_MAPPING,
     PG_CATALOG_VIEW_PG_SETTINGS,
+    PG_CATALOG_VIEW_PG_VIEWS,
+    PG_CATALOG_VIEW_PG_MATVIEWS,
 ];
 
 #[derive(Debug, Hash, Eq, PartialEq, PartialOrd, Ord)]
@@ -316,44 +322,63 @@ impl SchemaProvider for PgCatalogSchemaProvider {
                     self.oid_counter.clone(),
                     self.oid_cache.clone(),
                 ));
-                Ok(Some(Arc::new(
-                    StreamingTable::try_new(Arc::clone(table.schema()), vec![table]).unwrap(),
-                )))
+                Ok(Some(Arc::new(StreamingTable::try_new(
+                    Arc::clone(table.schema()),
+                    vec![table],
+                )?)))
             }
             PG_CATALOG_TABLE_PG_CLASS => {
                 let table = Arc::new(pg_class::PgClassTable::new(
                     self.catalog_list.clone(),
                     self.oid_counter.clone(),
                     self.oid_cache.clone(),
                 ));
-                Ok(Some(Arc::new(
-                    StreamingTable::try_new(Arc::clone(table.schema()), vec![table]).unwrap(),
-                )))
+                Ok(Some(Arc::new(StreamingTable::try_new(
+                    Arc::clone(table.schema()),
+                    vec![table],
+                )?)))
             }
             PG_CATALOG_TABLE_PG_DATABASE => {
                 let table = Arc::new(pg_database::PgDatabaseTable::new(
                     self.catalog_list.clone(),
                     self.oid_counter.clone(),
                     self.oid_cache.clone(),
                 ));
-                Ok(Some(Arc::new(
-                    StreamingTable::try_new(Arc::clone(table.schema()), vec![table]).unwrap(),
-                )))
+                Ok(Some(Arc::new(StreamingTable::try_new(
+                    Arc::clone(table.schema()),
+                    vec![table],
+                )?)))
             }
             PG_CATALOG_TABLE_PG_NAMESPACE => {
                 let table = Arc::new(pg_namespace::PgNamespaceTable::new(
                     self.catalog_list.clone(),
                     self.oid_counter.clone(),
                     self.oid_cache.clone(),
                 ));
-                Ok(Some(Arc::new(
-                    StreamingTable::try_new(Arc::clone(table.schema()), vec![table]).unwrap(),
-                )))
+                Ok(Some(Arc::new(StreamingTable::try_new(
+                    Arc::clone(table.schema()),
+                    vec![table],
+                )?)))
+            }
+            PG_CATALOG_VIEW_PG_TABLES => {
+                let table = Arc::new(pg_tables::PgTablesTable::new(self.catalog_list.clone()));
+                Ok(Some(Arc::new(StreamingTable::try_new(
+                    Arc::clone(table.schema()),
+                    vec![table],
+                )?)))
             }
             PG_CATALOG_VIEW_PG_SETTINGS => {
                 let table = pg_settings::PgSettingsView::try_new()?;
                 Ok(Some(Arc::new(table.try_into_memtable()?)))
             }
+            PG_CATALOG_VIEW_PG_VIEWS => {
+                let table = pg_views::PgViewsTable::new();
+                Ok(Some(Arc::new(table.try_into_memtable()?)))
+            }
+            PG_CATALOG_VIEW_PG_MATVIEWS => {
+                let table = pg_views::PgMatviewsTable::new();
+                Ok(Some(Arc::new(table.try_into_memtable()?)))
+            }
 
             _ => Ok(None),
         }

datafusion-postgres/src/pg_catalog/pg_tables.rs

@@ -0,0 +1,102 @@
+use std::sync::Arc;
+
+use datafusion::arrow::array::{ArrayRef, BooleanArray, RecordBatch, StringArray};
+use datafusion::arrow::datatypes::{DataType, Field, Schema, SchemaRef};
+use datafusion::catalog::CatalogProviderList;
+use datafusion::error::Result;
+use datafusion::execution::{SendableRecordBatchStream, TaskContext};
+use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
+use datafusion::physical_plan::streaming::PartitionStream;
+
+#[derive(Debug, Clone)]
+pub(crate) struct PgTablesTable {
+    schema: SchemaRef,
+    catalog_list: Arc<dyn CatalogProviderList>,
+}
+
+impl PgTablesTable {
+    pub(crate) fn new(catalog_list: Arc<dyn CatalogProviderList>) -> PgTablesTable {
+        // Define the schema for pg_class
+        // This matches key columns from PostgreSQL's pg_class
+        let schema = Arc::new(Schema::new(vec![
+            Field::new("schemaname", DataType::Utf8, false),
+            Field::new("tablename", DataType::Utf8, false),
+            Field::new("tableowner", DataType::Utf8, false),
+            Field::new("tablespace", DataType::Utf8, true),
+            Field::new("hasindex", DataType::Boolean, false),
+            Field::new("hasrules", DataType::Boolean, false),
+            Field::new("hastriggers", DataType::Boolean, false),
+            Field::new("rowsecurity", DataType::Boolean, false),
+        ]));
+
+        Self {
+            schema,
+            catalog_list,
+        }
+    }
+
+    /// Generate record batches based on the current state of the catalog
+    async fn get_data(this: PgTablesTable) -> Result<RecordBatch> {
+        // Vectors to store column data
+        let mut schema_names = Vec::new();
+        let mut table_names = Vec::new();
+        let mut table_owners = Vec::new();
+        let mut table_spaces: Vec<Option<String>> = Vec::new();
+        let mut has_index = Vec::new();
+        let mut has_rules = Vec::new();
+        let mut has_triggers = Vec::new();
+        let mut row_security = Vec::new();
+
+        // Iterate through all catalogs and schemas
+        for catalog_name in this.catalog_list.catalog_names() {
+            if let Some(catalog) = this.catalog_list.catalog(&catalog_name) {
+                for schema_name in catalog.schema_names() {
+                    if let Some(schema) = catalog.schema(&schema_name) {
+                        // Now process all tables in this schema
+                        for table_name in schema.table_names() {
+                            schema_names.push(schema_name.to_string());
+                            table_names.push(table_name.to_string());
+                            table_owners.push("postgres".to_string());
+                            table_spaces.push(None);
+                            has_index.push(false);
+                            has_rules.push(false);
+                            has_triggers.push(false);
+                            row_security.push(false);
+                        }
+                    }
+                }
+            }
+        }
+
+        // Create Arrow arrays from the collected data
+        let arrays: Vec<ArrayRef> = vec![
+            Arc::new(StringArray::from(schema_names)),
+            Arc::new(StringArray::from(table_names)),
+            Arc::new(StringArray::from(table_owners)),
+            Arc::new(StringArray::from(table_spaces)),
+            Arc::new(BooleanArray::from(has_index)),
+            Arc::new(BooleanArray::from(has_rules)),
+            Arc::new(BooleanArray::from(has_triggers)),
+            Arc::new(BooleanArray::from(row_security)),
+        ];
+
+        // Create a record batch
+        let batch = RecordBatch::try_new(this.schema.clone(), arrays)?;
+
+        Ok(batch)
+    }
+}
+
+impl PartitionStream for PgTablesTable {
+    fn schema(&self) -> &SchemaRef {
+        &self.schema
+    }
+
+    fn execute(&self, _ctx: Arc<TaskContext>) -> SendableRecordBatchStream {
+        let this = self.clone();
+        Box::pin(RecordBatchStreamAdapter::new(
+            this.schema.clone(),
+            futures::stream::once(async move { PgTablesTable::get_data(this).await }),
+        ))
+    }
+}