saas-rs-sdk 0.6.3

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

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

#[derive(Clone, Debug)]
pub(crate) struct ForeignKey {
    pub table_name: String,
    #[allow(unused)]
    pub column_name: String,
    pub foreign_table_name: String,
    #[allow(unused)]
    pub foreign_column_name: String,
}

#[derive(Clone, Debug)]
pub(crate) struct JoinTable {
    #[allow(unused)]
    pub table_name: String,
    pub foreign_keys: Vec<ForeignKey>,
}

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

pub(crate) 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? {
        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? {
        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)
}

pub(crate) async fn load_foreign_keys(pool: &Pool<Postgres>) -> Result<Vec<ForeignKey>, Error> {
    // Query for foreign key constraints
    let query = "
        SELECT
            tc.table_name,
            kcu.column_name,
            ccu.table_name AS foreign_table_name,
            ccu.column_name AS foreign_column_name
        FROM information_schema.table_constraints AS tc
        JOIN information_schema.key_column_usage AS kcu
            ON tc.constraint_name = kcu.constraint_name
            AND tc.table_schema = kcu.table_schema
        JOIN information_schema.constraint_column_usage AS ccu
            ON ccu.constraint_name = tc.constraint_name
        WHERE tc.constraint_type = 'FOREIGN KEY'
    ";
    let mut rows = sqlx::query(query).fetch(pool);

    // Collect results
    let mut foreign_keys = vec![];
    while let Some(row) = rows.try_next().await? {
        let table_name: String = row.get(0);
        let column_name: String = row.get(1);
        let foreign_table_name: String = row.get(2);
        let foreign_column_name: String = row.get(3);
        let foreign_key = ForeignKey {
            table_name,
            column_name,
            foreign_table_name,
            foreign_column_name,
        };
        foreign_keys.push(foreign_key);
    }
    Ok(foreign_keys)
}

// Find all the join tables (ones with a 2-column compound primary key)
pub(crate) fn find_join_tables(
    tables_by_name: &HashMap<String, Table>,
    foreign_keys: &[ForeignKey],
) -> HashMap<String, JoinTable> {
    let mut join_tables_by_name: HashMap<String, JoinTable> = HashMap::new();
    for fk in foreign_keys.iter().filter_map(|fk| {
        tables_by_name
            .get(&fk.table_name)
            .filter(|&table| table.primary_key.len() == 2)
            .map(|_table| fk)
    }) {
        if let Some(join_table) = join_tables_by_name.get_mut(&fk.table_name) {
            join_table.foreign_keys.push(fk.clone());
        } else {
            let join_table = crate::storage::config_store::adapters::postgres::schema::JoinTable {
                table_name: fk.table_name.clone(),
                foreign_keys: vec![fk.clone()],
            };
            join_tables_by_name.insert(fk.table_name.clone(), join_table);
        }
    }
    join_tables_by_name
}