rquest/src-tauri/src/db.rs

use crate::{
    auth::{Auth, Authentication},
    error::AppError,
    request::{
        EntryRequestBody, RequestBody, RequestHeader, RequestHeaderInsert, RequestHeaderUpdate,
        RequestParams, RequestPathParam, WorkspaceRequest,
    },
    workspace::{
        Workspace, WorkspaceEntry, WorkspaceEntryBase, WorkspaceEntryCreate, WorkspaceEntryType,
        WorkspaceEntryUpdate, WorkspaceEnvVariable, WorkspaceEnvironment,
    },
    AppResult,
};
use serde::Deserialize;
use sqlx::{sqlite::SqlitePool, types::Json, QueryBuilder};
use std::collections::HashMap;
use tauri_plugin_log::log;

/// Used in update DTOs for **optional** properties. A value indicates a parameter needs to be updated to whatever is contained in it, a null indicates to set the field to null.
#[derive(Debug, Deserialize)]
pub enum Update<T> {
    Value(T),
    Null,
}

impl<T: Clone> Clone for Update<T> {
    fn clone(&self) -> Self {
        match self {
            Self::Value(arg0) => Self::Value(arg0.clone()),
            Self::Null => Self::Null,
        }
    }
}

impl<T: Copy> Copy for Update<T> {}

impl<T: Copy> Update<T> {
    pub fn value(self) -> Option<T> {
        match self {
            Update::Value(v) => Some(v),
            Update::Null => None,
        }
    }
}

pub async fn init(url: &str) -> SqlitePool {
    let pool = SqlitePool::connect(url)
        .await
        .expect("error while connecting to db");

    sqlx::migrate!()
        .run(&pool)
        .await
        .expect("error in migrations");

    pool
}

pub async fn create_workspace(db: SqlitePool, name: String) -> Result<Workspace, String> {
    match sqlx::query_as!(
        Workspace,
        "INSERT INTO workspaces (name) VALUES (?) RETURNING id, name",
        name
    )
    .fetch_one(&db)
    .await
    {
        Ok(workspace) => Ok(workspace),
        Err(e) => Err(e.to_string()),
    }
}

pub async fn list_workspaces(db: SqlitePool) -> AppResult<Vec<Workspace>> {
    Ok(
        sqlx::query_as!(Workspace, "SELECT id, name FROM workspaces")
            .fetch_all(&db)
            .await?,
    )
}

/// Check whether the entry whose `id` == `parent_id` supports children.
async fn check_parent(db: &SqlitePool, parent_id: Option<i64>) -> AppResult<()> {
    let Some(parent_id) = parent_id else {
        return Ok(());
    };

    let ty = sqlx::query!("SELECT type FROM workspace_entries WHERE id = ?", parent_id)
        .fetch_one(db)
        .await?
        .r#type;

    if !matches!(WorkspaceEntryType::from(ty), WorkspaceEntryType::Collection) {
        return Err(AppError::InvalidUpdate(format!(
            "{parent_id} is not a valid parent ID (type: {ty})"
        )));
    }

    Ok(())
}

