From 3649d1dc05e3833fd0c9f21f4099b58a429616ac Mon Sep 17 00:00:00 2001
From: Stratoula Kalafateli <efstratia.kalafateli@elastic.co>
Date: Mon, 22 Apr 2024 17:09:09 +0100
Subject: [PATCH] [Obs AI assistant] Relax validation on missing fallbacks
 (#181324)

## Summary

Relaxes the validation which means that we can't quick fix the quotes
problem but we are still fixing other syntax errors such as functions
misspell etc.

(cherry picked from commit 0643f63d77aa0fd43c8bc380f72f4d9e53874768)
---
 .../query/correct_query_with_actions.test.ts  | 46 ++++---------------
 .../query/correct_query_with_actions.ts       |  4 +-
 2 files changed, 12 insertions(+), 38 deletions(-)

diff --git a/x-pack/plugins/observability_solution/observability_ai_assistant_app/server/functions/query/correct_query_with_actions.test.ts b/x-pack/plugins/observability_solution/observability_ai_assistant_app/server/functions/query/correct_query_with_actions.test.ts
index c2cf207a6925d..818f4854e038d 100644
--- a/x-pack/plugins/observability_solution/observability_ai_assistant_app/server/functions/query/correct_query_with_actions.test.ts
+++ b/x-pack/plugins/observability_solution/observability_ai_assistant_app/server/functions/query/correct_query_with_actions.test.ts
@@ -33,42 +33,14 @@ describe('correctQueryWithActions', () => {
     expect(fixedQuery).toBe('from logstash-* | stats var0 = max(bytes) | eval abs(var0) | limit 1');
   });
 
-  it(`fixes errors correctly for a query with missing quotes`, async () => {
-    const fixedQuery = await correctQueryWithActions('from logstash-* | keep field-1');
-    expect(fixedQuery).toBe('from logstash-* | keep `field-1`');
-  });
-
-  it(`fixes errors correctly for a query with missing quotes in multiple commands`, async () => {
-    const fixedQuery = await correctQueryWithActions(
-      'from logstash-* | stats avg(field-1) | eval abs(field-2)'
-    );
-    expect(fixedQuery).toBe('from logstash-* | stats avg(`field-1`) | eval abs(`field-2`)');
-  });
-
-  it(`fixes errors correctly for a query with missing quotes and keep with multiple fields`, async () => {
-    const fixedQuery = await correctQueryWithActions('from logstash-* | keep field-1, field-2');
-    expect(fixedQuery).toBe('from logstash-* | keep `field-1`, `field-2`');
-  });
-
-  it(`fixes errors correctly for a query with missing quotes with variable assignment`, async () => {
-    const fixedQuery = await correctQueryWithActions('from logstash-* | stats var1 = avg(field-1)');
-    expect(fixedQuery).toBe('from logstash-* | stats var1 = avg(`field-1`)');
-  });
-
-  it(`fixes errors correctly for a query with missing quotes in an aggregation`, async () => {
-    const fixedQuery = await correctQueryWithActions('from logstash-* | stats avg(field-1)');
-    expect(fixedQuery).toBe('from logstash-* | stats avg(`field-1`)');
-  });
-
-  it(`fixes errors correctly for a query with typo on stats and wrong quotes`, async () => {
-    const fixedQuery = await correctQueryWithActions('from logstash-* | stats aveg(field-1)');
-    expect(fixedQuery).toBe('from logstash-* | stats avg(`field-1`)');
-  });
-
-  it(`fixes errors correctly for a query with missing quotes on stats and keep`, async () => {
-    const fixedQuery = await correctQueryWithActions(
-      'from logstash-* | stats avg(field-1) | keep field-2'
-    );
-    expect(fixedQuery).toBe('from logstash-* | stats avg(`field-1`) | keep `field-2`');
+  it(`doesnt complain for @timestamp column`, async () => {
+    const queryWithTimestamp = `FROM logstash-*
+    | WHERE @timestamp >= NOW() - 15 minutes
+    | EVAL bucket = DATE_TRUNC(1 minute, @timestamp)
+    | STATS avg_cpu = AVG(system.cpu.total.norm.pct) BY service.name, bucket
+    | SORT avg_cpu DESC
+    | LIMIT 10`;
+    const fixedQuery = await correctQueryWithActions(queryWithTimestamp);
+    expect(fixedQuery).toBe(queryWithTimestamp);
   });
 });
diff --git a/x-pack/plugins/observability_solution/observability_ai_assistant_app/server/functions/query/correct_query_with_actions.ts b/x-pack/plugins/observability_solution/observability_ai_assistant_app/server/functions/query/correct_query_with_actions.ts
index bc684d13841eb..213b7e967970a 100644
--- a/x-pack/plugins/observability_solution/observability_ai_assistant_app/server/functions/query/correct_query_with_actions.ts
+++ b/x-pack/plugins/observability_solution/observability_ai_assistant_app/server/functions/query/correct_query_with_actions.ts
@@ -8,7 +8,9 @@ import { validateQuery, getActions } from '@kbn/esql-validation-autocomplete';
 import { getAstAndSyntaxErrors } from '@kbn/esql-ast';
 
 const fixedQueryByOneAction = async (queryString: string) => {
-  const { errors } = await validateQuery(queryString, getAstAndSyntaxErrors);
+  const { errors } = await validateQuery(queryString, getAstAndSyntaxErrors, {
+    ignoreOnMissingCallbacks: true,
+  });
 
   const actions = await getActions(queryString, errors, getAstAndSyntaxErrors, {
     relaxOnMissingCallbacks: true,