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;
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 {
let query = 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()]
);
}
}