From 24654b8eee8706e8ae98948733a28b56df83536b Mon Sep 17 00:00:00 2001 From: Jere Vaara Date: Fri, 18 Oct 2024 13:25:45 +0300 Subject: [PATCH] compute_ctl: Add endpoint that allows setting role grants (#9395) This PR introduces a `/grants` endpoint which allows setting specific `privileges` to certain `role` for a certain `schema`. Related to #9344 Together these endpoints will be used to configure JWT extension and set correct usage to its schema to specific roles that will need them. --------- Co-authored-by: Conrad Ludgate --- compute_tools/src/compute.rs | 43 ++++++++++++ compute_tools/src/http/api.rs | 48 ++++++++++++- compute_tools/src/http/openapi_spec.yaml | 89 ++++++++++++++++++++++++ libs/compute_api/src/lib.rs | 1 + libs/compute_api/src/privilege.rs | 35 ++++++++++ libs/compute_api/src/requests.rs | 13 +++- libs/compute_api/src/responses.rs | 13 +++- test_runner/fixtures/endpoint/http.py | 8 +++ test_runner/regress/test_role_grants.py | 41 +++++++++++ 9 files changed, 287 insertions(+), 4 deletions(-) create mode 100644 libs/compute_api/src/privilege.rs create mode 100644 test_runner/regress/test_role_grants.py diff --git a/compute_tools/src/compute.rs b/compute_tools/src/compute.rs index 6aec008f3a..11fee73f03 100644 --- a/compute_tools/src/compute.rs +++ b/compute_tools/src/compute.rs @@ -15,6 +15,7 @@ use std::time::Instant; use anyhow::{Context, Result}; use chrono::{DateTime, Utc}; +use compute_api::spec::PgIdent; use futures::future::join_all; use futures::stream::FuturesUnordered; use futures::StreamExt; @@ -25,6 +26,7 @@ use tracing::{debug, error, info, instrument, warn}; use utils::id::{TenantId, TimelineId}; use utils::lsn::Lsn; +use compute_api::privilege::Privilege; use compute_api::responses::{ComputeMetrics, ComputeStatus}; use compute_api::spec::{ComputeFeature, ComputeMode, ComputeSpec}; use utils::measured_stream::MeasuredReader; @@ -1373,6 +1375,47 @@ LIMIT 100", download_size } + pub async fn set_role_grants( + &self, + db_name: &PgIdent, + schema_name: &PgIdent, + privileges: &[Privilege], + role_name: &PgIdent, + ) -> Result<()> { + use tokio_postgres::config::Config; + use tokio_postgres::NoTls; + + let mut conf = Config::from_str(self.connstr.as_str()).unwrap(); + conf.dbname(db_name); + + let (db_client, conn) = conf + .connect(NoTls) + .await + .context("Failed to connect to the database")?; + 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) + .await + .with_context(|| format!("Failed to execute query: {}", query))?; + + Ok(()) + } + #[tokio::main] pub async fn prepare_preload_libraries( &self, diff --git a/compute_tools/src/http/api.rs b/compute_tools/src/http/api.rs index 79e6158081..133ab9f5af 100644 --- a/compute_tools/src/http/api.rs +++ b/compute_tools/src/http/api.rs @@ -9,8 +9,10 @@ use crate::catalog::SchemaDumpError; use crate::catalog::{get_database_schema, get_dbs_and_roles}; use crate::compute::forward_termination_signal; use crate::compute::{ComputeNode, ComputeState, ParsedSpec}; -use compute_api::requests::ConfigurationRequest; -use compute_api::responses::{ComputeStatus, ComputeStatusResponse, GenericAPIError}; +use compute_api::requests::{ConfigurationRequest, SetRoleGrantsRequest}; +use compute_api::responses::{ + ComputeStatus, ComputeStatusResponse, GenericAPIError, SetRoleGrantsResponse, +}; use anyhow::Result; use hyper::header::CONTENT_TYPE; @@ -165,6 +167,48 @@ async fn routes(req: Request, compute: &Arc) -> Response { + info!("serving /grants POST request"); + let status = compute.get_status(); + if status != ComputeStatus::Running { + let msg = format!( + "invalid compute status for set_role_grants request: {:?}", + status + ); + error!(msg); + return render_json_error(&msg, StatusCode::PRECONDITION_FAILED); + } + + let request = hyper::body::to_bytes(req.into_body()).await.unwrap(); + let request = serde_json::from_slice::(&request).unwrap(); + + let res = compute + .set_role_grants( + &request.database, + &request.schema, + &request.privileges, + &request.role, + ) + .await; + match res { + Ok(()) => render_json(Body::from( + serde_json::to_string(&SetRoleGrantsResponse { + database: request.database, + schema: request.schema, + role: request.role, + privileges: request.privileges, + }) + .unwrap(), + )), + Err(e) => render_json_error( + &format!("could not grant role privileges to the schema: {e}"), + // TODO: can we filter on role/schema not found errors + // and return appropriate error code? + StatusCode::INTERNAL_SERVER_ERROR, + ), + } + } + // get the list of installed extensions // currently only used in python tests // TODO: call it from cplane diff --git a/compute_tools/src/http/openapi_spec.yaml b/compute_tools/src/http/openapi_spec.yaml index e9fa66b323..73dbdc3ee9 100644 --- a/compute_tools/src/http/openapi_spec.yaml +++ b/compute_tools/src/http/openapi_spec.yaml @@ -127,6 +127,41 @@ paths: schema: $ref: "#/components/schemas/GenericError" + /grants: + post: + tags: + - Grants + summary: Apply grants to the database. + description: "" + operationId: setRoleGrants + requestBody: + description: Grants request. + required: true + content: + application/json: + schema: + $ref: "#/components/schemas/SetRoleGrantsRequest" + responses: + 200: + description: Grants applied. + content: + application/json: + schema: + $ref: "#/components/schemas/SetRoleGrantsResponse" + 412: + description: | + Compute is not in the right state for processing the request. + content: + application/json: + schema: + $ref: "#/components/schemas/GenericError" + 500: + description: Error occurred during grants application. + content: + application/json: + schema: + $ref: "#/components/schemas/GenericError" + /check_writability: post: tags: @@ -427,6 +462,60 @@ components: n_databases: type: integer + SetRoleGrantsRequest: + type: object + required: + - database + - schema + - privileges + - role + properties: + database: + type: string + description: Database name. + example: "neondb" + schema: + type: string + description: Schema name. + example: "public" + privileges: + type: array + items: + type: string + description: List of privileges to set. + example: ["SELECT", "INSERT"] + role: + type: string + description: Role name. + example: "neon" + + SetRoleGrantsResponse: + type: object + required: + - database + - schema + - privileges + - role + properties: + database: + type: string + description: Database name. + example: "neondb" + schema: + type: string + description: Schema name. + example: "public" + privileges: + type: array + items: + type: string + description: List of privileges set. + example: ["SELECT", "INSERT"] + role: + type: string + description: Role name. + example: "neon" + # # Errors # diff --git a/libs/compute_api/src/lib.rs b/libs/compute_api/src/lib.rs index 210a52d089..f4f3d92fc6 100644 --- a/libs/compute_api/src/lib.rs +++ b/libs/compute_api/src/lib.rs @@ -1,5 +1,6 @@ #![deny(unsafe_code)] #![deny(clippy::undocumented_unsafe_blocks)] +pub mod privilege; pub mod requests; pub mod responses; pub mod spec; diff --git a/libs/compute_api/src/privilege.rs b/libs/compute_api/src/privilege.rs new file mode 100644 index 0000000000..dc0d870946 --- /dev/null +++ b/libs/compute_api/src/privilege.rs @@ -0,0 +1,35 @@ +#[derive(Debug, Clone, serde::Deserialize, serde::Serialize)] +#[serde(rename_all = "UPPERCASE")] +pub enum Privilege { + Select, + Insert, + Update, + Delete, + Truncate, + References, + Trigger, + Usage, + Create, + Connect, + Temporary, + Execute, +} + +impl Privilege { + pub fn as_str(&self) -> &'static str { + match self { + Privilege::Select => "SELECT", + Privilege::Insert => "INSERT", + Privilege::Update => "UPDATE", + Privilege::Delete => "DELETE", + Privilege::Truncate => "TRUNCATE", + Privilege::References => "REFERENCES", + Privilege::Trigger => "TRIGGER", + Privilege::Usage => "USAGE", + Privilege::Create => "CREATE", + Privilege::Connect => "CONNECT", + Privilege::Temporary => "TEMPORARY", + Privilege::Execute => "EXECUTE", + } + } +} diff --git a/libs/compute_api/src/requests.rs b/libs/compute_api/src/requests.rs index 5896c7dc65..fbc7577dd9 100644 --- a/libs/compute_api/src/requests.rs +++ b/libs/compute_api/src/requests.rs @@ -1,6 +1,9 @@ //! Structs representing the JSON formats used in the compute_ctl's HTTP API. -use crate::spec::ComputeSpec; +use crate::{ + privilege::Privilege, + spec::{ComputeSpec, PgIdent}, +}; use serde::Deserialize; /// Request of the /configure API @@ -12,3 +15,11 @@ use serde::Deserialize; pub struct ConfigurationRequest { pub spec: ComputeSpec, } + +#[derive(Deserialize, Debug)] +pub struct SetRoleGrantsRequest { + pub database: PgIdent, + pub schema: PgIdent, + pub privileges: Vec, + pub role: PgIdent, +} diff --git a/libs/compute_api/src/responses.rs b/libs/compute_api/src/responses.rs index 5023fce003..fadf524273 100644 --- a/libs/compute_api/src/responses.rs +++ b/libs/compute_api/src/responses.rs @@ -6,7 +6,10 @@ use std::fmt::Display; use chrono::{DateTime, Utc}; use serde::{Deserialize, Serialize, Serializer}; -use crate::spec::{ComputeSpec, Database, Role}; +use crate::{ + privilege::Privilege, + spec::{ComputeSpec, Database, PgIdent, Role}, +}; #[derive(Serialize, Debug, Deserialize)] pub struct GenericAPIError { @@ -168,3 +171,11 @@ pub struct InstalledExtension { pub struct InstalledExtensions { pub extensions: Vec, } + +#[derive(Clone, Debug, Default, Serialize)] +pub struct SetRoleGrantsResponse { + pub database: PgIdent, + pub schema: PgIdent, + pub privileges: Vec, + pub role: PgIdent, +} diff --git a/test_runner/fixtures/endpoint/http.py b/test_runner/fixtures/endpoint/http.py index 26895df8a6..e7b014b4a9 100644 --- a/test_runner/fixtures/endpoint/http.py +++ b/test_runner/fixtures/endpoint/http.py @@ -28,3 +28,11 @@ class EndpointHttpClient(requests.Session): res = self.get(f"http://localhost:{self.port}/installed_extensions") res.raise_for_status() return res.json() + + def set_role_grants(self, database: str, role: str, schema: str, privileges: list[str]): + res = self.post( + f"http://localhost:{self.port}/grants", + json={"database": database, "schema": schema, "role": role, "privileges": privileges}, + ) + res.raise_for_status() + return res.json() diff --git a/test_runner/regress/test_role_grants.py b/test_runner/regress/test_role_grants.py new file mode 100644 index 0000000000..b2251875f0 --- /dev/null +++ b/test_runner/regress/test_role_grants.py @@ -0,0 +1,41 @@ +import psycopg2 +from fixtures.neon_fixtures import NeonEnv + + +def test_role_grants(neon_simple_env: NeonEnv): + """basic test for the endpoint that grants permissions for a role against a schema""" + + env = neon_simple_env + + env.create_branch("test_role_grants") + + endpoint = env.endpoints.create_start("test_role_grants") + + endpoint.safe_psql("CREATE DATABASE test_role_grants") + endpoint.safe_psql("CREATE SCHEMA IF NOT EXISTS test_schema", dbname="test_role_grants") + endpoint.safe_psql("CREATE ROLE test_role WITH LOGIN", dbname="test_role_grants") + + # confirm we do not yet have access + pg_conn = endpoint.connect(dbname="test_role_grants", user="test_role") + with pg_conn.cursor() as cur: + try: + cur.execute('CREATE TABLE "test_schema"."test_table" (id integer primary key)') + raise ValueError("create table should not succeed") + except psycopg2.errors.InsufficientPrivilege: + pass + except BaseException as e: + raise e + + client = endpoint.http_client() + res = client.set_role_grants( + "test_role_grants", "test_role", "test_schema", ["CREATE", "USAGE"] + ) + + # confirm we have access + with pg_conn.cursor() as cur: + cur.execute('CREATE TABLE "test_schema"."test_table" (id integer primary key)') + cur.execute('INSERT INTO "test_schema"."test_table" (id) VALUES (1)') + cur.execute('SELECT id from "test_schema"."test_table"') + res = cur.fetchall() + + assert res == [(1,)], "select should not succeed"