Skip to content

Commit 690a499

Browse files
committed
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<T> - 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.
1 parent 2b51166 commit 690a499

File tree

3 files changed

+73
-17
lines changed

3 files changed

+73
-17
lines changed

datafusion-postgres-cli/src/main.rs

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@ use datafusion::execution::options::{
66
ArrowReadOptions, AvroReadOptions, CsvReadOptions, NdJsonReadOptions, ParquetReadOptions,
77
};
88
use datafusion::prelude::{SessionConfig, SessionContext};
9-
use datafusion_postgres::auth::AuthManager;
9+
use datafusion_postgres::auth::{AuthManager, RoleProviderBridge};
1010
use datafusion_postgres::datafusion_pg_catalog::setup_pg_catalog;
1111
use datafusion_postgres::{serve, ServerOptions};
1212
use env_logger::Env;
@@ -180,8 +180,9 @@ async fn setup_session_context(
180180
info!("Loaded {table_path} as table {table_name}");
181181
}
182182

183-
// Register pg_catalog
184-
setup_pg_catalog(session_context, "datafusion", auth_manager)?;
183+
// Register pg_catalog using RoleProviderBridge for decoupled auth integration
184+
let role_bridge = RoleProviderBridge::new(auth_manager);
185+
setup_pg_catalog(session_context, "datafusion", role_bridge)?;
185186

186187
Ok(())
187188
}
@@ -199,7 +200,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
199200
let session_config = SessionConfig::new().with_information_schema(true);
200201
let session_context = SessionContext::new_with_config(session_config);
201202

202-
// TODO: remove or replace AuthManager for pg_catalog
203+
// Create AuthManager for pg_catalog role information
203204
let auth_manager = Arc::new(AuthManager::new());
204205
setup_session_context(&session_context, &opts, Arc::clone(&auth_manager)).await?;
205206

datafusion-postgres/src/auth.rs

Lines changed: 60 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,32 @@ use pgwire::api::auth::{AuthSource, LoginInfo, Password};
66
use pgwire::error::{PgWireError, PgWireResult};
77
use tokio::sync::RwLock;
88

9-
use datafusion_pg_catalog::pg_catalog::context::*;
9+
// Re-export Role and related types from pg_catalog context for use by other modules
10+
pub use datafusion_pg_catalog::pg_catalog::context::{
11+
Grant, Permission, ResourceType, Role, RoleConfig, User,
12+
};
13+
14+
/// Trait for providing role information to pg_catalog
15+
/// This decouples AuthManager from PgCatalogContextProvider
16+
#[async_trait]
17+
pub trait RoleProvider: Send + Sync {
18+
/// List all role names
19+
async fn list_roles(&self) -> Vec<String>;
20+
/// Get role details by name
21+
async fn get_role(&self, name: &str) -> Option<Role>;
22+
}
23+
24+
/// Implement RoleProvider for Arc<T> to allow shared ownership
25+
#[async_trait]
26+
impl<T: RoleProvider + ?Sized> RoleProvider for Arc<T> {
27+
async fn list_roles(&self) -> Vec<String> {
28+
(**self).list_roles().await
29+
}
30+
31+
async fn get_role(&self, name: &str) -> Option<Role> {
32+
(**self).get_role(name).await
33+
}
34+
}
1035

1136
/// Authentication manager that handles users and roles
1237
#[derive(Debug, Clone)]
@@ -446,15 +471,44 @@ impl AuthManager {
446471
}
447472

448473
#[async_trait]
449-
impl PgCatalogContextProvider for AuthManager {
450-
// retrieve all database role names
474+
impl RoleProvider for AuthManager {
475+
async fn list_roles(&self) -> Vec<String> {
476+
let roles = self.roles.read().await;
477+
roles.keys().cloned().collect()
478+
}
479+
480+
async fn get_role(&self, name: &str) -> Option<Role> {
481+
let roles = self.roles.read().await;
482+
roles.get(name).cloned()
483+
}
484+
}
485+
486+
/// Bridge adapter that implements PgCatalogContextProvider using a RoleProvider
487+
/// This allows decoupling of AuthManager from pg_catalog while still providing
488+
/// role information to pg_roles table
489+
#[derive(Debug, Clone)]
490+
pub struct RoleProviderBridge<R: RoleProvider + Clone> {
491+
role_provider: R,
492+
}
493+
494+
impl<R: RoleProvider + Clone> RoleProviderBridge<R> {
495+
pub fn new(role_provider: R) -> Self {
496+
Self { role_provider }
497+
}
498+
}
499+
500+
use datafusion_pg_catalog::pg_catalog::context::PgCatalogContextProvider;
501+
502+
#[async_trait]
503+
impl<R: RoleProvider + Clone + std::fmt::Debug + 'static> PgCatalogContextProvider
504+
for RoleProviderBridge<R>
505+
{
451506
async fn roles(&self) -> Vec<String> {
452-
self.list_roles().await
507+
self.role_provider.list_roles().await
453508
}
454509

455-
// retrieve database role information
456510
async fn role(&self, name: &str) -> Option<Role> {
457-
self.get_role(name).await
511+
self.role_provider.get_role(name).await
458512
}
459513
}
460514

datafusion-postgres/src/testing.rs

Lines changed: 8 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -10,18 +10,19 @@ use pgwire::{
1010
},
1111
};
1212

13-
use crate::{auth::AuthManager, DfSessionService};
13+
use crate::auth::{AuthManager, RoleProviderBridge};
14+
use crate::DfSessionService;
1415

1516
pub fn setup_handlers() -> DfSessionService {
1617
let session_config = SessionConfig::new().with_information_schema(true);
1718
let session_context = SessionContext::new_with_config(session_config);
1819

19-
setup_pg_catalog(
20-
&session_context,
21-
"datafusion",
22-
Arc::new(AuthManager::default()),
23-
)
24-
.expect("Failed to setup sesession context");
20+
// Use RoleProviderBridge to decouple AuthManager from pg_catalog
21+
let auth_manager = Arc::new(AuthManager::default());
22+
let role_bridge = RoleProviderBridge::new(auth_manager);
23+
24+
setup_pg_catalog(&session_context, "datafusion", role_bridge)
25+
.expect("Failed to setup sesession context");
2526

2627
DfSessionService::new(Arc::new(session_context))
2728
}

0 commit comments

Comments
 (0)