From 37b2f72814632d089b11b7892bdc66d518b7428d Mon Sep 17 00:00:00 2001 From: Oluwapeluwa Ibrahim Date: Mon, 15 Dec 2025 11:19:20 +0100 Subject: [PATCH 1/4] refactor: use AST pattern matching instead of string matching for statement detection Fixes #159 The codebase was using starts_with() string checks to detect statement types (SET, SHOW, INSERT, etc.) which is fragile - it can break with leading whitespace, comments, or case variations. Since we already have the parsed Statement AST from sqlparser, this switches to using pattern matching directly on the Statement enum. Changes in permissions.rs: - Renamed check_query_permission to check_statement_permission, now takes &Statement instead of &str - Permission detection (SELECT/INSERT/UPDATE/DELETE/CREATE/DROP/ALTER) now uses match on Statement variants - Added should_skip_permission_check() helper using matches!() macro for SET, SHOW, and transaction statements - Removed all the to_lowercase().starts_with() chains Changes in handlers.rs: - INSERT detection now uses matches!(statement, Statement::Insert(_)) - Removed unnecessary query_lower variable construction - Fixed extended query handler to properly destructure statement from the portal tuple All 12 unit tests pass. The existing integration test failures (dbeaver, metabase, psql) are unrelated - they fail due to missing DataFusion functions like array_length and array_contains. --- datafusion-postgres/src/handlers.rs | 16 +--- datafusion-postgres/src/hooks/permissions.rs | 98 ++++++++++---------- 2 files changed, 54 insertions(+), 60 deletions(-) diff --git a/datafusion-postgres/src/handlers.rs b/datafusion-postgres/src/handlers.rs index 2f2c349..b3880ca 100644 --- a/datafusion-postgres/src/handlers.rs +++ b/datafusion-postgres/src/handlers.rs @@ -139,9 +139,7 @@ impl SimpleQueryHandler for DfSessionService { let mut results = vec![]; 'stmt: for statement in statements { - // TODO: improve statement check by using statement directly let query = statement.to_string(); - let query_lower = query.to_lowercase().trim().to_string(); // Call query hooks with the parsed statement for hook in &self.query_hooks { @@ -179,7 +177,7 @@ impl SimpleQueryHandler for DfSessionService { } }; - if query_lower.starts_with("insert into") { + if matches!(statement, sqlparser::ast::Statement::Insert(_)) { let resp = map_rows_affected_for_insert(&df).await?; results.push(resp); } else { @@ -265,13 +263,7 @@ impl ExtendedQueryHandler for DfSessionService { where C: ClientInfo + Unpin + Send + Sync, { - let query = portal - .statement - .statement - .0 - .to_lowercase() - .trim() - .to_string(); + let query = portal.statement.statement.0.to_string(); log::debug!("Received execute extended query: {query}"); // Log for debugging // Check query hooks first @@ -302,7 +294,7 @@ impl ExtendedQueryHandler for DfSessionService { } } - if let (_, Some((_, plan))) = &portal.statement.statement { + if let (_, Some((statement, plan))) = &portal.statement.statement { let param_types = plan .get_parameter_types() .map_err(|e| PgWireError::ApiError(Box::new(e)))?; @@ -345,7 +337,7 @@ impl ExtendedQueryHandler for DfSessionService { } }; - if query.starts_with("insert into") { + if matches!(statement, sqlparser::ast::Statement::Insert(_)) { let resp = map_rows_affected_for_insert(&dataframe).await?; Ok(resp) diff --git a/datafusion-postgres/src/hooks/permissions.rs b/datafusion-postgres/src/hooks/permissions.rs index ac59348..aa49b65 100644 --- a/datafusion-postgres/src/hooks/permissions.rs +++ b/datafusion-postgres/src/hooks/permissions.rs @@ -23,8 +23,12 @@ impl PermissionsHook { PermissionsHook { auth_manager } } - /// Check if the current user has permission to execute a query - async fn check_query_permission(&self, client: &C, query: &str) -> PgWireResult<()> + /// Check if the current user has permission to execute a statement + async fn check_statement_permission( + &self, + client: &C, + statement: &Statement, + ) -> PgWireResult<()> where C: ClientInfo + ?Sized, { @@ -35,29 +39,19 @@ impl PermissionsHook { .map(|s| s.as_str()) .unwrap_or("anonymous"); - // Parse query to determine required permissions - let query_lower = query.to_lowercase(); - let query_trimmed = query_lower.trim(); - - let (required_permission, resource) = if query_trimmed.starts_with("select") { - (Permission::Select, ResourceType::All) - } else if query_trimmed.starts_with("insert") { - (Permission::Insert, ResourceType::All) - } else if query_trimmed.starts_with("update") { - (Permission::Update, ResourceType::All) - } else if query_trimmed.starts_with("delete") { - (Permission::Delete, ResourceType::All) - } else if query_trimmed.starts_with("create table") - || query_trimmed.starts_with("create view") - { - (Permission::Create, ResourceType::All) - } else if query_trimmed.starts_with("drop") { - (Permission::Drop, ResourceType::All) - } else if query_trimmed.starts_with("alter") { - (Permission::Alter, ResourceType::All) - } else { - // For other queries (SHOW, EXPLAIN, etc.), allow all users - return Ok(()); + // Determine required permissions based on Statement type + let (required_permission, resource) = match statement { + Statement::Query(_) => (Permission::Select, ResourceType::All), + Statement::Insert(_) => (Permission::Insert, ResourceType::All), + Statement::Update { .. } => (Permission::Update, ResourceType::All), + Statement::Delete(_) => (Permission::Delete, ResourceType::All), + Statement::CreateTable { .. } | Statement::CreateView { .. } => { + (Permission::Create, ResourceType::All) + } + Statement::Drop { .. } => (Permission::Drop, ResourceType::All), + Statement::AlterTable { .. } => (Permission::Alter, ResourceType::All), + // For other statements (SET, SHOW, EXPLAIN, transactions, etc.), allow all users + _ => return Ok(()), }; // Check permission @@ -78,6 +72,21 @@ impl PermissionsHook { Ok(()) } + + /// Check if a statement should skip permission checks + fn should_skip_permission_check(statement: &Statement) -> bool { + matches!( + statement, + Statement::Set { .. } + | Statement::ShowVariable { .. } + | Statement::ShowStatus { .. } + | Statement::StartTransaction { .. } + | Statement::Commit { .. } + | Statement::Rollback { .. } + | Statement::Savepoint { .. } + | Statement::ReleaseSavepoint { .. } + ) + } } #[async_trait] @@ -89,22 +98,14 @@ impl QueryHook for PermissionsHook { _session_context: &SessionContext, client: &mut (dyn ClientInfo + Send + Sync), ) -> Option> { - let query_lower = statement.to_string().to_lowercase(); - - // Check permissions for the query (skip for SET, transaction, and SHOW statements) - if !query_lower.starts_with("set") - && !query_lower.starts_with("begin") - && !query_lower.starts_with("commit") - && !query_lower.starts_with("rollback") - && !query_lower.starts_with("start") - && !query_lower.starts_with("end") - && !query_lower.starts_with("abort") - && !query_lower.starts_with("show") - { - let res = self.check_query_permission(&*client, &query_lower).await; - if let Err(e) = res { - return Some(Err(e)); - } + // Skip permission checks for SET, SHOW, and transaction statements + if Self::should_skip_permission_check(statement) { + return None; + } + + // Check permissions for other statements + if let Err(e) = self.check_statement_permission(&*client, statement).await { + return Some(Err(e)); } None @@ -127,15 +128,16 @@ impl QueryHook for PermissionsHook { _session_context: &SessionContext, client: &mut (dyn ClientInfo + Send + Sync), ) -> Option> { - let query = statement.to_string().to_lowercase(); + // Skip permission checks for SET and SHOW statements + if Self::should_skip_permission_check(statement) { + return None; + } - // Check permissions for the query (skip for SET and SHOW statements) - if !query.starts_with("set") && !query.starts_with("show") { - let res = self.check_query_permission(&*client, &query).await; - if let Err(e) = res { - return Some(Err(e)); - } + // Check permissions for other statements + if let Err(e) = self.check_statement_permission(&*client, statement).await { + return Some(Err(e)); } + None } } From 49c9132ab3da9c81f0452811e823ffbb6bd1b07a Mon Sep 17 00:00:00 2001 From: Oluwapeluwa Ibrahim Date: Mon, 15 Dec 2025 11:46:23 +0100 Subject: [PATCH 2/4] test: add pgAdmin startup queries test for issue #178 Adds a test file to verify pgAdmin startup queries work correctly. The test covers: - SELECT version() query - The CASE expression checking pg_extension for 'bdr' extension and pg_replication_slots for replication slot count Both pg_extension and pg_replication_slots tables are already implemented in pg_catalog, so these queries now pass. --- datafusion-postgres/tests/pgadmin.rs | 32 ++++++++++++++++++++++++++++ 1 file changed, 32 insertions(+) create mode 100644 datafusion-postgres/tests/pgadmin.rs diff --git a/datafusion-postgres/tests/pgadmin.rs b/datafusion-postgres/tests/pgadmin.rs new file mode 100644 index 0000000..cf846b1 --- /dev/null +++ b/datafusion-postgres/tests/pgadmin.rs @@ -0,0 +1,32 @@ +use pgwire::api::query::SimpleQueryHandler; + +use datafusion_postgres::testing::*; + +// pgAdmin startup queries from issue #178 +// https://github.com/datafusion-contrib/datafusion-postgres/issues/178 +const PGADMIN_QUERIES: &[&str] = &[ + // Basic version query (fixed by #179) + "SELECT version()", + // Query to check for BDR extension and replication slots + r#"SELECT CASE + WHEN (SELECT count(extname) FROM pg_catalog.pg_extension WHERE extname='bdr') > 0 + THEN 'pgd' + WHEN (SELECT COUNT(*) FROM pg_replication_slots) > 0 + THEN 'log' + ELSE NULL + END as type"#, +]; + +#[tokio::test] +pub async fn test_pgadmin_startup_sql() { + let service = setup_handlers(); + let mut client = MockClient::new(); + + for query in PGADMIN_QUERIES { + SimpleQueryHandler::do_query(&service, &mut client, query) + .await + .unwrap_or_else(|e| { + panic!("failed to run sql:\n--------------\n{query}\n--------------\n{e}") + }); + } +} From b6cdeb5bd288da4d0704079bc637c0fb0bd398b1 Mon Sep 17 00:00:00 2001 From: Oluwapeluwa Ibrahim Date: Mon, 15 Dec 2025 14:42:30 +0100 Subject: [PATCH 3/4] refactor: avoid String allocation in extended query handler Use reference instead of calling .to_string() on query string as suggested by mjgarton - the query is only used for logging and doesn't need to be owned. --- datafusion-postgres/src/handlers.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/datafusion-postgres/src/handlers.rs b/datafusion-postgres/src/handlers.rs index b3880ca..731dc8e 100644 --- a/datafusion-postgres/src/handlers.rs +++ b/datafusion-postgres/src/handlers.rs @@ -263,7 +263,7 @@ impl ExtendedQueryHandler for DfSessionService { where C: ClientInfo + Unpin + Send + Sync, { - let query = portal.statement.statement.0.to_string(); + let query = &portal.statement.statement.0; log::debug!("Received execute extended query: {query}"); // Log for debugging // Check query hooks first From 0b01b1ae6594f7c2c62012c8e6f5648840113076 Mon Sep 17 00:00:00 2001 From: Oluwapeluwa Ibrahim Date: Mon, 15 Dec 2025 19:26:39 +0100 Subject: [PATCH 4/4] refactor: remove redundant comments in permissions hook Remove comments that don't fully describe what's being skipped. The function name should_skip_permission_check is self-documenting. Suggested by mjgarton. --- datafusion-postgres/src/hooks/permissions.rs | 2 -- 1 file changed, 2 deletions(-) diff --git a/datafusion-postgres/src/hooks/permissions.rs b/datafusion-postgres/src/hooks/permissions.rs index aa49b65..ff5a7f6 100644 --- a/datafusion-postgres/src/hooks/permissions.rs +++ b/datafusion-postgres/src/hooks/permissions.rs @@ -98,7 +98,6 @@ impl QueryHook for PermissionsHook { _session_context: &SessionContext, client: &mut (dyn ClientInfo + Send + Sync), ) -> Option> { - // Skip permission checks for SET, SHOW, and transaction statements if Self::should_skip_permission_check(statement) { return None; } @@ -128,7 +127,6 @@ impl QueryHook for PermissionsHook { _session_context: &SessionContext, client: &mut (dyn ClientInfo + Send + Sync), ) -> Option> { - // Skip permission checks for SET and SHOW statements if Self::should_skip_permission_check(statement) { return None; }