saas-rs-sdk 0.6.0

use crate::storage::Error;
use change_case::snake_case;
use futures_util::TryStreamExt;
use sqlx::{Pool, Postgres, Row};
use std::collections::HashMap;

#[derive(Clone, Debug)]
pub struct Field {
    pub column_name: String,
    pub data_type: String,
    pub is_nullable: bool,
}

#[derive(Clone, Debug)]
pub struct Table {
    pub name: String,
    pub fields: Vec<Field>,
    pub fields_by_name: HashMap<String, Field>,
    pub primary_key: Vec<String>,
}

pub async fn load_schema(pool: &Pool<Postgres>) -> Result<HashMap<String, Table>, Error> {
    // Load fields for all tables
    let query = "
        SELECT table_name, column_name, data_type, is_nullable
        FROM information_schema.columns
        WHERE table_schema = 'public'
    ";
    let mut rows = sqlx::query(query).fetch(pool);
    let mut tables_by_name: HashMap<String, Table> = HashMap::new();
    while let Some(row) = rows.try_next().await.map_err(|e| Error::internal(e.to_string()))? {
        let table_name: String = row.get(0);
        let column_name: String = row.get(1);
        let data_type: String = row.get(2);
        let is_nullable: bool = row.get::<&str, _>(3) == "YES";
        let column_name_snake_case = snake_case(&column_name);
        let field = Field {
            column_name,
            data_type,
            is_nullable,
        };
        if let Some(table) = tables_by_name.get_mut(&table_name) {
            table.fields.push(field.clone());
            table.fields_by_name.insert(column_name_snake_case, field);
        } else {
            let mut fields_by_name = HashMap::new();
            fields_by_name.insert(column_name_snake_case, field.clone());
            let table = Table {
                name: table_name.clone(),
                fields: vec![field],
                fields_by_name,
                primary_key: vec![],
            };
            tables_by_name.insert(table_name, table);
        }
    }

    // Load primary key info for all tables
    let query = "
        SELECT tc.table_name, c.column_name
        FROM information_schema.table_constraints tc
        JOIN information_schema.constraint_column_usage AS ccu USING (constraint_schema, constraint_name)
        JOIN information_schema.columns AS c
            ON c.table_schema = tc.constraint_schema
            AND tc.table_name = c.table_name AND ccu.column_name = c.column_name
            WHERE constraint_type = 'PRIMARY KEY'
    ";
    let mut rows = sqlx::query(query).fetch(pool);
    while let Some(row) = rows.try_next().await.map_err(|e| Error::internal(e.to_string()))? {
        let table_name: String = row.get(0);
        let column_name: String = row.get(1);
        let column_name_snake_case = snake_case(&column_name);
        if let Some(table) = tables_by_name.get_mut(&table_name) {
            table.primary_key.push(column_name_snake_case);
        }
    }

    Ok(tables_by_name)
}