From 690a4990c0a9c1ac0fdba9b1c99bdba1f3021703 Mon Sep 17 00:00:00 2001 From: Oluwapeluwa Ibrahim Date: Tue, 16 Dec 2025 07:43:45 +0100 Subject: [PATCH] refactor: decouple AuthManager from pg_catalog Fixes #191 This change decouples AuthManager from PgCatalogContextProvider to unblock issue #189 (moving pg_catalog to a separate module). Changes: - Add RoleProvider trait for providing role information to pg_catalog - Implement RoleProvider for AuthManager and Arc - Add RoleProviderBridge adapter that wraps RoleProvider and implements PgCatalogContextProvider - this bridges the two worlds without tight coupling - Remove direct PgCatalogContextProvider impl from AuthManager - Update testing.rs and CLI to use RoleProviderBridge The architecture now looks like: AuthManager -> RoleProvider trait RoleProviderBridge -> PgCatalogContextProvider trait pg_catalog uses PgCatalogContextProvider for pg_roles table This keeps auth concerns separate from pg_catalog while still allowing role information to flow to the pg_roles table. --- datafusion-postgres-cli/src/main.rs | 9 ++-- datafusion-postgres/src/auth.rs | 66 ++++++++++++++++++++++++++--- datafusion-postgres/src/testing.rs | 15 ++++--- 3 files changed, 73 insertions(+), 17 deletions(-) diff --git a/datafusion-postgres-cli/src/main.rs b/datafusion-postgres-cli/src/main.rs index 7b0f296..e95fc23 100644 --- a/datafusion-postgres-cli/src/main.rs +++ b/datafusion-postgres-cli/src/main.rs @@ -6,7 +6,7 @@ use datafusion::execution::options::{ ArrowReadOptions, AvroReadOptions, CsvReadOptions, NdJsonReadOptions, ParquetReadOptions, }; use datafusion::prelude::{SessionConfig, SessionContext}; -use datafusion_postgres::auth::AuthManager; +use datafusion_postgres::auth::{AuthManager, RoleProviderBridge}; use datafusion_postgres::datafusion_pg_catalog::setup_pg_catalog; use datafusion_postgres::{serve, ServerOptions}; use env_logger::Env; @@ -180,8 +180,9 @@ async fn setup_session_context( info!("Loaded {table_path} as table {table_name}"); } - // Register pg_catalog - setup_pg_catalog(session_context, "datafusion", auth_manager)?; + // Register pg_catalog using RoleProviderBridge for decoupled auth integration + let role_bridge = RoleProviderBridge::new(auth_manager); + setup_pg_catalog(session_context, "datafusion", role_bridge)?; Ok(()) } @@ -199,7 +200,7 @@ async fn main() -> Result<(), Box> { let session_config = SessionConfig::new().with_information_schema(true); let session_context = SessionContext::new_with_config(session_config); - // TODO: remove or replace AuthManager for pg_catalog + // Create AuthManager for pg_catalog role information let auth_manager = Arc::new(AuthManager::new()); setup_session_context(&session_context, &opts, Arc::clone(&auth_manager)).await?; diff --git a/datafusion-postgres/src/auth.rs b/datafusion-postgres/src/auth.rs index 034bf3f..54e1a02 100644 --- a/datafusion-postgres/src/auth.rs +++ b/datafusion-postgres/src/auth.rs @@ -6,7 +6,32 @@ use pgwire::api::auth::{AuthSource, LoginInfo, Password}; use pgwire::error::{PgWireError, PgWireResult}; use tokio::sync::RwLock; -use datafusion_pg_catalog::pg_catalog::context::*; +// Re-export Role and related types from pg_catalog context for use by other modules +pub use datafusion_pg_catalog::pg_catalog::context::{ + Grant, Permission, ResourceType, Role, RoleConfig, User, +}; + +/// Trait for providing role information to pg_catalog +/// This decouples AuthManager from PgCatalogContextProvider +#[async_trait] +pub trait RoleProvider: Send + Sync { + /// List all role names + async fn list_roles(&self) -> Vec; + /// Get role details by name + async fn get_role(&self, name: &str) -> Option; +} + +/// Implement RoleProvider for Arc to allow shared ownership +#[async_trait] +impl RoleProvider for Arc { + async fn list_roles(&self) -> Vec { + (**self).list_roles().await + } + + async fn get_role(&self, name: &str) -> Option { + (**self).get_role(name).await + } +} /// Authentication manager that handles users and roles #[derive(Debug, Clone)] @@ -446,15 +471,44 @@ impl AuthManager { } #[async_trait] -impl PgCatalogContextProvider for AuthManager { - // retrieve all database role names +impl RoleProvider for AuthManager { + async fn list_roles(&self) -> Vec { + let roles = self.roles.read().await; + roles.keys().cloned().collect() + } + + async fn get_role(&self, name: &str) -> Option { + let roles = self.roles.read().await; + roles.get(name).cloned() + } +} + +/// Bridge adapter that implements PgCatalogContextProvider using a RoleProvider +/// This allows decoupling of AuthManager from pg_catalog while still providing +/// role information to pg_roles table +#[derive(Debug, Clone)] +pub struct RoleProviderBridge { + role_provider: R, +} + +impl RoleProviderBridge { + pub fn new(role_provider: R) -> Self { + Self { role_provider } + } +} + +use datafusion_pg_catalog::pg_catalog::context::PgCatalogContextProvider; + +#[async_trait] +impl PgCatalogContextProvider + for RoleProviderBridge +{ async fn roles(&self) -> Vec { - self.list_roles().await + self.role_provider.list_roles().await } - // retrieve database role information async fn role(&self, name: &str) -> Option { - self.get_role(name).await + self.role_provider.get_role(name).await } } diff --git a/datafusion-postgres/src/testing.rs b/datafusion-postgres/src/testing.rs index f2d53db..545b225 100644 --- a/datafusion-postgres/src/testing.rs +++ b/datafusion-postgres/src/testing.rs @@ -10,18 +10,19 @@ use pgwire::{ }, }; -use crate::{auth::AuthManager, DfSessionService}; +use crate::auth::{AuthManager, RoleProviderBridge}; +use crate::DfSessionService; pub fn setup_handlers() -> DfSessionService { let session_config = SessionConfig::new().with_information_schema(true); let session_context = SessionContext::new_with_config(session_config); - setup_pg_catalog( - &session_context, - "datafusion", - Arc::new(AuthManager::default()), - ) - .expect("Failed to setup sesession context"); + // Use RoleProviderBridge to decouple AuthManager from pg_catalog + let auth_manager = Arc::new(AuthManager::default()); + let role_bridge = RoleProviderBridge::new(auth_manager); + + setup_pg_catalog(&session_context, "datafusion", role_bridge) + .expect("Failed to setup sesession context"); DfSessionService::new(Arc::new(session_context)) }