Merge remote-tracking branch 'origin/grants-endpoint' into test-local-proxy-jwt-ext-install

2026-07-07 06:00:38 +00:00 · 2024-10-17 10:15:37 +01:00
parent 3f3cb744f0 6de90c7ce4
commit 8a7ed8ce72
7 changed files with 205 additions and 3 deletions
--- a/compute_tools/src/compute.rs
+++ b/compute_tools/src/compute.rs
@@ -25,6 +25,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;
@@ -1417,6 +1418,46 @@ LIMIT 100",
        Ok(ext_version.to_string())
    }

+    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 (db_client, conn) = conf
+            .connect(NoTls)
+            .await
+            .context("Failed to connect to the database")?;
+        tokio::spawn(conn);
+
+        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| serde_json::to_string(p).unwrap())
+                .collect::<Vec<String>>()
+                .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)
+            .await
+            .context(format!("Failed to execute query: {}", query))?;
+
+        Ok(())
+    }
+
    #[tokio::main]
    pub async fn prepare_preload_libraries(
        &self,
--- a/compute_tools/src/http/api.rs
+++ b/compute_tools/src/http/api.rs
@@ -9,9 +9,12 @@ 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, ExtensionInstallRequest};
+use compute_api::requests::{
+    ExtensionInstallRequest, {ConfigurationRequest, SetRoleGrantsRequest},
+};
 use compute_api::responses::{
    ComputeStatus, ComputeStatusResponse, ExtensionInstallResult, GenericAPIError,
+    SetRoleGrantsResponse,
 };

 use anyhow::Result;
@@ -199,6 +202,46 @@ async fn routes(req: Request<Body>, compute: &Arc<ComputeNode>) -> Response<Body
            }
        }

+        (&Method::POST, "/grants") => {
+            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::<SetRoleGrantsRequest>(&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) => {
+                    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
--- 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:
@@ -487,6 +515,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
    #
--- 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;
--- a/libs/compute_api/src/privilege.rs
+++ b/libs/compute_api/src/privilege.rs
@@ -0,0 +1,16 @@
+#[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,
+}
--- 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
@@ -19,3 +19,11 @@ pub struct ExtensionInstallRequest {
    pub database: String,
    pub version: String,
 }
+
+#[derive(Deserialize, Debug)]
+pub struct SetRoleGrantsRequest {
+    pub database: String,
+    pub schema: String,
+    pub privileges: Vec<Privilege>,
+    pub role: String,
+}
--- 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 {
@@ -174,3 +177,11 @@ pub struct ExtensionInstallResult {
    pub extension: String,
    pub version: String,
 }
+
+#[derive(Clone, Debug, Default, Serialize)]
+pub struct SetRoleGrantsResponse {
+    pub database: String,
+    pub schema: String,
+    pub privileges: Vec<Privilege>,
+    pub role: String,
+}