athena_rs 2.0.2

use once_cell::sync::Lazy;
use sqlx::PgPool;
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::RwLock;

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 = TABLE_COLUMN_TYPE_CACHE.read().await;
        if let Some(columns) = cache.get(table_name) {
            return Ok(columns.clone());
        }
    }

    let rows = 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 = rows.into_iter().collect::<HashMap<_, _>>();

    let mut cache = TABLE_COLUMN_TYPE_CACHE.write().await;
    cache.insert(table_name.to_string(), columns.clone());

    Ok(columns)
}