athena_rs 0.75.4

#![allow(unused)]

//! Execute CQL queries against ScyllaDB and return rows as JSON.
//!
//! Connects to `SCYLLA_URI` (default `127.0.0.1:9042`).
use crate::config::Config;
use scylla::client::session::Session;
use scylla::client::session_builder::SessionBuilder;
use scylla::response::query_result::QueryResult;
use serde_json::{Map, Number, Value, json};

use scylla::deserialize::result::TypedRowIterator;
use scylla::response::query_result::{ColumnSpecs, QueryRowsResult};
use scylla::statement::Consistency;
use scylla::statement::Statement;
use scylla::value::{CqlValue, Row};
use std::collections::HashMap;
use std::error::Error;

/// Execute a CQL statement and return `(rows, columns)` where:
/// - `rows` is a vector of JSON objects representing result rows
/// - `columns` is the ordered list of column names in the result set
#[allow(unused_variables)]
pub async fn execute_query(query: String) -> Result<(Vec<Value>, Vec<String>), Box<dyn Error>> {
    // Create a new Session which connects to node at 127.0.0.1:9042
    // (or SCYLLA_URI if specified)
    let config: Config = Config::load().unwrap_or_else(|_| panic!("Failed to load config"));
    let uri: String = config
        .get_host("scylladb")
        .unwrap_or(&"167.235.9.113:9042".to_string())
        .clone();

    let empty_map: HashMap<String, String> = HashMap::new();
    let authenticator: &HashMap<String, String> =
        config.get_authenticator("scylladb").unwrap_or(&empty_map);
    let empty_string: String = String::new();
    let session: Session = SessionBuilder::new()
        .known_node("167.235.9.113:9042")
        .build()
        .await?;

    // Execute the provided query and collect results
    let mut results: Vec<Value> = Vec::new();
    let mut columns: Vec<String> = Vec::new();

    // Create query with consistency level and unpaged settings

    let mut cql_query: Statement = Statement::new(query.as_str());
    cql_query.set_consistency(Consistency::One);
    cql_query.set_page_size(1);

    // Execute the query
    let query_result: QueryResult = session.query_unpaged(cql_query, &[]).await?;

    // Check if this is a rows result (SELECT query)
    if query_result.is_rows() {
        // Convert the QueryResult to QueryRowsResult to access rows
        let rows_result: QueryRowsResult = query_result.into_rows_result()?;

        // Get column specifications
        let column_specs: ColumnSpecs<'_, '_> = rows_result.column_specs();

        // Capture and log column names
        columns = column_specs
            .iter()
            .map(|spec| spec.name().to_string())
            .collect();

        // Parse rows using the generic approach with Value
        let rows: TypedRowIterator<'_, '_, Row> = rows_result.rows::<Row>()?;

        let mut row_count: i32 = 0;
        for row_result in rows {
            let row: Row = row_result?;
            row_count += 1;
            let mut row_map: Map<String, Value> = Map::new();
            for (i, column_spec) in column_specs.iter().enumerate() {
                let column_name: String = column_spec.name().to_string();
                let value: Value = match row.columns.get(i) {
                    Some(Some(cql_value)) => {
                        match cql_value {
                            CqlValue::Text(s) => Value::String(s.clone()),
                            CqlValue::Int(i) => Value::Number((*i).into()),
                            CqlValue::BigInt(i) => Value::Number((*i).into()),
                            CqlValue::Boolean(b) => Value::Bool(*b),
                            CqlValue::Float(f) => Value::Number(
                                Number::from_f64(*f as f64).unwrap_or(Number::from(0)),
                            ),
                            CqlValue::Double(d) => {
                                Value::Number(Number::from_f64(*d).unwrap_or(Number::from(0)))
                            }
                            CqlValue::Decimal(decimal) => {
                                // Convert CqlDecimal to a string representation and then parse as f64
                                let decimal_str: String = format!("{:?}", decimal);

                                // Try to extract numeric value from the debug representation
                                // This is a fallback approach since CqlDecimal doesn't implement Display
                                let decimal_value: f64 =
                                    if let Some(start) = decimal_str.find("value: ") {
                                        let value_part: &str = &decimal_str[start + 7..];
                                        if let Some(end) = value_part.find(',') {
                                            value_part[..end].parse().unwrap_or(0.0)
                                        } else {
                                            0.0
                                        }
                                    } else {
                                        0.0
                                    };

                                Value::Number(
                                    Number::from_f64(decimal_value).unwrap_or(Number::from(0)),
                                )
                            }
                            _ => Value::String(format!("{:?}", cql_value)),
                        }
                    }
                    Some(None) => Value::Null,
                    None => Value::Null,
                };

                row_map.insert(column_name, value);
            }

            results.push(Value::Object(row_map));
        }
    } else {
        // For non-SELECT queries, return success indicator
        results.push(json!({
            "status": "success",
            "message": "Query executed successfully"
        }));
    }

    Ok((results, columns))
}