athena_rs 3.3.0

use once_cell::sync::Lazy;
use sqlx::PgPool;
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::RwLockWriteGuard;
use tokio::sync::{RwLock, RwLockReadGuard};

type TableSchemaCache = HashMap<String, HashMap<String, String>>;

static TABLE_COLUMN_TYPE_CACHE: Lazy<Arc<RwLock<TableSchemaCache>>> =
    Lazy::new(|| Arc::new(RwLock::new(HashMap::new())));

pub async fn get_public_table_column_types(
    pool: &PgPool,
    table_name: &str,
) -> Result<HashMap<String, String>, sqlx::Error> {
    {
        let cache: RwLockReadGuard<'_, HashMap<String, HashMap<String, String>>> =
            TABLE_COLUMN_TYPE_CACHE.read().await;
        if let Some(columns) = cache.get(table_name) {
            return Ok(columns.clone());
        }
    }

    let rows: Vec<(String, String)> = sqlx::query_as::<_, (String, String)>(
        r#"
        SELECT column_name, data_type
        FROM information_schema.columns
        WHERE table_schema = 'public'
          AND table_name = $1
        "#,
    )
    .bind(table_name)
    .fetch_all(pool)
    .await?;

    let columns: HashMap<String, String> = rows.into_iter().collect::<HashMap<_, _>>();

    let mut cache: RwLockWriteGuard<'_, HashMap<String, HashMap<String, String>>> =
        TABLE_COLUMN_TYPE_CACHE.write().await;
    cache.insert(table_name.to_string(), columns.clone());

    Ok(columns)
}