From 5cd7f936f90978673a1f6a1dc64765e701035aa4 Mon Sep 17 00:00:00 2001 From: Conrad Ludgate Date: Fri, 9 May 2025 08:48:30 +0100 Subject: [PATCH] fix(neon-rls): optimistically assume role grants are already assigned for replicas (#11811) ## Problem Read replicas cannot grant permissions for roles for Neon RLS. Usually the permission is already granted, so we can optimistically check. See INC-509 ## Summary of changes Perform a permission lookup prior to actually executing any grants. --- Cargo.lock | 1 + compute_tools/Cargo.toml | 1 + compute_tools/src/compute.rs | 52 +++++++++++++++++-------- test_runner/fixtures/neon_fixtures.py | 10 ++++- test_runner/regress/test_role_grants.py | 7 ++++ 5 files changed, 52 insertions(+), 19 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index fe4cc35029..7083baa092 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1303,6 +1303,7 @@ dependencies = [ "futures", "http 1.1.0", "indexmap 2.0.1", + "itertools 0.10.5", "jsonwebtoken", "metrics", "nix 0.27.1", diff --git a/compute_tools/Cargo.toml b/compute_tools/Cargo.toml index 8ee5dd0665..f9da3ba700 100644 --- a/compute_tools/Cargo.toml +++ b/compute_tools/Cargo.toml @@ -28,6 +28,7 @@ flate2.workspace = true futures.workspace = true http.workspace = true indexmap.workspace = true +itertools.workspace = true jsonwebtoken.workspace = true metrics.workspace = true nix.workspace = true diff --git a/compute_tools/src/compute.rs b/compute_tools/src/compute.rs index 25920675c1..f494e2444a 100644 --- a/compute_tools/src/compute.rs +++ b/compute_tools/src/compute.rs @@ -11,6 +11,7 @@ use compute_api::spec::{ use futures::StreamExt; use futures::future::join_all; use futures::stream::FuturesUnordered; +use itertools::Itertools; use nix::sys::signal::{Signal, kill}; use nix::unistd::Pid; use once_cell::sync::Lazy; @@ -18,7 +19,7 @@ use postgres; use postgres::NoTls; use postgres::error::SqlState; use remote_storage::{DownloadError, RemotePath}; -use std::collections::HashMap; +use std::collections::{HashMap, HashSet}; use std::net::SocketAddr; use std::os::unix::fs::{PermissionsExt, symlink}; use std::path::Path; @@ -1995,23 +1996,40 @@ LIMIT 100", tokio::spawn(conn); // TODO: support other types of grants apart from schemas? - let query = format!( - "GRANT {} ON SCHEMA {} TO {}", - privileges - .iter() - // should not be quoted as it's part of the command. - // is already sanitized so it's ok - .map(|p| p.as_str()) - .collect::>() - .join(", "), - // quote the schema and role name as identifiers to sanitize them. - schema_name.pg_quote(), - role_name.pg_quote(), - ); - db_client - .simple_query(&query) + + // check the role grants first - to gracefully handle read-replicas. + let select = "SELECT privilege_type + FROM pg_namespace + JOIN LATERAL (SELECT * FROM aclexplode(nspacl) AS x) acl ON true + JOIN pg_user users ON acl.grantee = users.usesysid + WHERE users.usename = $1 + AND nspname = $2"; + let rows = db_client + .query(select, &[role_name, schema_name]) .await - .with_context(|| format!("Failed to execute query: {}", query))?; + .with_context(|| format!("Failed to execute query: {select}"))?; + + let already_granted: HashSet = rows.into_iter().map(|row| row.get(0)).collect(); + + let grants = privileges + .iter() + .filter(|p| !already_granted.contains(p.as_str())) + // should not be quoted as it's part of the command. + // is already sanitized so it's ok + .map(|p| p.as_str()) + .join(", "); + + if !grants.is_empty() { + // quote the schema and role name as identifiers to sanitize them. + let schema_name = schema_name.pg_quote(); + let role_name = role_name.pg_quote(); + + let query = format!("GRANT {grants} ON SCHEMA {schema_name} TO {role_name}",); + db_client + .simple_query(&query) + .await + .with_context(|| format!("Failed to execute query: {}", query))?; + } Ok(()) } diff --git a/test_runner/fixtures/neon_fixtures.py b/test_runner/fixtures/neon_fixtures.py index 1b4562c0b3..131820f23e 100644 --- a/test_runner/fixtures/neon_fixtures.py +++ b/test_runner/fixtures/neon_fixtures.py @@ -4613,7 +4613,10 @@ class EndpointFactory: return self def new_replica( - self, origin: Endpoint, endpoint_id: str, config_lines: list[str] | None = None + self, + origin: Endpoint, + endpoint_id: str | None = None, + config_lines: list[str] | None = None, ): branch_name = origin.branch_name assert origin in self.endpoints @@ -4629,7 +4632,10 @@ class EndpointFactory: ) def new_replica_start( - self, origin: Endpoint, endpoint_id: str, config_lines: list[str] | None = None + self, + origin: Endpoint, + endpoint_id: str | None = None, + config_lines: list[str] | None = None, ): branch_name = origin.branch_name assert origin in self.endpoints diff --git a/test_runner/regress/test_role_grants.py b/test_runner/regress/test_role_grants.py index b2251875f0..5b13d461f0 100644 --- a/test_runner/regress/test_role_grants.py +++ b/test_runner/regress/test_role_grants.py @@ -39,3 +39,10 @@ def test_role_grants(neon_simple_env: NeonEnv): res = cur.fetchall() assert res == [(1,)], "select should not succeed" + + # confirm that replicas can also ensure the grants are correctly set. + replica = env.endpoints.new_replica_start(endpoint) + replica_client = replica.http_client() + replica_client.set_role_grants( + "test_role_grants", "test_role", "test_schema", ["CREATE", "USAGE"] + )