Add dataset_groups and dataset_permissions tables with organization_id references

- Created dataset_groups and dataset_permissions tables in the database schema, including organization_id as a foreign key with ON DELETE CASCADE.
- Added corresponding indexes for organization_id in both tables to optimize query performance.
- Updated the Rust models and schema to reflect the new tables and their relationships.
- Integrated dataset_groups into the API routes for improved data organization and management.

These changes enhance the database structure and facilitate better handling of dataset-related permissions and groupings.

api/migrations/2025-01-08-163102_add_dataset_groups_adjust_teams_adjust_permissions/up.sql

@@ -39,6 +39,7 @@ ALTER TABLE users_to_organizations
 -- Create dataset_groups table
 CREATE TABLE dataset_groups (
     id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
+    organization_id UUID NOT NULL REFERENCES organizations(id) ON DELETE CASCADE,
     name VARCHAR NOT NULL,
     created_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT NOW(),
     updated_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT NOW(),
@@ -48,6 +49,9 @@ CREATE TABLE dataset_groups (
 -- Add index on deleted_at for soft delete queries
 CREATE INDEX dataset_groups_deleted_at_idx ON dataset_groups(deleted_at);
 
+-- Add indexes
+CREATE INDEX dataset_groups_organization_id_idx ON dataset_groups(organization_id);
+
 -- Create datasets_to_dataset_groups join table
 CREATE TABLE datasets_to_dataset_groups (
     dataset_id UUID NOT NULL REFERENCES datasets(id) ON DELETE CASCADE,
@@ -73,6 +77,7 @@ CREATE INDEX permission_groups_to_users_user_id_idx ON permission_groups_to_user
 -- Create dataset_permissions table
 CREATE TABLE dataset_permissions (
     id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
+    organization_id UUID NOT NULL REFERENCES organizations(id) ON DELETE CASCADE,
     dataset_id UUID NOT NULL REFERENCES datasets(id) ON DELETE CASCADE,
     permission_id UUID NOT NULL,
     permission_type VARCHAR NOT NULL CHECK (
@@ -89,6 +94,9 @@ CREATE INDEX dataset_permissions_deleted_at_idx ON dataset_permissions(deleted_a
 CREATE INDEX dataset_permissions_permission_lookup_idx ON dataset_permissions(permission_id, permission_type);
 CREATE INDEX dataset_permissions_dataset_id_idx ON dataset_permissions(dataset_id);
 
+-- Add indexes
+CREATE INDEX dataset_permissions_organization_id_idx ON dataset_permissions(organization_id);
+
 -- Drop default before type change
 ALTER TABLE teams_to_users ALTER COLUMN role DROP DEFAULT;
 

api/src/database/models.rs

@@ -490,3 +490,14 @@ pub struct EntityRelationship {
     pub relationship_type: String,
     pub created_at: DateTime<Utc>,
 }
+
+#[derive(Queryable, Insertable, Debug)]
+#[diesel(table_name = dataset_groups)]
+pub struct DatasetGroup {
+    pub id: Uuid,
+    pub organization_id: Uuid,
+    pub name: String,
+    pub created_at: DateTime<Utc>,
+    pub updated_at: DateTime<Utc>,
+    pub deleted_at: Option<DateTime<Utc>>,
+}

api/src/database/schema.rs

@@ -193,6 +193,7 @@ diesel::table! {
 diesel::table! {
     dataset_groups (id) {
         id -> Uuid,
+        organization_id -> Uuid,
         name -> Varchar,
         created_at -> Timestamptz,
         updated_at -> Timestamptz,
@@ -203,6 +204,7 @@ diesel::table! {
 diesel::table! {
     dataset_permissions (id) {
         id -> Uuid,
+        organization_id -> Uuid,
         dataset_id -> Uuid,
         permission_id -> Uuid,
         permission_type -> Varchar,
@@ -501,7 +503,9 @@ diesel::joinable!(collections -> organizations (organization_id));
 diesel::joinable!(dashboard_versions -> dashboards (dashboard_id));
 diesel::joinable!(dashboards -> organizations (organization_id));
 diesel::joinable!(data_sources -> organizations (organization_id));
+diesel::joinable!(dataset_groups -> organizations (organization_id));
 diesel::joinable!(dataset_permissions -> datasets (dataset_id));
+diesel::joinable!(dataset_permissions -> organizations (organization_id));
 diesel::joinable!(datasets -> data_sources (data_source_id));
 diesel::joinable!(datasets -> organizations (organization_id));
 diesel::joinable!(datasets_to_dataset_groups -> dataset_groups (dataset_group_id));

api/src/routes/rest/routes/dataset_groups/delete_dataset_group.rs

@@ -0,0 +1,60 @@
+use anyhow::Result;
+use axum::{extract::Path, http::StatusCode, Extension};
+use chrono::Utc;
+use diesel::prelude::*;
+use diesel_async::RunQueryDsl;
+use uuid::Uuid;
+
+use crate::database::lib::get_pg_pool;
+use crate::database::models::User;
+use crate::database::schema::dataset_groups;
+use crate::routes::rest::ApiResponse;
+use crate::utils::security::checks::is_user_workspace_admin_or_data_admin;
+
+pub async fn delete_dataset_group(
+    Extension(user): Extension<User>,
+    Path(dataset_group_id): Path<Uuid>,
+) -> Result<ApiResponse<()>, (StatusCode, &'static str)> {
+    // Check if user is workspace admin or data admin
+    match is_user_workspace_admin_or_data_admin(&user.id).await {
+        Ok(true) => (),
+        Ok(false) => return Err((StatusCode::FORBIDDEN, "Insufficient permissions")),
+        Err(e) => {
+            tracing::error!("Error checking user permissions: {:?}", e);
+            return Err((
+                StatusCode::INTERNAL_SERVER_ERROR,
+                "Error checking user permissions",
+            ));
+        }
+    }
+
+    match delete_dataset_group_handler(dataset_group_id).await {
+        Ok(_) => Ok(ApiResponse::NoContent),
+        Err(e) => {
+            tracing::error!("Error deleting dataset group: {:?}", e);
+            Err((
+                StatusCode::INTERNAL_SERVER_ERROR,
+                "Error deleting dataset group",
+            ))
+        }
+    }
+}
+
+async fn delete_dataset_group_handler(dataset_group_id: Uuid) -> Result<()> {
+    let mut conn = get_pg_pool().get().await?;
+
+    let rows_affected = diesel::update(
+        dataset_groups::table
+            .filter(dataset_groups::id.eq(dataset_group_id))
+            .filter(dataset_groups::deleted_at.is_null()),
+    )
+    .set(dataset_groups::deleted_at.eq(Some(Utc::now())))
+    .execute(&mut *conn)
+    .await?;
+
+    if rows_affected == 0 {
+        return Err(anyhow::anyhow!("Dataset group not found"));
+    }
+
+    Ok(())
+} 

api/src/routes/rest/routes/dataset_groups/get_dataset_group.rs

@@ -0,0 +1,48 @@
+use anyhow::Result;
+use axum::{
+    extract::Path,
+    http::StatusCode,
+    Extension,
+};
+use diesel::prelude::*;
+use diesel_async::RunQueryDsl;
+use uuid::Uuid;
+
+use crate::database::lib::get_pg_pool;
+use crate::database::models::{DatasetGroup, User};
+use crate::database::schema::dataset_groups;
+use crate::routes::rest::ApiResponse;
+use super::list_dataset_groups::DatasetGroupInfo;
+
+pub async fn get_dataset_group(
+    Extension(user): Extension<User>,
+    Path(dataset_group_id): Path<Uuid>,
+) -> Result<ApiResponse<DatasetGroupInfo>, (StatusCode, &'static str)> {
+    let dataset_group = match get_dataset_group_handler(dataset_group_id).await {
+        Ok(group) => group,
+        Err(e) => {
+            tracing::error!("Error getting dataset group: {:?}", e);
+            return Err((StatusCode::INTERNAL_SERVER_ERROR, "Error getting dataset group"));
+        }
+    };
+
+    Ok(ApiResponse::JsonData(dataset_group))
+}
+
+async fn get_dataset_group_handler(dataset_group_id: Uuid) -> Result<DatasetGroupInfo> {
+    let mut conn = get_pg_pool().get().await?;
+
+    let dataset_group = dataset_groups::table
+        .filter(dataset_groups::id.eq(dataset_group_id))
+        .filter(dataset_groups::deleted_at.is_null())
+        .first::<DatasetGroup>(&mut *conn)
+        .await
+        .map_err(|_| anyhow::anyhow!("Dataset group not found"))?;
+
+    Ok(DatasetGroupInfo {
+        id: dataset_group.id,
+        name: dataset_group.name,
+        created_at: dataset_group.created_at,
+        updated_at: dataset_group.updated_at,
+    })
+} 

api/src/routes/rest/routes/dataset_groups/list_dataset_groups.rs

@@ -0,0 +1,61 @@
+use anyhow::Result;
+use axum::http::StatusCode;
+use axum::Extension;
+use chrono::{DateTime, Utc};
+use diesel::prelude::*;
+use diesel_async::RunQueryDsl;
+use serde::Serialize;
+use uuid::Uuid;
+
+use crate::database::lib::get_pg_pool;
+use crate::database::models::User;
+use crate::database::schema::dataset_groups;
+use crate::routes::rest::ApiResponse;
+use crate::utils::user::user_info::get_user_organization_id;
+
+#[derive(Debug, Serialize)]
+pub struct DatasetGroupInfo {
+    pub id: Uuid,
+    pub name: String,
+    pub created_at: DateTime<Utc>,
+    pub updated_at: DateTime<Utc>,
+}
+
+pub async fn list_dataset_groups(
+    Extension(user): Extension<User>,
+) -> Result<ApiResponse<Vec<DatasetGroupInfo>>, (StatusCode, &'static str)> {
+    let dataset_groups = match list_dataset_groups_handler(user).await {
+        Ok(groups) => groups,
+        Err(e) => {
+            tracing::error!("Error listing dataset groups: {:?}", e);
+            return Err((
+                StatusCode::INTERNAL_SERVER_ERROR,
+                "Error listing dataset groups",
+            ));
+        }
+    };
+
+    Ok(ApiResponse::JsonData(dataset_groups))
+}
+
+async fn list_dataset_groups_handler(user: User) -> Result<Vec<DatasetGroupInfo>> {
+    let mut conn = get_pg_pool().get().await?;
+    let organization_id = get_user_organization_id(&user.id).await?;
+
+    let dataset_groups = dataset_groups::table
+        .filter(dataset_groups::deleted_at.is_null())
+        .filter(dataset_groups::organization_id.eq(organization_id))
+        .order_by(dataset_groups::created_at.desc())
+        .load::<crate::database::models::DatasetGroup>(&mut *conn)
+        .await?;
+
+    Ok(dataset_groups
+        .into_iter()
+        .map(|group| DatasetGroupInfo {
+            id: group.id,
+            name: group.name.to_string(),
+            created_at: group.created_at,
+            updated_at: group.updated_at,
+        })
+        .collect())
+}

api/src/routes/rest/routes/dataset_groups/mod.rs

@@ -0,0 +1,31 @@
+pub mod delete_dataset_group;
+pub mod get_dataset_group;
+pub mod list_dataset_groups;
+pub mod post_dataset_group;
+pub mod put_dataset_group;
+
+use axum::{
+    middleware,
+    routing::{delete, get, post, put},
+    Router,
+};
+
+use crate::buster_middleware::auth::auth;
+
+use self::{
+    delete_dataset_group::delete_dataset_group,
+    get_dataset_group::get_dataset_group,
+    list_dataset_groups::list_dataset_groups,
+    post_dataset_group::post_dataset_group,
+    put_dataset_group::put_dataset_group,
+};
+
+pub fn router() -> Router {
+    Router::new()
+        .route("/", post(post_dataset_group))
+        .route("/", get(list_dataset_groups))
+        .route("/:dataset_group_id", get(get_dataset_group))
+        .route("/:dataset_group_id", delete(delete_dataset_group))
+        .route("/", put(put_dataset_group))
+        .route_layer(middleware::from_fn(auth))
+} 

api/src/routes/rest/routes/dataset_groups/post_dataset_group.rs

@@ -0,0 +1,89 @@
+use anyhow::Result;
+use axum::http::StatusCode;
+use axum::{Extension, Json};
+use chrono::Utc;
+use diesel::insert_into;
+use diesel_async::RunQueryDsl;
+use serde::{Deserialize, Serialize};
+use uuid::Uuid;
+
+use crate::database::lib::get_pg_pool;
+use crate::database::models::{DatasetGroup, User};
+use crate::database::schema::dataset_groups;
+use crate::routes::rest::ApiResponse;
+use crate::utils::security::checks::is_user_workspace_admin_or_data_admin;
+use crate::utils::user::user_info::get_user_organization_id;
+
+#[derive(Debug, Deserialize)]
+pub struct PostDatasetGroupRequest {
+    pub name: String,
+}
+
+#[derive(Debug, Serialize)]
+pub struct PostDatasetGroupResponse {
+    pub id: Uuid,
+    pub name: String,
+    pub created_at: chrono::DateTime<Utc>,
+    pub updated_at: chrono::DateTime<Utc>,
+}
+
+pub async fn post_dataset_group(
+    Extension(user): Extension<User>,
+    Json(request): Json<PostDatasetGroupRequest>,
+) -> Result<ApiResponse<PostDatasetGroupResponse>, (StatusCode, &'static str)> {
+    // Check if user is workspace admin or data admin
+    match is_user_workspace_admin_or_data_admin(&user.id).await {
+        Ok(true) => (),
+        Ok(false) => return Err((StatusCode::FORBIDDEN, "Insufficient permissions")),
+        Err(e) => {
+            tracing::error!("Error checking user permissions: {:?}", e);
+            return Err((
+                StatusCode::INTERNAL_SERVER_ERROR,
+                "Error checking user permissions",
+            ));
+        }
+    }
+
+    let dataset_group = match post_dataset_group_handler(request, user).await {
+        Ok(group) => group,
+        Err(e) => {
+            tracing::error!("Error creating dataset group: {:?}", e);
+            return Err((
+                StatusCode::INTERNAL_SERVER_ERROR,
+                "Error creating dataset group",
+            ));
+        }
+    };
+
+    Ok(ApiResponse::JsonData(PostDatasetGroupResponse {
+        id: dataset_group.id,
+        name: dataset_group.name,
+        created_at: dataset_group.created_at,
+        updated_at: dataset_group.updated_at,
+    }))
+}
+
+async fn post_dataset_group_handler(
+    request: PostDatasetGroupRequest,
+    user: User,
+) -> Result<DatasetGroup> {
+    let mut conn = get_pg_pool().get().await?;
+
+    let organization_id = get_user_organization_id(&user.id).await?;
+
+    let dataset_group = DatasetGroup {
+        id: Uuid::new_v4(),
+        organization_id,
+        name: request.name,
+        created_at: Utc::now(),
+        updated_at: Utc::now(),
+        deleted_at: None,
+    };
+
+    insert_into(dataset_groups::table)
+        .values(&dataset_group)
+        .execute(&mut *conn)
+        .await?;
+
+    Ok(dataset_group)
+}

api/src/routes/rest/routes/dataset_groups/put_dataset_group.rs

@@ -0,0 +1,84 @@
+use anyhow::Result;
+use axum::{http::StatusCode, Extension, Json};
+use chrono::Utc;
+use diesel::prelude::*;
+use diesel_async::RunQueryDsl;
+use serde::Deserialize;
+use uuid::Uuid;
+
+use crate::database::lib::get_pg_pool;
+use crate::database::models::User;
+use crate::database::schema::dataset_groups;
+use crate::routes::rest::ApiResponse;
+use crate::utils::security::checks::is_user_workspace_admin_or_data_admin;
+
+#[derive(Debug, Deserialize, Clone)]
+pub struct DatasetGroupUpdate {
+    pub id: Uuid,
+    pub name: String,
+}
+
+pub async fn put_dataset_group(
+    Extension(user): Extension<User>,
+    Json(request): Json<Vec<DatasetGroupUpdate>>,
+) -> Result<ApiResponse<()>, (StatusCode, &'static str)> {
+    // Check if user is workspace admin or data admin
+    match is_user_workspace_admin_or_data_admin(&user.id).await {
+        Ok(true) => (),
+        Ok(false) => return Err((StatusCode::FORBIDDEN, "Insufficient permissions")),
+        Err(e) => {
+            tracing::error!("Error checking user permissions: {:?}", e);
+            return Err((
+                StatusCode::INTERNAL_SERVER_ERROR,
+                "Error checking user permissions",
+            ));
+        }
+    }
+
+    match put_dataset_group_handler(request).await {
+        Ok(_) => Ok(ApiResponse::NoContent),
+        Err(e) => {
+            tracing::error!("Error updating dataset groups: {:?}", e);
+            Err((
+                StatusCode::INTERNAL_SERVER_ERROR,
+                "Error updating dataset groups",
+            ))
+        }
+    }
+}
+
+async fn put_dataset_group_handler(request: Vec<DatasetGroupUpdate>) -> Result<()> {
+    let now = Utc::now();
+
+    // Process in chunks of 25
+    let mut handles = vec![];
+    for chunk in request.chunks(25) {
+        let updates = chunk.to_vec();
+        let timestamp = now;
+
+        let handle = tokio::spawn(async move {
+            let mut conn = get_pg_pool().get().await?;
+
+            for update in updates {
+                diesel::update(dataset_groups::table)
+                    .filter(dataset_groups::id.eq(update.id))
+                    .filter(dataset_groups::deleted_at.is_null())
+                    .set((
+                        dataset_groups::name.eq(update.name),
+                        dataset_groups::updated_at.eq(timestamp),
+                    ))
+                    .execute(&mut *conn)
+                    .await?;
+            }
+            Ok::<_, anyhow::Error>(())
+        });
+        handles.push(handle);
+    }
+
+    // Wait for all tasks to complete
+    for handle in handles {
+        handle.await??;
+    }
+
+    Ok(())
+} 

api/src/routes/rest/routes/mod.rs

@@ -1,6 +1,7 @@
 mod api_keys;
 mod assets;
 mod data_sources;
+mod dataset_groups;
 mod datasets;
 mod permission_groups;
 mod users;
@@ -17,6 +18,7 @@ pub fn router() -> Router {
             .nest("/datasets", datasets::router())
             .nest("/data_sources", data_sources::router())
             .nest("/permission_groups", permission_groups::router())
+            .nest("/dataset_groups", dataset_groups::router())
             .route_layer(middleware::from_fn(auth)),
     )
 }