Update ElasticSQL handlers

Author: kavyabhat02

asap-query-engine/src/data_model/streaming_config.rs

@@ -72,7 +72,7 @@ impl StreamingConfig {
                 SchemaConfig::PromQL(_) => QueryLanguage::promql,
                 SchemaConfig::SQL(_) => QueryLanguage::sql,
                 SchemaConfig::ElasticQueryDSL => QueryLanguage::elastic_querydsl,
-                SchemaConfig::ElasticSQL => QueryLanguage::elastic_sql,
+                SchemaConfig::ElasticSQL(_) => QueryLanguage::elastic_sql,
             })
             .unwrap_or(QueryLanguage::promql); // Default to promql if no inference_config
 

asap-query-engine/src/drivers/query/adapters/elastic_http.rs

@@ -1,5 +1,6 @@
 use super::config::AdapterConfig;
 use super::traits::*;
+use crate::QueryResult;
 use crate::data_model::QueryLanguage;
 use async_trait::async_trait;
 use axum::{
@@ -11,7 +12,7 @@ use axum::{
 use serde_json::{json, Value};
 use std::collections::HashMap;
 use std::sync::Arc;
-use tracing::{debug, error};
+use tracing::{debug, error, info};
 
 /// Elasticsearch HTTP protocol adapter
 pub struct ElasticHttpAdapter {
@@ -26,34 +27,30 @@ impl ElasticHttpAdapter {
 
     /// Parse Elasticsearch query from JSON body
     fn parse_elasticsearch_query(&self, body: &Bytes) -> Result<ParsedQueryRequest, AdapterError> {
-        debug!(
-            "Elasticsearch adapter: parsing query for language {:?}",
-            self.config.language
-        );
-
-        // Parse the JSON body
         let json_body: Value = serde_json::from_slice(body)
             .map_err(|e| AdapterError::ParseError(format!("Invalid JSON: {}", e)))?;
 
-        // Store the entire query as a JSON string
-        let query = serde_json::to_string(&json_body)
-            .map_err(|e| AdapterError::ParseError(format!("Failed to serialize query: {}", e)))?;
+        // Extract the SQL string from the "query" field
+        let query = match &self.config.language {
+            QueryLanguage::elastic_sql => {
+                json_body
+                    .get("query")
+                    .and_then(|v| v.as_str())
+                    .ok_or_else(|| AdapterError::MissingParameter("query".to_string()))?
+                    .to_string()
+            }
+            _ => {
+                // For QueryDSL, keep the full JSON as before
+                serde_json::to_string(&json_body)
+                    .map_err(|e| AdapterError::ParseError(format!("Failed to serialize query: {}", e)))?
+            }
+        };
 
         let time = std::time::SystemTime::now()
             .duration_since(std::time::UNIX_EPOCH)
             .unwrap_or_default()
             .as_secs_f64();
 
-        debug!(
-            "Elasticsearch adapter: parsed {} query with time={}",
-            if matches!(self.config.language, QueryLanguage::elastic_sql) {
-                "SQL"
-            } else {
-                "Query DSL"
-            },
-            time
-        );
-
         Ok(ParsedQueryRequest { query, time })
     }
 }
@@ -125,35 +122,53 @@ impl QueryRequestAdapter for ElasticHttpAdapter {
 #[async_trait]
 impl QueryResponseAdapter for ElasticHttpAdapter {
     async fn format_success_response(
-        &self,
-        _result: &QueryExecutionResult,
-    ) -> Result<Response, StatusCode> {
-        debug!("Elasticsearch adapter: formatting success response");
-
-        // For now, since we're falling back for every query,
-        // the result from the fallback will be passed through
-        // In the future, this could transform local execution results
-        // to Elasticsearch format
+    &self,
+    result: &QueryExecutionResult,
+) -> Result<Response, StatusCode> {
+        info!("SKETCH HIT: serving {} rows from precomputed sketches", 
+          match &result.query_result {
+              QueryResult::Vector(v) => v.values.len(),
+              QueryResult::Matrix(_) => 0,
+          });
+
+        let label_names = &result.query_output_labels.labels;
+
+        let hits: Vec<Value> = match &result.query_result {
+            QueryResult::Vector(instant_vector) => {
+                instant_vector
+                    .values
+                    .iter()
+                    .map(|element| {
+                        let mut source = serde_json::Map::new();
+                        for (i, label_name) in label_names.iter().enumerate() {
+                            let label_value =
+                                element.labels.get(i).map(|s| s.as_str()).unwrap_or("");
+                            source.insert(label_name.clone(), json!(label_value));
+                        }
+                        source.insert("value".to_string(), json!(element.value));
+                        json!({
+                            "_source": source
+                        })
+                    })
+                    .collect()
+            }
+            QueryResult::Matrix(_) => {
+                return Err(StatusCode::NOT_IMPLEMENTED);
+            }
+        };
 
-        // Return a stub Elasticsearch-style response for now.
         let response = json!({
             "took": 0,
             "timed_out": false,
             "hits": {
                 "total": {
-                    "value": 0,
+                    "value": hits.len(),
                     "relation": "eq"
                 },
-                "hits": []
+                "hits": hits
             }
         });
 
-        debug!(
-            "Elasticsearch adapter: returning stub response: {}",
-            serde_json::to_string_pretty(&response)
-                .unwrap_or_else(|_| "Unable to format".to_string())
-        );
-
         Ok(Json(response).into_response())
     }
 

asap-query-engine/src/drivers/query/fallback/elastic.rs

@@ -80,11 +80,23 @@ impl FallbackClient for ElasticHttpFallback {
 
         debug!("Full forwarding URL: {}", full_url);
 
-        let query_body: Value = match serde_json::from_str(&request.query) {
-            Ok(json) => json,
-            Err(e) => {
-                error!("Failed to parse query as JSON: {}", e);
-                return Err(StatusCode::INTERNAL_SERVER_ERROR);
+        let query_body: Value = match self.language {
+            QueryLanguage::elastic_sql => {
+                // query is a raw SQL string, need to wrap it for the ES SQL endpoint
+                serde_json::json!({
+                    "query": request.query.trim().trim_end_matches(';'),
+                    "fetch_size": 1000,
+                })
+            }
+            _ => {
+                // query is already a JSON string (Query DSL)
+                match serde_json::from_str(&request.query) {
+                    Ok(json) => json,
+                    Err(e) => {
+                        error!("Failed to parse query as JSON: {}", e);
+                        return Err(StatusCode::INTERNAL_SERVER_ERROR);
+                    }
+                }
             }
         };
 

asap-query-engine/src/engines/simple_engine.rs

@@ -1132,7 +1132,7 @@ impl SimpleEngine {
                 return None;
             }
             &SchemaConfig::ElasticQueryDSL => todo!(),
-            &SchemaConfig::ElasticSQL => todo!(),
+            SchemaConfig::ElasticSQL(sql_schema) => sql_schema.clone(),
         };
 
         let statements = parser::parse_sql(&GenericDialect {}, query.as_str()).unwrap();
@@ -1345,7 +1345,7 @@ impl SimpleEngine {
             .streaming_config
             .get_aggregation_config(agg_info.aggregation_id_for_key)
             .map(|config| config.aggregated_labels.clone())
-            .unwrap_or_else(KeyByLabelNames::empty);
+            .unwrap_or_else(KeyByLabelNames::empty); 
 
         Some(QueryExecutionContext {
             metric: metric.to_string(),
@@ -1473,14 +1473,14 @@ impl SimpleEngine {
             QueryLanguage::promql => self.handle_query_promql(query, time),
             QueryLanguage::sql => self.handle_query_sql(query, time),
             QueryLanguage::elastic_querydsl => self.handle_query_elastic(),
-            QueryLanguage::elastic_sql => self.handle_query_elastic(),
+            QueryLanguage::elastic_sql => self.handle_query_sql(query, time),
         }
     }
 
     pub fn handle_query_elastic(&self) -> Option<(KeyByLabelNames, QueryResult)> {
         None
     }
-
+    
     // /// Try to extract sketch query components from a PromQL query string.
     // ///
     // /// Attempts the standard AST parser first. If that fails (e.g. for custom
@@ -1919,7 +1919,7 @@ impl SimpleEngine {
                 warn!("PromQL query requested but config has ElasticQueryDSL schema");
                 return None;
             }
-            &SchemaConfig::ElasticSQL => {
+            SchemaConfig::ElasticSQL(_) => {
                 warn!("PromQL query requested but config has ElasticSQL schema");
                 return None;
             }