pub async fn create_workspace_entry(
    db: SqlitePool,
    entry: WorkspaceEntryCreate,
) -> AppResult<WorkspaceEntryBase> {
    match entry {
        WorkspaceEntryCreate::Collection {
            name,
            workspace_id,
            parent_id,
        } => {
            check_parent(&db, parent_id).await?;
            let entry = sqlx::query_as!(
                WorkspaceEntryBase,
                r#"INSERT INTO workspace_entries(name, workspace_id, parent_id, type) VALUES (?, ?, ?, ?) 
                   RETURNING id, workspace_id, parent_id, name, type, auth, auth_inherit"#,
                name,
                workspace_id,
                parent_id,
                1)
            .fetch_one(&db).await?;

            Ok(entry)
        }
        WorkspaceEntryCreate::Request {
            name,
            workspace_id,
            parent_id,
            method,
            url,
        } => {
            if let Some(parent) = parent_id {
                let ty = sqlx::query!("SELECT type FROM workspace_entries WHERE id = ?", parent)
                    .fetch_one(&db)
                    .await?
                    .r#type;

                if !matches!(WorkspaceEntryType::from(ty), WorkspaceEntryType::Collection) {
                    return Err(AppError::InvalidUpdate(format!(
                        "{parent} is not a valid parent ID (type: {ty})"
                    )));
                }
            }

            let mut tx = db.begin().await?;

            let entry = match sqlx::query_as!(
                WorkspaceEntryBase,
                r#"INSERT INTO workspace_entries(name, workspace_id, parent_id, type) VALUES (?, ?, ?, ?) 
                   RETURNING id, workspace_id, name, parent_id, type, auth, auth_inherit"#,
                name,
                workspace_id,
                parent_id,
                0)
            .fetch_one(&mut *tx).await {
                Ok(entry) => entry,
                Err(e) => {
                    tx.rollback().await?;
                    return Err(e.into());
                }

            };

            match sqlx::query!(
                "INSERT INTO request_params(workspace_id, request_id, method, url) VALUES (?, ?, ?, ?)",
                workspace_id,
                entry.id,
                method,
                url
            )
            .execute(&mut *tx)
            .await {
                Ok(_) => {},
                Err(e) => {
                    tx.rollback().await?;
                    return Err(e.into());
                }
            }
            ;

            tx.commit().await?;

            Ok(entry)
        }
    }
}

