From 4e15a68ffd7e2b14bc255956ae5d2130fbfcf2b9 Mon Sep 17 00:00:00 2001 From: Jere Vaara Date: Mon, 14 Oct 2024 20:01:02 +0300 Subject: [PATCH 1/8] Add endpoint to set role grants --- compute_tools/src/compute.rs | 23 +++++++++++++ compute_tools/src/http/api.rs | 35 ++++++++++++++++++-- compute_tools/src/http/openapi_spec.yaml | 42 ++++++++++++++++++------ libs/compute_api/src/requests.rs | 8 +++++ libs/compute_api/src/responses.rs | 5 +++ 5 files changed, 101 insertions(+), 12 deletions(-) diff --git a/compute_tools/src/compute.rs b/compute_tools/src/compute.rs index 285be56264..329cc78573 100644 --- a/compute_tools/src/compute.rs +++ b/compute_tools/src/compute.rs @@ -19,6 +19,7 @@ use futures::future::join_all; use futures::stream::FuturesUnordered; use futures::StreamExt; use nix::unistd::Pid; +use postgres::config::Config; use postgres::error::SqlState; use postgres::{Client, NoTls}; use tracing::{debug, error, info, instrument, warn}; @@ -1367,6 +1368,28 @@ LIMIT 100", download_size } + pub fn set_role_grants( + &self, + db_name: &str, + schema_name: &str, + privilege: &str, + role_name: &str, + ) -> Result<()> { + let mut conf = Config::from_str(self.connstr.as_str()).unwrap(); + conf.dbname(db_name); + + let mut db_client = conf + .connect(NoTls) + .context("Failed to connect to the database")?; + + let query = "GRANT $1 ON SCHEMA $2 TO $3"; + db_client + .execute(query, &[&schema_name, &privilege, &role_name]) + .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..60d542af4c 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, SetRoleGrantsResult, +}; use anyhow::Result; use hyper::header::CONTENT_TYPE; @@ -165,6 +167,35 @@ 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 Response::new(Body::from(msg)); + } + + 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.privilege, + &request.role, + ); + match res { + Ok(_) => Response::new(Body::from("true")), + Err(e) => { + error!("set_role_grants failed: {}", e); + Response::new(Body::from(e.to_string())) + } + } + } + // 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..6b54a9e986 100644 --- a/compute_tools/src/http/openapi_spec.yaml +++ b/compute_tools/src/http/openapi_spec.yaml @@ -10,7 +10,7 @@ paths: /status: get: tags: - - Info + - Info summary: Get compute node internal status. description: "" operationId: getComputeStatus @@ -25,7 +25,7 @@ paths: /metrics.json: get: tags: - - Info + - Info summary: Get compute node startup metrics in JSON format. description: "" operationId: getComputeMetricsJSON @@ -40,7 +40,7 @@ paths: /insights: get: tags: - - Info + - Info summary: Get current compute insights in JSON format. description: | Note, that this doesn't include any historical data. @@ -56,7 +56,7 @@ paths: /installed_extensions: get: tags: - - Info + - Info summary: Get installed extensions. description: "" operationId: getInstalledExtensions @@ -70,7 +70,7 @@ paths: /info: get: tags: - - Info + - Info summary: Get info about the compute pod / VM. description: "" operationId: getInfo @@ -130,7 +130,7 @@ paths: /check_writability: post: tags: - - Check + - Check summary: Check that we can write new data on this compute. description: "" operationId: checkComputeWritability @@ -147,7 +147,7 @@ paths: /configure: post: tags: - - Configure + - Configure summary: Perform compute node configuration. description: | This is a blocking API endpoint, i.e. it blocks waiting until @@ -201,7 +201,7 @@ paths: /extension_server: post: tags: - - Extension + - Extension summary: Download extension from S3 to local folder. description: "" operationId: downloadExtension @@ -230,7 +230,7 @@ paths: /terminate: post: tags: - - Terminate + - Terminate summary: Terminate Postgres and wait for it to exit description: "" operationId: terminate @@ -369,7 +369,7 @@ components: moment, when spec was received. example: "2022-10-12T07:20:50.52Z" status: - $ref: '#/components/schemas/ComputeStatus' + $ref: "#/components/schemas/ComputeStatus" last_active: type: string description: | @@ -427,6 +427,28 @@ components: n_databases: type: integer + SetRoleGrantsRequest: + type: object + required: + - database + - role + - grants + properties: + database: + type: string + description: Database name. + example: "neondb" + role: + type: string + description: Role name. + example: "neon" + grants: + type: array + items: + type: string + description: List of grants to set. + example: ["SELECT", "INSERT"] + # # Errors # diff --git a/libs/compute_api/src/requests.rs b/libs/compute_api/src/requests.rs index 5896c7dc65..963c5992af 100644 --- a/libs/compute_api/src/requests.rs +++ b/libs/compute_api/src/requests.rs @@ -12,3 +12,11 @@ use serde::Deserialize; pub struct ConfigurationRequest { pub spec: ComputeSpec, } + +#[derive(Deserialize, Debug)] +pub struct SetRoleGrantsRequest { + pub database: String, + pub schema: String, + pub privilege: String, + pub role: String, +} diff --git a/libs/compute_api/src/responses.rs b/libs/compute_api/src/responses.rs index 5023fce003..4b4b16ebc4 100644 --- a/libs/compute_api/src/responses.rs +++ b/libs/compute_api/src/responses.rs @@ -168,3 +168,8 @@ pub struct InstalledExtension { pub struct InstalledExtensions { pub extensions: Vec, } + +#[derive(Clone, Debug, Default, Serialize)] +pub struct SetRoleGrantsResult { + pub extension: String, +} From c6c29f86dad6eeaa33265fb8b03770847060e646 Mon Sep 17 00:00:00 2001 From: Jere Vaara Date: Tue, 15 Oct 2024 11:47:55 +0300 Subject: [PATCH 2/8] Handle privileges more strictly --- compute_tools/src/compute.rs | 16 +++++- compute_tools/src/http/api.rs | 31 +++++++++-- compute_tools/src/http/openapi_spec.yaml | 66 ++++++++++++++++++++++-- compute_tools/src/lib.rs | 1 + compute_tools/src/privilege.rs | 58 +++++++++++++++++++++ libs/compute_api/src/requests.rs | 2 +- libs/compute_api/src/responses.rs | 7 ++- 7 files changed, 170 insertions(+), 11 deletions(-) create mode 100644 compute_tools/src/privilege.rs diff --git a/compute_tools/src/compute.rs b/compute_tools/src/compute.rs index 329cc78573..bc8f115973 100644 --- a/compute_tools/src/compute.rs +++ b/compute_tools/src/compute.rs @@ -38,6 +38,7 @@ use crate::checker::create_availability_check_data; use crate::local_proxy; use crate::logger::inlinify; use crate::pg_helpers::*; +use crate::privilege::Privilege; use crate::spec::*; use crate::sync_sk::{check_if_synced, ping_safekeeper}; use crate::{config, extension_server}; @@ -1372,7 +1373,7 @@ LIMIT 100", &self, db_name: &str, schema_name: &str, - privilege: &str, + privileges: &[Privilege], role_name: &str, ) -> Result<()> { let mut conf = Config::from_str(self.connstr.as_str()).unwrap(); @@ -1384,7 +1385,18 @@ LIMIT 100", let query = "GRANT $1 ON SCHEMA $2 TO $3"; db_client - .execute(query, &[&schema_name, &privilege, &role_name]) + .execute( + query, + &[ + &privileges + .iter() + .map(|p| p.as_str()) + .collect::>() + .join(", "), + &schema_name, + &role_name, + ], + ) .context(format!("Failed to execute query: {}", query))?; Ok(()) diff --git a/compute_tools/src/http/api.rs b/compute_tools/src/http/api.rs index 60d542af4c..17b5d16459 100644 --- a/compute_tools/src/http/api.rs +++ b/compute_tools/src/http/api.rs @@ -2,6 +2,7 @@ use std::convert::Infallible; use std::net::IpAddr; use std::net::Ipv6Addr; use std::net::SocketAddr; +use std::str::FromStr; use std::sync::Arc; use std::thread; @@ -9,9 +10,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 crate::privilege::Privilege; use compute_api::requests::{ConfigurationRequest, SetRoleGrantsRequest}; use compute_api::responses::{ - ComputeStatus, ComputeStatusResponse, GenericAPIError, SetRoleGrantsResult, + ComputeStatus, ComputeStatusResponse, GenericAPIError, SetRoleGrantsResponse, }; use anyhow::Result; @@ -181,14 +183,37 @@ async fn routes(req: Request, compute: &Arc) -> Response(&request).unwrap(); + + let privileges: Result, _> = request + .privileges + .iter() + .map(|p| Privilege::from_str(p.as_str())) + .collect(); + let privileges = match privileges { + Ok(privs) => privs, + Err(_) => { + let msg = format!("Invalid privilege in request: {:?}", &request.privileges); + error!(msg); + return Response::new(Body::from(msg)); + } + }; + let res = compute.set_role_grants( &request.database, &request.schema, - &request.privilege, + &privileges, &request.role, ); match res { - Ok(_) => Response::new(Body::from("true")), + 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) => { error!("set_role_grants failed: {}", e); Response::new(Body::from(e.to_string())) diff --git a/compute_tools/src/http/openapi_spec.yaml b/compute_tools/src/http/openapi_spec.yaml index 6b54a9e986..aa85168e44 100644 --- a/compute_tools/src/http/openapi_spec.yaml +++ b/compute_tools/src/http/openapi_spec.yaml @@ -127,6 +127,34 @@ 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: SetRoleGrantsRequest + responses: + 200: + description: Grants applied. + content: + application/json: + schema: + $ref: "#/components/schemas/SetRoleGrantsResponse" + 500: + description: Error occurred during grants application. + content: + application/json: + schema: + $ref: "#/components/schemas/GenericError" + /check_writability: post: tags: @@ -431,23 +459,55 @@ components: type: object required: - database + - schema + - privileges - role - - grants 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" - grants: + + 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 grants to set. + description: List of privileges set. example: ["SELECT", "INSERT"] + role: + type: string + description: Role name. + example: "neon" # # Errors diff --git a/compute_tools/src/lib.rs b/compute_tools/src/lib.rs index d27ae58fa2..8245dc9888 100644 --- a/compute_tools/src/lib.rs +++ b/compute_tools/src/lib.rs @@ -22,6 +22,7 @@ mod migration; pub mod monitor; pub mod params; pub mod pg_helpers; +pub mod privilege; pub mod spec; pub mod swap; pub mod sync_sk; diff --git a/compute_tools/src/privilege.rs b/compute_tools/src/privilege.rs new file mode 100644 index 0000000000..defeae5259 --- /dev/null +++ b/compute_tools/src/privilege.rs @@ -0,0 +1,58 @@ +use std::str::FromStr; + +#[derive(Debug)] +pub enum Privilege { + Select, + Insert, + Update, + Delete, + Truncate, + References, + Trigger, + Usage, + Create, + Connect, + Temporary, + Execute, +} + +impl FromStr for Privilege { + type Err = &'static str; + + fn from_str(s: &str) -> Result { + match s.to_uppercase().as_str() { + "SELECT" => Ok(Privilege::Select), + "INSERT" => Ok(Privilege::Insert), + "UPDATE" => Ok(Privilege::Update), + "DELETE" => Ok(Privilege::Delete), + "TRUNCATE" => Ok(Privilege::Truncate), + "REFERENCES" => Ok(Privilege::References), + "TRIGGER" => Ok(Privilege::Trigger), + "USAGE" => Ok(Privilege::Usage), + "CREATE" => Ok(Privilege::Create), + "CONNECT" => Ok(Privilege::Connect), + "TEMPORARY" => Ok(Privilege::Temporary), + "EXECUTE" => Ok(Privilege::Execute), + _ => Err("Invalid privilege"), + } + } +} + +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 963c5992af..d2e60cc092 100644 --- a/libs/compute_api/src/requests.rs +++ b/libs/compute_api/src/requests.rs @@ -17,6 +17,6 @@ pub struct ConfigurationRequest { pub struct SetRoleGrantsRequest { pub database: String, pub schema: String, - pub privilege: String, + pub privileges: Vec, pub role: String, } diff --git a/libs/compute_api/src/responses.rs b/libs/compute_api/src/responses.rs index 4b4b16ebc4..ef21233f87 100644 --- a/libs/compute_api/src/responses.rs +++ b/libs/compute_api/src/responses.rs @@ -170,6 +170,9 @@ pub struct InstalledExtensions { } #[derive(Clone, Debug, Default, Serialize)] -pub struct SetRoleGrantsResult { - pub extension: String, +pub struct SetRoleGrantsResponse { + pub database: String, + pub schema: String, + pub privileges: Vec, + pub role: String, } From 3f92477af038135750ed3ed81c9f015aa0cadfcf Mon Sep 17 00:00:00 2001 From: Jere Vaara Date: Tue, 15 Oct 2024 16:05:52 +0300 Subject: [PATCH 3/8] Use string formatted query instead --- compute_tools/src/compute.rs | 24 +++++++++++------------- 1 file changed, 11 insertions(+), 13 deletions(-) diff --git a/compute_tools/src/compute.rs b/compute_tools/src/compute.rs index bc8f115973..ec83efbb17 100644 --- a/compute_tools/src/compute.rs +++ b/compute_tools/src/compute.rs @@ -1383,20 +1383,18 @@ LIMIT 100", .connect(NoTls) .context("Failed to connect to the database")?; - let query = "GRANT $1 ON SCHEMA $2 TO $3"; + let query = format!( + "GRANT {} ON SCHEMA {} TO {}", + privileges + .iter() + .map(|p| p.as_str().to_string()) + .collect::>() + .join(", "), + schema_name, + role_name + ); db_client - .execute( - query, - &[ - &privileges - .iter() - .map(|p| p.as_str()) - .collect::>() - .join(", "), - &schema_name, - &role_name, - ], - ) + .simple_query(&query) .context(format!("Failed to execute query: {}", query))?; Ok(()) From e495b99abee45657883e762d3069af80d76faba8 Mon Sep 17 00:00:00 2001 From: Jere Vaara Date: Tue, 15 Oct 2024 16:07:41 +0300 Subject: [PATCH 4/8] Add to_string for privilege --- compute_tools/src/compute.rs | 2 +- compute_tools/src/privilege.rs | 3 +++ 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/compute_tools/src/compute.rs b/compute_tools/src/compute.rs index ec83efbb17..ae2682b682 100644 --- a/compute_tools/src/compute.rs +++ b/compute_tools/src/compute.rs @@ -1387,7 +1387,7 @@ LIMIT 100", "GRANT {} ON SCHEMA {} TO {}", privileges .iter() - .map(|p| p.as_str().to_string()) + .map(|p| p.to_string()) .collect::>() .join(", "), schema_name, diff --git a/compute_tools/src/privilege.rs b/compute_tools/src/privilege.rs index defeae5259..dc4a768651 100644 --- a/compute_tools/src/privilege.rs +++ b/compute_tools/src/privilege.rs @@ -55,4 +55,7 @@ impl Privilege { Privilege::Execute => "EXECUTE", } } + pub fn to_string(&self) -> String { + self.as_str().to_string() + } } From 39e0e3160521ac4b685ad994b500a195ccd23745 Mon Sep 17 00:00:00 2001 From: Jere Vaara Date: Tue, 15 Oct 2024 16:21:52 +0300 Subject: [PATCH 5/8] Move privilege to compute api --- compute_tools/src/compute.rs | 2 +- compute_tools/src/lib.rs | 1 - libs/compute_api/src/lib.rs | 1 + {compute_tools => libs/compute_api}/src/privilege.rs | 0 4 files changed, 2 insertions(+), 2 deletions(-) rename {compute_tools => libs/compute_api}/src/privilege.rs (100%) diff --git a/compute_tools/src/compute.rs b/compute_tools/src/compute.rs index ae2682b682..1f19ed3644 100644 --- a/compute_tools/src/compute.rs +++ b/compute_tools/src/compute.rs @@ -26,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; @@ -38,7 +39,6 @@ use crate::checker::create_availability_check_data; use crate::local_proxy; use crate::logger::inlinify; use crate::pg_helpers::*; -use crate::privilege::Privilege; use crate::spec::*; use crate::sync_sk::{check_if_synced, ping_safekeeper}; use crate::{config, extension_server}; diff --git a/compute_tools/src/lib.rs b/compute_tools/src/lib.rs index 8245dc9888..d27ae58fa2 100644 --- a/compute_tools/src/lib.rs +++ b/compute_tools/src/lib.rs @@ -22,7 +22,6 @@ mod migration; pub mod monitor; pub mod params; pub mod pg_helpers; -pub mod privilege; pub mod spec; pub mod swap; pub mod sync_sk; 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/compute_tools/src/privilege.rs b/libs/compute_api/src/privilege.rs similarity index 100% rename from compute_tools/src/privilege.rs rename to libs/compute_api/src/privilege.rs From d0ca79aeb3b4cfa6edf96fa34f4a2f878fa35d36 Mon Sep 17 00:00:00 2001 From: Jere Vaara Date: Tue, 15 Oct 2024 17:07:23 +0300 Subject: [PATCH 6/8] Use serde to serialize/deserialize Privilege instead of manual --- compute_tools/src/compute.rs | 2 +- compute_tools/src/http/api.rs | 18 +----------- libs/compute_api/src/privilege.rs | 49 ++----------------------------- libs/compute_api/src/requests.rs | 4 +-- libs/compute_api/src/responses.rs | 7 +++-- 5 files changed, 11 insertions(+), 69 deletions(-) diff --git a/compute_tools/src/compute.rs b/compute_tools/src/compute.rs index 1f19ed3644..f5e65a2cec 100644 --- a/compute_tools/src/compute.rs +++ b/compute_tools/src/compute.rs @@ -1387,7 +1387,7 @@ LIMIT 100", "GRANT {} ON SCHEMA {} TO {}", privileges .iter() - .map(|p| p.to_string()) + .map(|p| serde_json::to_string(p).unwrap()) .collect::>() .join(", "), schema_name, diff --git a/compute_tools/src/http/api.rs b/compute_tools/src/http/api.rs index 17b5d16459..2f7b3feef1 100644 --- a/compute_tools/src/http/api.rs +++ b/compute_tools/src/http/api.rs @@ -2,7 +2,6 @@ use std::convert::Infallible; use std::net::IpAddr; use std::net::Ipv6Addr; use std::net::SocketAddr; -use std::str::FromStr; use std::sync::Arc; use std::thread; @@ -10,7 +9,6 @@ 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 crate::privilege::Privilege; use compute_api::requests::{ConfigurationRequest, SetRoleGrantsRequest}; use compute_api::responses::{ ComputeStatus, ComputeStatusResponse, GenericAPIError, SetRoleGrantsResponse, @@ -184,24 +182,10 @@ async fn routes(req: Request, compute: &Arc) -> Response(&request).unwrap(); - let privileges: Result, _> = request - .privileges - .iter() - .map(|p| Privilege::from_str(p.as_str())) - .collect(); - let privileges = match privileges { - Ok(privs) => privs, - Err(_) => { - let msg = format!("Invalid privilege in request: {:?}", &request.privileges); - error!(msg); - return Response::new(Body::from(msg)); - } - }; - let res = compute.set_role_grants( &request.database, &request.schema, - &privileges, + &request.privileges, &request.role, ); match res { diff --git a/libs/compute_api/src/privilege.rs b/libs/compute_api/src/privilege.rs index dc4a768651..bdb8cfa753 100644 --- a/libs/compute_api/src/privilege.rs +++ b/libs/compute_api/src/privilege.rs @@ -1,6 +1,5 @@ -use std::str::FromStr; - -#[derive(Debug)] +#[derive(Debug, Clone, serde::Deserialize, serde::Serialize)] +#[serde(rename_all = "UPPERCASE")] pub enum Privilege { Select, Insert, @@ -15,47 +14,3 @@ pub enum Privilege { Temporary, Execute, } - -impl FromStr for Privilege { - type Err = &'static str; - - fn from_str(s: &str) -> Result { - match s.to_uppercase().as_str() { - "SELECT" => Ok(Privilege::Select), - "INSERT" => Ok(Privilege::Insert), - "UPDATE" => Ok(Privilege::Update), - "DELETE" => Ok(Privilege::Delete), - "TRUNCATE" => Ok(Privilege::Truncate), - "REFERENCES" => Ok(Privilege::References), - "TRIGGER" => Ok(Privilege::Trigger), - "USAGE" => Ok(Privilege::Usage), - "CREATE" => Ok(Privilege::Create), - "CONNECT" => Ok(Privilege::Connect), - "TEMPORARY" => Ok(Privilege::Temporary), - "EXECUTE" => Ok(Privilege::Execute), - _ => Err("Invalid privilege"), - } - } -} - -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", - } - } - pub fn to_string(&self) -> String { - self.as_str().to_string() - } -} diff --git a/libs/compute_api/src/requests.rs b/libs/compute_api/src/requests.rs index d2e60cc092..5e14ccfeeb 100644 --- a/libs/compute_api/src/requests.rs +++ b/libs/compute_api/src/requests.rs @@ -1,6 +1,6 @@ //! Structs representing the JSON formats used in the compute_ctl's HTTP API. -use crate::spec::ComputeSpec; +use crate::{privilege::Privilege, spec::ComputeSpec}; use serde::Deserialize; /// Request of the /configure API @@ -17,6 +17,6 @@ pub struct ConfigurationRequest { pub struct SetRoleGrantsRequest { pub database: String, pub schema: String, - pub privileges: Vec, + pub privileges: Vec, pub role: String, } diff --git a/libs/compute_api/src/responses.rs b/libs/compute_api/src/responses.rs index ef21233f87..18ffcd5071 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, Role}, +}; #[derive(Serialize, Debug, Deserialize)] pub struct GenericAPIError { @@ -173,6 +176,6 @@ pub struct InstalledExtensions { pub struct SetRoleGrantsResponse { pub database: String, pub schema: String, - pub privileges: Vec, + pub privileges: Vec, pub role: String, } From b57f2c60f297e724d86fb249cfffb267a7f15869 Mon Sep 17 00:00:00 2001 From: Jere Vaara Date: Wed, 16 Oct 2024 13:25:19 +0300 Subject: [PATCH 7/8] Make set role grants fn async --- compute_tools/src/compute.rs | 17 ++++++++++------- compute_tools/src/http/api.rs | 16 +++++++++------- 2 files changed, 19 insertions(+), 14 deletions(-) diff --git a/compute_tools/src/compute.rs b/compute_tools/src/compute.rs index f5e65a2cec..b3902a6c4a 100644 --- a/compute_tools/src/compute.rs +++ b/compute_tools/src/compute.rs @@ -1369,32 +1369,35 @@ LIMIT 100", download_size } - pub fn set_role_grants( + pub async fn set_role_grants( &self, db_name: &str, schema_name: &str, privileges: &[Privilege], role_name: &str, ) -> 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 mut db_client = conf + let (db_client, conn) = conf .connect(NoTls) + .await .context("Failed to connect to the database")?; + tokio::spawn(conn); + // todo: pg_quote let query = format!( "GRANT {} ON SCHEMA {} TO {}", - privileges - .iter() - .map(|p| serde_json::to_string(p).unwrap()) - .collect::>() - .join(", "), + privileges.iter().map(|p| p.as_str()).join(", "), schema_name, role_name ); db_client .simple_query(&query) + .await .context(format!("Failed to execute query: {}", query))?; Ok(()) diff --git a/compute_tools/src/http/api.rs b/compute_tools/src/http/api.rs index 2f7b3feef1..03a3633342 100644 --- a/compute_tools/src/http/api.rs +++ b/compute_tools/src/http/api.rs @@ -182,14 +182,16 @@ async fn routes(req: Request, compute: &Arc) -> Response(&request).unwrap(); - let res = compute.set_role_grants( - &request.database, - &request.schema, - &request.privileges, - &request.role, - ); + let res = compute + .set_role_grants( + &request.database, + &request.schema, + &request.privileges, + &request.role, + ) + .await; match res { - Ok(_) => render_json(Body::from( + Ok(()) => render_json(Body::from( serde_json::to_string(&SetRoleGrantsResponse { database: request.database, schema: request.schema, From 6de90c7ce42bce9c71bd645523e0d97ac7bb19f9 Mon Sep 17 00:00:00 2001 From: Conrad Ludgate Date: Wed, 16 Oct 2024 12:09:25 +0100 Subject: [PATCH 8/8] quote identifiers --- compute_tools/src/compute.rs | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/compute_tools/src/compute.rs b/compute_tools/src/compute.rs index b3902a6c4a..771e0db118 100644 --- a/compute_tools/src/compute.rs +++ b/compute_tools/src/compute.rs @@ -19,7 +19,6 @@ use futures::future::join_all; use futures::stream::FuturesUnordered; use futures::StreamExt; use nix::unistd::Pid; -use postgres::config::Config; use postgres::error::SqlState; use postgres::{Client, NoTls}; use tracing::{debug, error, info, instrument, warn}; @@ -1388,12 +1387,18 @@ LIMIT 100", .context("Failed to connect to the database")?; tokio::spawn(conn); - // todo: pg_quote let query = format!( "GRANT {} ON SCHEMA {} TO {}", - privileges.iter().map(|p| p.as_str()).join(", "), - schema_name, - role_name + privileges + .iter() + // should not be quoted as it's part of the command. + // is already sanitized so it's ok + .map(|p| serde_json::to_string(p).unwrap()) + .collect::>() + .join(", "), + // quote the schema and role name as identifiers to sanitize them. + schema_name.to_string().pg_quote(), + role_name.to_string().pg_quote(), ); db_client .simple_query(&query)