athena_rs 3.23.0

Hyper performant polyglot Database driver
Documentation
//! Compatibility facade for the gateway resource-ID fallback resolver.
//!
//! The cache, UUID-column parsing, and closest-column heuristic now live in
//! `athena-gateway`. This module keeps the historical import path stable while
//! `athena_rs` remains responsible for the runtime Scylla schema query.

pub use athena_gateway::{find_closest_uuid_column, parse_uuid_columns_from_schema_rows};

use athena_gateway::get_resource_id_key_with_uuid_loader;

use crate::drivers::scylla::client::execute_query;

/// Retrieves the resource-ID column name for the given gateway table.
///
/// Athena still loads UUID columns through the runtime Scylla client wrapper
/// here, while the extracted `athena-gateway` crate owns the cache and fallback
/// heuristic that turn those UUID columns into a stable resource identifier.
pub async fn get_resource_id_key(table_name: &str) -> String {
    get_resource_id_key_with_uuid_loader(table_name, |safe_table_name| async move {
        // TODO: Raw sql query, should get typed
        let query: String = format!(
            "SELECT column_name, type FROM system_schema.columns WHERE keyspace_name = 'athena_rs' AND table_name = '{}' ALLOW FILTERING",
            safe_table_name
        );

        match execute_query(query).await {
            Ok((rows, _)) => parse_uuid_columns_from_schema_rows(&rows),
            Err(err) => {
                tracing::warn!(
                    "Failed to query schema for table {}: {}",
                    safe_table_name,
                    err
                );
                Vec::new()
            }
        }
    })
    .await
}

#[cfg(test)]
mod tests {
    use super::{find_closest_uuid_column, parse_uuid_columns_from_schema_rows};
    use serde_json::json;

    #[test]
    fn closest_uuid_column_is_selected_via_wrapper() {
        let columns = vec![
            "id".to_string(),
            "order_item_id".to_string(),
            "order_id".to_string(),
        ];

        let closest = find_closest_uuid_column("orders", &columns);
        assert_eq!(closest.as_deref(), Some("order_id"));
    }

    #[test]
    fn uuid_columns_are_parsed_via_wrapper() {
        let rows = vec![
            json!({ "column_name": "user_id", "type": "uuid" }),
            json!({ "column_name": "created_at", "type": "timestamp" }),
        ];

        assert_eq!(
            parse_uuid_columns_from_schema_rows(&rows),
            vec!["user_id".to_string()]
        );
    }
}