pub async fn insert_request_body(
    db: SqlitePool,
    entry_id: i64,
    body: RequestBody,
) -> AppResult<EntryRequestBody> {
    Ok(sqlx::query_as!(EntryRequestBody, r#"INSERT INTO request_bodies(request_id, content_type, body) VALUES (?, ?, ?) RETURNING id, content_type AS "content_type: _", body"#, entry_id, body.ty, body.content).fetch_one(&db).await?)
}

pub async fn update_request_body(
    db: SqlitePool,
    id: i64,
    body: Update<RequestBody>,
) -> AppResult<()> {
    match body {
        Update::Value(body) => {
            sqlx::query!(
                "UPDATE request_bodies SET content_type = ?, body = ? WHERE id = ?",
                body.ty,
                body.content,
                id,
            )
            .execute(&db)
            .await?;
        }
        Update::Null => {
            sqlx::query!("DELETE FROM request_bodies WHERE id = ?", id)
                .execute(&db)
                .await?;
        }
    }
    Ok(())
}

pub async fn update_workspace_entry(
    db: SqlitePool,
    entry_id: i64,
    update: WorkspaceEntryUpdate,
) -> AppResult<()> {
    match update {
        WorkspaceEntryUpdate::Collection(update) => {
            let mut sql = sqlx::query_builder::QueryBuilder::new("UPDATE workspace_entries SET ");

            if let Some(parent) = update.parent_id {
                check_parent(&db, parent.value()).await?;
            }

            match (update.name, update.parent_id) {
                (None, None) => {
                    return Err(AppError::InvalidUpdate(
                        "cannot update entry: no updates present".to_string(),
                    ))
                }
                (None, Some(parent_id)) => match parent_id {
                    Update::Value(v) => {
                        sql.push("parent_id = ").push_bind(v);
                    }
                    Update::Null => {
                        sql.push("parent_id = NULL ");
                    }
                },
                (Some(name), None) => {
                    sql.push("name = ").push_bind(name);
                }
                (Some(name), Some(parent_id)) => {
                    match parent_id {
                        Update::Value(v) => {
                            sql.push("parent_id = ").push_bind(v);
                        }
                        Update::Null => {
                            sql.push("parent_id = NULL ");
                        }
                    };
                    sql.push(", name = ").push_bind(name);
                }
            }

            sql.push("WHERE id = ")
                .push_bind(entry_id)
                .build()
                .execute(&db)
                .await?;

            Ok(())
        }
        WorkspaceEntryUpdate::Request {
            base,
            method,
            url,
            path_params,
        } => {
            let mut tx = db.begin().await?;

            'entry: {
                if let Some(parent) = base.parent_id {
                    check_parent(&db, parent.value()).await?;
                }

                let mut sql =
                    sqlx::query_builder::QueryBuilder::new("UPDATE workspace_entries SET ");

                match (base.name, base.parent_id) {
                    (None, None) => break 'entry,
                    (None, Some(parent_id)) => match parent_id {
                        Update::Value(v) => {
                            sql.push("parent_id = ").push_bind(v);
                        }
                        Update::Null => {
                            sql.push("parent_id = NULL ");
                        }
                    },
                    (Some(name), None) => {
                        sql.push("name = ").push_bind(name);
                    }
                    (Some(name), Some(parent_id)) => {
                        match parent_id {
                            Update::Value(v) => {
                                sql.push("parent_id = ").push_bind(v);
                            }
                            Update::Null => {
                                sql.push("parent_id = NULL ");
                            }
                        };
                        sql.push(", name = ").push_bind(name);
                    }
                }

                sql.push("WHERE id = ")
                    .push_bind(entry_id)
                    .build()
                    .execute(&mut *tx)
                    .await?;
            };

            'param: {
                let mut sql = sqlx::query_builder::QueryBuilder::new("UPDATE request_params ");

                match (method, url) {
                    (None, None) => break 'param,
                    (None, Some(url)) => {
                        sql.push("SET url = ").push_bind(url);
                    }
                    (Some(method), None) => {
                        sql.push("SET method = ").push_bind(method);
                    }
                    (Some(method), Some(url)) => {
                        sql.push("SET method = ")
                            .push_bind(method)
                            .push(", url = ")
                            .push_bind(url);
                    }
                }

                sql.push("WHERE request_id = ")
                    .push_bind(entry_id)
                    .build()
                    .execute(&mut *tx)
                    .await?;
            };

            if let Some(path_params) = path_params {
                if path_params.is_empty() {
                    sqlx::query!(
                        "DELETE FROM request_path_params WHERE request_id = ?",
                        entry_id
                    )
                    .execute(&mut *tx)
                    .await?;
                } else {
                    let mut sql = QueryBuilder::new(
                        "INSERT INTO request_path_params(position, request_id, name, value) ",
                    );

                    sql.push_values(path_params.iter(), |mut b, path| {
                        b.push_bind(path.position as i64)
                            .push_bind(entry_id)
                            .push_bind(&path.name);

                        if let Some(ref path) = path.value {
                            b.push_bind(path);
                        } else {
                            b.push_bind("");
                        }
                    });

                    sql.push(
                        r#"
                        ON CONFLICT(position, request_id) DO UPDATE 
                        SET 
                            value = excluded.value,
                            name = excluded.name;

                        DELETE FROM request_path_params 
                        WHERE request_id = "#,
                    )
                    .push_bind(entry_id)
                    .push(" AND position NOT IN (");

                    let mut sep = sql.separated(", ");

                    for param in path_params.iter() {
                        sep.push_bind(param.position as i64);
                    }

                    sep.push_unseparated(")");

                    sql.build().execute(&mut *tx).await?;
                }
            }

            tx.commit().await?;

            Ok(())
        }
    }
}

pub async fn get_workspace_request(db: SqlitePool, id: i64) -> AppResult<WorkspaceRequest> {
    let entry = sqlx::query_as!(
        WorkspaceEntryBase,
        "SELECT id, workspace_id, parent_id, name, type, auth, auth_inherit FROM workspace_entries WHERE id = ?",
        id,
    )
    .fetch_one(&db)
    .await?;

    let params = sqlx::query_as!(
        RequestParams,
        r#"
           SELECT 
            rp.request_id as id,
            method as 'method!',
            url as 'url!',
            content_type as "content_type: _",
            body AS "body: _",
            rb.id AS "body_id: _"
           FROM request_params rp
           LEFT JOIN request_bodies rb ON rp.request_id = rb.request_id
           WHERE rp.request_id = ? 
        "#,
        id
    )
    .fetch_one(&db)
    .await?;

    let headers = sqlx::query_as!(
        RequestHeader,
        "SELECT id, name, value FROM request_headers WHERE request_id = ?",
        entry.id
    )
    .fetch_all(&db)
    .await?;

    let path_params = sqlx::query_as!(
        RequestPathParam,
        "SELECT position, name, value FROM request_path_params WHERE request_id = ?",
        entry.id
    )
    .fetch_all(&db)
    .await?;

    Ok(WorkspaceRequest::from_params_and_headers(
        entry,
        params,
        headers,
        path_params,
    ))
}

