athena_rs 1.1.0

//! Helpers that map table names to their resource ID columns via `table_id_map.yaml`.
//!
//! Falls back to finding the UUID column with the closest Levenshtein distance
//! to the table name when no explicit mapping exists.
use once_cell::sync::Lazy;
use serde::Deserialize;
use serde_yaml::from_str;
use std::collections::HashMap;
use strsim::levenshtein;

static TABLE_ID_YAML: &str = include_str!("../../../../table_id_map.yaml");

#[derive(Debug, Deserialize)]
/// Deserializes the YAML mapping between table names and ID columns.
struct TableIdConfig {
    /// Entries loaded from `table_id_map.yaml`.
    mappings: HashMap<String, String>,
}

/// Lazy-loaded map of table names -> resource ID column names.
static TABLE_ID_MAP: Lazy<HashMap<String, String>> = Lazy::new(|| {
    let cfg: TableIdConfig = from_str(TABLE_ID_YAML).expect("Invalid table_id_map.yaml");
    cfg.mappings
});

/// Retrieves UUID columns from a ScyllaDB table by querying system schema.
///
/// # Parameters
/// - `table_name`: The table name to query.
///
/// # Returns
/// Vector of column names that are of UUID type.
async fn get_uuid_columns_from_schema(table_name: &str) -> Vec<String> {
    let query: String = format!(
        "SELECT column_name, type FROM system_schema.columns WHERE keyspace_name = 'athena_rs' AND table_name = '{}' ALLOW FILTERING",
        table_name
    );

    match crate::drivers::scylla::client::execute_query(query).await {
        Ok((rows, _)) => rows
            .iter()
            .filter_map(|row| {
                let column_name = row.get("column_name")?.as_str()?;
                let column_type = row.get("type")?.as_str()?;

                if column_type == "uuid" {
                    Some(column_name.to_string())
                } else {
                    None
                }
            })
            .collect(),
        Err(err) => {
            tracing::warn!("Failed to query schema for table {}: {}", table_name, err);
            Vec::new()
        }
    }
}

/// Finds the UUID column with the smallest Levenshtein distance to the table name.
///
/// # Parameters
/// - `table_name`: The table name to match against.
/// - `uuid_columns`: List of UUID column names from the schema.
///
/// # Returns
/// The column name with the closest match, or `None` if no UUID columns exist.
#[doc(hidden)]
pub fn find_closest_uuid_column(table_name: &str, uuid_columns: &[String]) -> Option<String> {
    if uuid_columns.is_empty() {
        return None;
    }

    uuid_columns
        .iter()
        .map(|col| {
            let distance = levenshtein(table_name, col);
            (col.clone(), distance)
        })
        .min_by_key(|(_, distance)| *distance)
        .map(|(col, _)| col)
}

/// Retrieves the ID column name for the given table.
///
/// Resolution strategy:
/// 1. Check `table_id_map.yaml` for explicit mapping
/// 2. Query ScyllaDB schema for UUID columns
/// 3. Find UUID column with closest Levenshtein distance to table name
/// 4. Fallback to `"id"` if no matches found
///
/// # Parameters
/// - `table_name`: Table to look up inside `table_id_map.yaml`.
///
/// # Returns
/// Resource ID column name used during UUID enrichment.
///
/// # Example
/// ```text
/// get_resource_id_key("users").await // -> "user_id"
/// get_resource_id_key("unknown_table").await // -> closest UUID column or "id"
/// ```
pub async fn get_resource_id_key(table_name: &str) -> String {
    // First, check the static YAML map
    if let Some(mapped_key) = TABLE_ID_MAP.get(table_name) {
        return mapped_key.clone();
    }

    // Second, try to find the best UUID column from schema
    let uuid_columns: Vec<String> = get_uuid_columns_from_schema(table_name).await;
    if let Some(closest_column) = find_closest_uuid_column(table_name, &uuid_columns) {
        tracing::info!(
            "Dynamic fallback for table '{}': using column '{}'",
            table_name,
            closest_column
        );
        return closest_column;
    }

    // Final fallback
    tracing::warn!(
        "No UUID column found for table '{}', falling back to 'id'",
        table_name
    );
    "id".to_string()
}