Merge branch 'master' into refactor/pg-catalog-record-batch-access

sunng87 · sunng87 · commit a43f3678674a · 2025-09-24T05:53:18.000+08:00
diff --git a/datafusion-postgres-cli/src/main.rs b/datafusion-postgres-cli/src/main.rs
@@ -6,6 +6,7 @@ use datafusion::execution::options::{
     ArrowReadOptions, AvroReadOptions, CsvReadOptions, NdJsonReadOptions, ParquetReadOptions,
 };
 use datafusion::prelude::{SessionConfig, SessionContext};
+use datafusion_postgres::auth::AuthManager;
 use datafusion_postgres::pg_catalog::setup_pg_catalog;
 use datafusion_postgres::{serve, ServerOptions};
 use env_logger::Env;
@@ -124,6 +125,7 @@ impl Opt {
 async fn setup_session_context(
     session_context: &SessionContext,
     opts: &Opt,
+    auth_manager: Arc<AuthManager>,
 ) -> Result<(), Box<dyn std::error::Error>> {
     // Register CSV tables
     for (table_name, table_path) in opts.csv_tables.iter().map(|s| parse_table_def(s.as_ref())) {
@@ -179,7 +181,7 @@ async fn setup_session_context(
     }
 
     // Register pg_catalog
-    setup_pg_catalog(session_context, "datafusion")?;
+    setup_pg_catalog(session_context, "datafusion", auth_manager)?;
 
     Ok(())
 }
@@ -196,16 +198,17 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
 
     let session_config = SessionConfig::new().with_information_schema(true);
     let session_context = SessionContext::new_with_config(session_config);
+    let auth_manager = Arc::new(AuthManager::new());
 
-    setup_session_context(&session_context, &opts).await?;
+    setup_session_context(&session_context, &opts, Arc::clone(&auth_manager)).await?;
 
     let server_options = ServerOptions::new()
         .with_host(opts.host)
         .with_port(opts.port)
         .with_tls_cert_path(opts.tls_cert)
         .with_tls_key_path(opts.tls_key);
 
-    serve(Arc::new(session_context), &server_options)
+    serve(Arc::new(session_context), &server_options, auth_manager)
         .await
         .map_err(|e| format!("Failed to run server: {e}"))?;
 
diff --git a/datafusion-postgres/src/lib.rs b/datafusion-postgres/src/lib.rs
@@ -83,10 +83,8 @@ fn setup_tls(cert_path: &str, key_path: &str) -> Result<TlsAcceptor, IOError> {
 pub async fn serve(
     session_context: Arc<SessionContext>,
     opts: &ServerOptions,
+    auth_manager: Arc<AuthManager>,
 ) -> Result<(), std::io::Error> {
-    // Create authentication manager
-    let auth_manager = Arc::new(AuthManager::new());
-
     // Create the handler factory with authentication
     let factory = Arc::new(HandlerFactory::new(session_context, auth_manager));
 
diff --git a/datafusion-postgres/src/pg_catalog.rs b/datafusion-postgres/src/pg_catalog.rs
@@ -4,7 +4,8 @@ use std::sync::Arc;
 
 use async_trait::async_trait;
 use datafusion::arrow::array::{
-    as_boolean_array, ArrayRef, AsArray, BooleanBuilder, RecordBatch, StringArray, StringBuilder,
+    as_boolean_array, ArrayRef, AsArray, BooleanBuilder, Int32Builder, RecordBatch, StringArray,
+    StringBuilder,
 };
 use datafusion::arrow::datatypes::{DataType, Field, Int32Type, SchemaRef};
 use datafusion::arrow::ipc::reader::FileReader;
@@ -19,6 +20,7 @@ use datafusion::prelude::{create_udf, Expr, SessionContext};
 use postgres_types::Oid;
 use tokio::sync::RwLock;
 
+use crate::auth::AuthManager;
 use crate::pg_catalog::catalog_info::CatalogInfo;
 use crate::pg_catalog::empty_table::EmptyTable;
 
@@ -32,7 +34,9 @@ pub mod pg_database;
 pub mod pg_get_expr_udf;
 pub mod pg_namespace;
 pub mod pg_replication_slot;
+pub mod pg_roles;
 pub mod pg_settings;
+pub mod pg_stat_gssapi;
 pub mod pg_tables;
 pub mod pg_views;
 
@@ -100,7 +104,9 @@ 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_ROLES: &str = "pg_roles";
 const PG_CATALOG_VIEW_PG_TABLES: &str = "pg_tables";
+const PG_CATALOG_VIEW_PG_STAT_GSSAPI: &str = "pg_stat_gssapi";
 const PG_CATALOG_VIEW_PG_STAT_USER_TABLES: &str = "pg_stat_user_tables";
 const PG_CATALOG_VIEW_PG_REPLICATION_SLOTS: &str = "pg_replication_slots";
 
@@ -166,7 +172,10 @@ pub const PG_CATALOG_TABLES: &[&str] = &[
     PG_CATALOG_TABLE_PG_TABLESPACE,
     PG_CATALOG_TABLE_PG_TRIGGER,
     PG_CATALOG_TABLE_PG_USER_MAPPING,
+    PG_CATALOG_VIEW_PG_ROLES,
     PG_CATALOG_VIEW_PG_SETTINGS,
+    PG_CATALOG_VIEW_PG_STAT_GSSAPI,
+    PG_CATALOG_VIEW_PG_TABLES,
     PG_CATALOG_VIEW_PG_VIEWS,
     PG_CATALOG_VIEW_PG_MATVIEWS,
     PG_CATALOG_VIEW_PG_STAT_USER_TABLES,
@@ -188,6 +197,7 @@ pub struct PgCatalogSchemaProvider<C> {
     oid_counter: Arc<AtomicU32>,
     oid_cache: Arc<RwLock<HashMap<OidCacheKey, Oid>>>,
     static_tables: Arc<PgCatalogStaticTables>,
+    auth_manager: Arc<AuthManager>,
 }
 
 #[async_trait]
@@ -218,12 +228,14 @@ impl<C: CatalogInfo> PgCatalogSchemaProvider<C> {
     pub fn try_new(
         catalog_list: C,
         static_tables: Arc<PgCatalogStaticTables>,
+        auth_manager: Arc<AuthManager>,
     ) -> Result<PgCatalogSchemaProvider<C>> {
         Ok(Self {
             catalog_list,
             oid_counter: Arc::new(AtomicU32::new(16384)),
             oid_cache: Arc::new(RwLock::new(HashMap::new())),
             static_tables,
+            auth_manager,
         })
     }
 
@@ -395,6 +407,16 @@ impl<C: CatalogInfo> PgCatalogSchemaProvider<C> {
                 let table = Arc::new(pg_settings::PgSettingsView::new());
                 Ok(Some(PgCatalogTable::Dynamic(table)))
             }
+
+            PG_CATALOG_VIEW_PG_STAT_GSSAPI => {
+                let table = Arc::new(pg_stat_gssapi::PgStatGssApiTable::new());
+                Ok(Some(PgCatalogTable::Dynamic(table)))
+            }
+            PG_CATALOG_VIEW_PG_ROLES => {
+                let table = Arc::new(pg_roles::PgRolesTable::new(Arc::clone(&self.auth_manager)));
+                Ok(Some(PgCatalogTable::Dynamic(table)))
+            }
+
             PG_CATALOG_VIEW_PG_VIEWS => Ok(Some(pg_views::pg_views().into())),
             PG_CATALOG_VIEW_PG_MATVIEWS => Ok(Some(pg_views::pg_matviews().into())),
             PG_CATALOG_VIEW_PG_STAT_USER_TABLES => Ok(Some(pg_views::pg_stat_user_tables().into())),
@@ -1236,15 +1258,36 @@ pub fn create_pg_encoding_to_char_udf() -> ScalarUDF {
     )
 }
 
+pub fn create_pg_backend_pid_udf() -> ScalarUDF {
+    let func = move |_args: &[ColumnarValue]| {
+        let mut builder = Int32Builder::new();
+        builder.append_value(BACKEND_PID);
+        let array: ArrayRef = Arc::new(builder.finish());
+        Ok(ColumnarValue::Array(array))
+    };
+
+    create_udf(
+        "pg_backend_pid",
+        vec![],
+        DataType::Int32,
+        Volatility::Stable,
+        Arc::new(func),
+    )
+}
+
+const BACKEND_PID: i32 = 1;
+
 /// Install pg_catalog and postgres UDFs to current `SessionContext`
 pub fn setup_pg_catalog(
     session_context: &SessionContext,
     catalog_name: &str,
+    auth_manager: Arc<AuthManager>,
 ) -> Result<(), Box<DataFusionError>> {
     let static_tables = Arc::new(PgCatalogStaticTables::try_new()?);
     let pg_catalog = PgCatalogSchemaProvider::try_new(
         session_context.state().catalog_list().clone(),
         static_tables.clone(),
+        auth_manager,
     )?;
     session_context
         .catalog(catalog_name)
@@ -1281,6 +1324,7 @@ pub fn setup_pg_catalog(
     session_context.register_udf(create_pg_relation_is_publishable_udf());
     session_context.register_udf(create_pg_get_statisticsobjdef_columns_udf());
     session_context.register_udf(create_pg_encoding_to_char_udf());
+    session_context.register_udf(create_pg_backend_pid_udf());
 
     Ok(())
 }
diff --git a/datafusion-postgres/src/pg_catalog/pg_roles.rs b/datafusion-postgres/src/pg_catalog/pg_roles.rs
@@ -0,0 +1,143 @@
+use std::sync::Arc;
+
+use datafusion::arrow::array::{
+    ArrayRef, BooleanArray, Int32Array, ListBuilder, RecordBatch, StringArray, StringBuilder,
+    TimestampMicrosecondBuilder,
+};
+use datafusion::arrow::datatypes::{DataType, Field, Schema, SchemaRef, TimeUnit};
+use datafusion::error::Result;
+use datafusion::execution::{SendableRecordBatchStream, TaskContext};
+use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
+use datafusion::physical_plan::streaming::PartitionStream;
+
+use crate::auth::AuthManager;
+
+#[derive(Debug, Clone)]
+pub(crate) struct PgRolesTable {
+    schema: SchemaRef,
+    auth_manager: Arc<AuthManager>,
+}
+
+impl PgRolesTable {
+    pub(crate) fn new(auth_manager: Arc<AuthManager>) -> Self {
+        let schema = Arc::new(Schema::new(vec![
+            Field::new("rolname", DataType::Utf8, true),
+            Field::new("rolsuper", DataType::Boolean, true),
+            Field::new("rolinherit", DataType::Boolean, true),
+            Field::new("rolcreaterole", DataType::Boolean, true),
+            Field::new("rolcreatedb", DataType::Boolean, true),
+            Field::new("rolcanlogin", DataType::Boolean, true),
+            Field::new("rolreplication", DataType::Boolean, true),
+            Field::new("rolconnlimit", DataType::Int32, true),
+            Field::new("rolpassword", DataType::Utf8, true),
+            Field::new(
+                "rolvaliduntil",
+                DataType::Timestamp(
+                    datafusion::arrow::datatypes::TimeUnit::Microsecond,
+                    Some(Arc::from("UTC")),
+                ),
+                true,
+            ),
+            Field::new("rolbypassrls", DataType::Boolean, true),
+            Field::new(
+                "rolconfig",
+                DataType::List(Arc::new(Field::new("item", DataType::Utf8, true))),
+                true,
+            ),
+            Field::new("oid", DataType::Int32, true),
+        ]));
+
+        Self {
+            schema,
+            auth_manager,
+        }
+    }
+
+    async fn get_data(this: Self) -> Result<RecordBatch> {
+        let mut rolname = Vec::new();
+        let mut rolsuper = Vec::new();
+        let mut rolinherit = Vec::new();
+        let mut rolcreaterole = Vec::new();
+        let mut rolcreatedb = Vec::new();
+        let mut rolcanlogin = Vec::new();
+        let mut rolreplication = Vec::new();
+        let mut rolconnlimit = Vec::new();
+        let mut rolpassword: Vec<Option<String>> = Vec::new();
+        let mut rolvaliduntil: Vec<Option<i64>> = Vec::new();
+        let mut rolbypassrls = Vec::new();
+        let mut rolconfig: Vec<Option<Vec<String>>> = Vec::new();
+        let mut oid: Vec<i32> = Vec::new();
+
+        for role_name in &this.auth_manager.list_roles().await {
+            let role = &this.auth_manager.get_role(role_name).await.unwrap();
+            rolname.push(role.name.clone());
+            rolsuper.push(role.is_superuser);
+            rolinherit.push(true);
+            rolcreaterole.push(role.can_create_role);
+            rolcreatedb.push(role.can_create_db);
+            rolcanlogin.push(role.can_login);
+            rolreplication.push(role.can_replication);
+            rolconnlimit.push(-1);
+            rolpassword.push(None);
+            rolvaliduntil.push(None);
+            rolbypassrls.push(None);
+            rolconfig.push(None);
+            oid.push(0); // TODO: handle oid properly somehow
+        }
+
+        let arrays: Vec<ArrayRef> = vec![
+            Arc::new(StringArray::from(rolname)),
+            Arc::new(BooleanArray::from(rolsuper)),
+            Arc::new(BooleanArray::from(rolinherit)),
+            Arc::new(BooleanArray::from(rolcreaterole)),
+            Arc::new(BooleanArray::from(rolcreatedb)),
+            Arc::new(BooleanArray::from(rolcanlogin)),
+            Arc::new(BooleanArray::from(rolreplication)),
+            Arc::new(Int32Array::from(rolconnlimit)),
+            Arc::new(StringArray::from(rolpassword)),
+            Arc::new({
+                let mut builder =
+                    TimestampMicrosecondBuilder::with_capacity(rolconfig.len()).with_data_type(
+                        DataType::Timestamp(TimeUnit::Microsecond, Some(Arc::from("UTC"))),
+                    );
+                for field in &rolvaliduntil {
+                    builder.append_option(field.as_ref().copied());
+                }
+                builder.finish()
+            }),
+            Arc::new(BooleanArray::from(rolbypassrls)),
+            Arc::new({
+                let mut builder = ListBuilder::new(StringBuilder::new());
+                for field in &rolconfig {
+                    match field {
+                        Some(values) => {
+                            for s in values {
+                                builder.values().append_value(s);
+                            }
+                            builder.append(true);
+                        }
+                        None => builder.append(false),
+                    }
+                }
+                builder.finish()
+            }),
+            Arc::new(Int32Array::from(oid)),
+        ];
+
+        Ok(RecordBatch::try_new(this.schema.clone(), arrays)?)
+    }
+}
+
+impl PartitionStream for PgRolesTable {
+    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 { PgRolesTable::get_data(this).await }),
+        ))
+    }
+}
diff --git a/datafusion-postgres/src/pg_catalog/pg_stat_gssapi.rs b/datafusion-postgres/src/pg_catalog/pg_stat_gssapi.rs
diff --git a/datafusion-postgres/tests/common/mod.rs b/datafusion-postgres/tests/common/mod.rs