pub async fn list_workspace_entries(
    db: SqlitePool,
    workspace_id: i64,
) -> AppResult<Vec<WorkspaceEntry>> {
    let entries = sqlx::query_as!(
        WorkspaceEntryBase,
        "SELECT id, workspace_id, parent_id, name, type, auth, auth_inherit FROM workspace_entries WHERE workspace_id = ? ORDER BY type DESC",
        workspace_id,
    )
    .fetch_all(&db)
    .await?;

    let mut request_params: HashMap<i64, RequestParams> = sqlx::query_as!(
        RequestParams,
        r#"
           SELECT rp.request_id as id, method as 'method!', url as 'url!', content_type as "content_type: _", body, rb.id as "body_id: _"
           FROM request_params rp
           LEFT JOIN request_bodies rb ON rp.request_id = rb.request_id
           WHERE workspace_id = ? 
        "#,
        workspace_id
    )
    .fetch_all(&db)
    .await?
    .into_iter()
    .map(|req| (req.id, req))
    .collect();

    let mut out: Vec<WorkspaceEntry> = vec![];

    for entry in entries {
        match entry.r#type {
            WorkspaceEntryType::Request => {
                let headers = sqlx::query_as!(
                    RequestHeader,
                    "SELECT id, name, value FROM request_headers WHERE request_id = ?",
                    entry.id
                )
                .fetch_all(&db)
                .await?;

                let path_params = sqlx::query_as!(
                    RequestPathParam,
                    "SELECT position, name, value FROM request_path_params WHERE request_id = ?",
                    entry.id
                )
                .fetch_all(&db)
                .await?;

                let Some(params) = request_params.remove(&entry.id) else {
                    log::warn!("request {} has no params!", entry.id);
                    continue;
                };

                let req =
                    WorkspaceRequest::from_params_and_headers(entry, params, headers, path_params);

                out.push(WorkspaceEntry::new_req(req));
            }
            WorkspaceEntryType::Collection => {
                out.push(WorkspaceEntry::new_col(entry));
            }
        }
    }

    Ok(out)
}

pub async fn list_environments(
    db: SqlitePool,
    workspace_id: i64,
) -> AppResult<Vec<WorkspaceEnvironment>> {
    let records = sqlx::query!(
        r#"
          SELECT 
            env.workspace_id,
            env.id AS env_id,
            env.name AS env_name,
            var.id AS "var_id?",
            var.name AS "var_name?",
            var.value AS "var_value?",
            var.secret AS "var_secret?"
          FROM workspace_envs env
          LEFT JOIN workspace_env_variables var ON env.id = var.env_id
          WHERE env.workspace_id = $1"#,
        workspace_id
    )
    .fetch_all(&db)
    .await?;

    let mut environments: HashMap<i64, WorkspaceEnvironment> = HashMap::new();

    for record in records {
        if let Some(env) = environments.get_mut(&record.env_id) {
            if record.var_id.is_some() {
                env.variables.push(WorkspaceEnvVariable {
                    id: record.var_id.unwrap(),
                    workspace_id,
                    env_id: record.env_id,
                    name: record.var_name.unwrap(),
                    value: record.var_value.unwrap(),
                    secret: record.var_secret.unwrap(),
                })
            }
        } else {
            let mut env = WorkspaceEnvironment {
                id: record.env_id,
                name: record.env_name,
                workspace_id,
                variables: vec![],
            };
            if record.var_id.is_some() {
                env.variables.push(WorkspaceEnvVariable {
                    id: record.var_id.unwrap(),
                    workspace_id,
                    env_id: record.env_id,
                    name: record.var_name.unwrap(),
                    value: record.var_value.unwrap(),
                    secret: record.var_secret.unwrap(),
                })
            }
            environments.insert(record.env_id, env);
        }
    }

    Ok(environments.into_values().collect())
}

