Add support for pg_roles

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}"))?;
 

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));
 

datafusion-postgres/src/pg_catalog.rs

@@ -20,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;
 
 pub mod catalog_info;
@@ -32,6 +33,7 @@ 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;
@@ -101,6 +103,7 @@ 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";
@@ -190,6 +193,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]
@@ -349,6 +353,13 @@ impl<C: CatalogInfo> SchemaProvider for PgCatalogSchemaProvider<C> {
                     vec![table],
                 )?)))
             }
+            PG_CATALOG_VIEW_PG_ROLES => {
+                let table = Arc::new(pg_roles::PgRolesTable::new(Arc::clone(&self.auth_manager)));
+                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(
@@ -382,12 +393,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,
         })
     }
 }
@@ -1195,11 +1208,13 @@ const BACKEND_PID: i32 = 1;
 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)

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 }),
+        ))
+    }
+}

datafusion-postgres/tests/common/mod.rs

@@ -12,7 +12,12 @@ use pgwire::{
 
 pub fn setup_handlers() -> DfSessionService {
     let session_context = SessionContext::new();
-    setup_pg_catalog(&session_context, "datafusion").expect("Failed to setup sesession context");
+    setup_pg_catalog(
+        &session_context,
+        "datafusion",
+        Arc::new(AuthManager::default()),
+    )
+    .expect("Failed to setup sesession context");
 
     DfSessionService::new(Arc::new(session_context), Arc::new(AuthManager::new()))
 }