datafusion_postgres/hooks/
permissions.rs

1use std::sync::Arc;
2
3use async_trait::async_trait;
4use datafusion::common::ParamValues;
5use datafusion::logical_expr::LogicalPlan;
6use datafusion::prelude::SessionContext;
7use datafusion::sql::sqlparser::ast::Statement;
8use pgwire::api::ClientInfo;
9use pgwire::api::results::Response;
10use pgwire::error::{PgWireError, PgWireResult};
11
12use crate::QueryHook;
13use crate::auth::AuthManager;
14use crate::hooks::HookClient;
15
16use datafusion_pg_catalog::pg_catalog::context::{Permission, ResourceType};
17
18#[derive(Debug)]
19pub struct PermissionsHook {
20    auth_manager: Arc<AuthManager>,
21}
22
23impl PermissionsHook {
24    pub fn new(auth_manager: Arc<AuthManager>) -> Self {
25        PermissionsHook { auth_manager }
26    }
27
28    /// Check if the current user has permission to execute a statement
29    async fn check_statement_permission<C>(
30        &self,
31        client: &C,
32        statement: &Statement,
33    ) -> PgWireResult<()>
34    where
35        C: ClientInfo + ?Sized,
36    {
37        // Get the username from client metadata
38        let username = client
39            .metadata()
40            .get("user")
41            .map(|s| s.as_str())
42            .unwrap_or("anonymous");
43
44        // Determine required permissions based on Statement type
45        let (required_permission, resource) = match statement {
46            Statement::Query(_) => (Permission::Select, ResourceType::All),
47            Statement::Insert(_) => (Permission::Insert, ResourceType::All),
48            Statement::Update { .. } => (Permission::Update, ResourceType::All),
49            Statement::Delete(_) => (Permission::Delete, ResourceType::All),
50            Statement::CreateTable { .. } | Statement::CreateView { .. } => {
51                (Permission::Create, ResourceType::All)
52            }
53            Statement::Drop { .. } => (Permission::Drop, ResourceType::All),
54            Statement::AlterTable { .. } => (Permission::Alter, ResourceType::All),
55            // For other statements (SET, SHOW, EXPLAIN, transactions, etc.), allow all users
56            _ => return Ok(()),
57        };
58
59        // Check permission
60        let has_permission = self
61            .auth_manager
62            .check_permission(username, required_permission, resource)
63            .await;
64
65        if !has_permission {
66            return Err(PgWireError::UserError(Box::new(
67                pgwire::error::ErrorInfo::new(
68                    "ERROR".to_string(),
69                    "42501".to_string(), // insufficient_privilege
70                    format!("permission denied for user \"{username}\""),
71                ),
72            )));
73        }
74
75        Ok(())
76    }
77
78    /// Check if a statement should skip permission checks
79    fn should_skip_permission_check(statement: &Statement) -> bool {
80        matches!(
81            statement,
82            Statement::Set { .. }
83                | Statement::ShowVariable { .. }
84                | Statement::ShowStatus { .. }
85                | Statement::ShowCatalogs { .. }
86                | Statement::StartTransaction { .. }
87                | Statement::Commit { .. }
88                | Statement::Rollback { .. }
89                | Statement::Savepoint { .. }
90                | Statement::ReleaseSavepoint { .. }
91        )
92    }
93}
94
95#[async_trait]
96impl QueryHook for PermissionsHook {
97    /// called in simple query handler to return response directly
98    async fn handle_simple_query(
99        &self,
100        statement: &Statement,
101        _session_context: &SessionContext,
102        client: &mut dyn HookClient,
103    ) -> Option<PgWireResult<Response>> {
104        if Self::should_skip_permission_check(statement) {
105            return None;
106        }
107
108        // Check permissions for other statements
109        if let Err(e) = self.check_statement_permission(&*client, statement).await {
110            return Some(Err(e));
111        }
112
113        None
114    }
115
116    async fn handle_extended_parse_query(
117        &self,
118        _stmt: &Statement,
119        _session_context: &SessionContext,
120        _client: &(dyn ClientInfo + Send + Sync),
121    ) -> Option<PgWireResult<LogicalPlan>> {
122        None
123    }
124
125    async fn handle_extended_query(
126        &self,
127        statement: &Statement,
128        _logical_plan: &LogicalPlan,
129        _params: &ParamValues,
130        _session_context: &SessionContext,
131        client: &mut dyn HookClient,
132    ) -> Option<PgWireResult<Response>> {
133        if Self::should_skip_permission_check(statement) {
134            return None;
135        }
136
137        // Check permissions for other statements
138        if let Err(e) = self.check_statement_permission(&*client, statement).await {
139            return Some(Err(e));
140        }
141
142        None
143    }
144}
datafusion_postgres/hooks/permissions.rs

datafusion_postgres/hooks/
permissions.rs