pub async fn get_env_variables(
    db: &SqlitePool,
    env_id: i64,
    names: &[&str],
) -> AppResult<Vec<(String, String)>> {
    let mut query =
        QueryBuilder::new("SELECT name, value FROM workspace_env_variables WHERE env_id = ");

    query.push_bind(env_id);

    let mut separated = query.push(" AND name IN (").separated(", ");

    for name in names {
        separated.push_bind(name);
    }

    separated.push_unseparated(")");

    Ok(query
        .build_query_as::<(String, String)>()
        .fetch_all(db)
        .await?)
}

pub async fn create_environment(
    db: SqlitePool,
    workspace_id: i64,
    name: String,
) -> AppResult<WorkspaceEnvironment> {
    let row = sqlx::query!(
        r#"
        INSERT INTO workspace_envs (workspace_id, name)
        VALUES (?, ?)
        RETURNING id, workspace_id, name
        "#,
        workspace_id,
        name
    )
    .fetch_one(&db)
    .await?;

    Ok(WorkspaceEnvironment {
        id: row.id,
        workspace_id: row.workspace_id,
        name: row.name,
        variables: vec![],
    })
}

pub async fn update_environment(db: SqlitePool, env_id: i64, name: String) -> AppResult<()> {
    sqlx::query!(
        r#"
        UPDATE workspace_envs SET name = ? WHERE id = ?
        "#,
        name,
        env_id,
    )
    .execute(&db)
    .await?;

    Ok(())
}

pub async fn insert_env_var(
    db: SqlitePool,
    workspace_id: i64,
    env_id: i64,
    name: String,
    value: String,
    secret: bool,
) -> AppResult<WorkspaceEnvVariable> {
    Ok(sqlx::query_as!(
        WorkspaceEnvVariable,
        r#"
        INSERT INTO workspace_env_variables (workspace_id, env_id, name, value, secret)
        VALUES (?, ?, ?, ?, ?)
        RETURNING id, workspace_id, env_id, name, value, secret
        "#,
        workspace_id,
        env_id,
        name,
        value,
        secret,
    )
    .fetch_one(&db)
    .await?)
}

pub async fn update_env_var(
    db: SqlitePool,
    id: i64,
    name: Option<String>,
    value: Option<String>,
    secret: Option<bool>,
) -> AppResult<()> {
    sqlx::query_as!(
        WorkspaceEnvVariable,
        r#"
        UPDATE workspace_env_variables
        SET
            name   = COALESCE(?, name),
            value  = COALESCE(?, value),
            secret = COALESCE(?, secret)
        WHERE id = ?
        "#,
        name,
        value,
        secret,
        id,
    )
    .execute(&db)
    .await?;
    Ok(())
}

pub async fn delete_env_var(db: SqlitePool, id: i64) -> AppResult<()> {
    sqlx::query_as!(
        WorkspaceEnvVariable,
        r#"
        DELETE FROM workspace_env_variables
        WHERE id = ?
        "#,
        id,
    )
    .execute(&db)
    .await?;
    Ok(())
}

pub async fn list_request_path_params(db: SqlitePool, id: i64) -> AppResult<Vec<RequestPathParam>> {
    Ok(sqlx::query_as!(
        RequestPathParam,
        "SELECT position, name, value FROM request_path_params WHERE request_id = ?",
        id
    )
    .fetch_all(&db)
    .await?)
}

pub async fn insert_headers(
    db: SqlitePool,
    entry_id: i64,
    headers: Vec<RequestHeaderInsert>,
) -> AppResult<RequestHeader> {
    let mut insert = QueryBuilder::new("INSERT INTO request_headers(request_id, name, value) ");
    insert.push_values(headers, |mut b, header| {
        b.push_bind(entry_id)
            .push_bind(header.name)
            .push_bind(header.value);
    });
    Ok(insert
        .push("RETURNING id, name, value")
        .build_query_as()
        .fetch_one(&db)
        .await?)
}

pub async fn update_header(db: SqlitePool, header: RequestHeaderUpdate) -> AppResult<()> {
    sqlx::query!(
        "UPDATE request_headers SET name = COALESCE(?, ''), value = COALESCE(?, '') WHERE id = ?",
        header.name,
        header.value,
        header.id
    )
    .execute(&db)
    .await?;
    Ok(())
}

pub async fn delete_header(db: SqlitePool, header_id: i64) -> AppResult<()> {
    sqlx::query!("DELETE FROM request_headers WHERE id = ?", header_id)
        .execute(&db)
        .await?;
    Ok(())
}

pub async fn insert_auth(
    db: SqlitePool,
    workspace_id: i64,
    params: Auth,
) -> AppResult<Authentication> {
    let json = Json(&params);

    let record = sqlx::query!(
        "INSERT INTO auth(workspace_id, name, params) VALUES (?, 'New authentication', ?) RETURNING id, name",
        workspace_id,
        json
    )
    .fetch_one(&db)
    .await?;

    Ok(Authentication {
        id: record.id,
        workspace_id,
        name: record.name,
        params,
    })
}

pub async fn delete_auth(db: SqlitePool, id: i64) -> AppResult<()> {
    sqlx::query!("DELETE FROM auth WHERE id = ?", id)
        .execute(&db)
        .await?;
    Ok(())
}

pub async fn list_auth(db: SqlitePool, workspace_id: i64) -> AppResult<Vec<Authentication>> {
    let records = sqlx::query!(
        r#"
        SELECT id, name, workspace_id, params as "params: Json<Auth>"
        FROM auth
        WHERE workspace_id = ?
        "#,
        workspace_id
    )
    .fetch_all(&db)
    .await?;

    Ok(records
        .into_iter()
        .map(|record| Authentication {
            id: record.id,
            name: record.name,
            workspace_id: record.workspace_id,
            params: record.params.0,
        })
        .collect())
}

pub async fn get_auth(db: SqlitePool, id: i64) -> AppResult<Authentication> {
    let record = sqlx::query!(
        r#"
        SELECT id, workspace_id, name, params as "params: Json<Auth>"
        FROM auth
        WHERE id = ?
        "#,
        id
    )
    .fetch_one(&db)
    .await?;

    Ok(Authentication {
        id: record.id,
        name: record.name,
        workspace_id: record.workspace_id,
        params: record.params.0,
    })
}

pub async fn set_workspace_entry_auth(
    db: SqlitePool,
    entry_id: i64,
    auth_id: Option<i64>,
    inherit: Option<bool>,
) -> AppResult<()> {
    sqlx::query!(
        "UPDATE workspace_entries SET auth = ?, auth_inherit = COALESCE(?, auth_inherit) WHERE id = ?",
        auth_id,
        inherit,
        entry_id
    )
    .execute(&db)
    .await?;

    Ok(())
}

pub async fn update_auth(db: SqlitePool, auth_id: i64, params: Auth) -> AppResult<()> {
    let params = Json(params);

    sqlx::query!("UPDATE auth SET params = ? WHERE id = ?", params, auth_id)
        .execute(&db)
        .await?;

    Ok(())
}

pub async fn rename_auth(db: SqlitePool, auth_id: i64, name: String) -> AppResult<()> {
    sqlx::query!("UPDATE auth SET name = ? WHERE id = ?", name, auth_id)
        .execute(&db)
        .await?;

    Ok(())
}

/// Check for the existence of an auth ID in the workspace entry. If one does not exist,
/// traverse its parents and attempt to find the first one that is present. If none exist,
/// returns `None`.
pub async fn get_auth_inherited(
    db: SqlitePool,
    mut parent_id: Option<i64>,
) -> AppResult<Option<i64>> {
    while let Some(id) = parent_id {
        let record = sqlx::query!(
            "SELECT auth, auth_inherit, parent_id FROM workspace_entries WHERE id = ?",
            id
        )
        .fetch_one(&db)
        .await?;

        if !record.auth_inherit {
            return Ok(record.auth);
        }

        parent_id = record.parent_id;
    }

    Ok(None)
}