From 3b8661201efe893cd80da5f8dbd3b21333a8eda9 Mon Sep 17 00:00:00 2001
From: Chase <62891993+engechas@users.noreply.github.com>
Date: Mon, 28 Oct 2024 14:34:49 -0700
Subject: [PATCH] Add validation method for Flint extension queries and wire it
 into the dispatcher (#3096)

* Add validation method for Flint extension queries and wire it into the dispatcher

Signed-off-by: Chase Engelbrecht <engechas@amazon.com>

* Add unit test

Signed-off-by: Chase Engelbrecht <engechas@amazon.com>

* Run spotless

Signed-off-by: Chase Engelbrecht <engechas@amazon.com>

* Fix NPE in test

Signed-off-by: Chase Engelbrecht <engechas@amazon.com>

* Add java doc

Signed-off-by: Chase Engelbrecht <engechas@amazon.com>

---------

Signed-off-by: Chase Engelbrecht <engechas@amazon.com>
---
 .../sql/spark/dispatcher/SparkQueryDispatcher.java     |  1 +
 .../sql/spark/validator/SQLQueryValidator.java         |  8 ++++++++
 .../sql/spark/dispatcher/SparkQueryDispatcherTest.java |  5 +++++
 .../sql/spark/validator/SQLQueryValidatorTest.java     | 10 ++++++++++
 4 files changed, 24 insertions(+)

diff --git a/async-query-core/src/main/java/org/opensearch/sql/spark/dispatcher/SparkQueryDispatcher.java b/async-query-core/src/main/java/org/opensearch/sql/spark/dispatcher/SparkQueryDispatcher.java
index ff8c8d1fe8..5f88ea9ca0 100644
--- a/async-query-core/src/main/java/org/opensearch/sql/spark/dispatcher/SparkQueryDispatcher.java
+++ b/async-query-core/src/main/java/org/opensearch/sql/spark/dispatcher/SparkQueryDispatcher.java
@@ -51,6 +51,7 @@ public DispatchQueryResponse dispatch(
       String query = dispatchQueryRequest.getQuery();
 
       if (SQLQueryUtils.isFlintExtensionQuery(query)) {
+        sqlQueryValidator.validateFlintExtensionQuery(query, dataSourceMetadata.getConnector());
         return handleFlintExtensionQuery(
             dispatchQueryRequest, asyncQueryRequestContext, dataSourceMetadata);
       }
diff --git a/async-query-core/src/main/java/org/opensearch/sql/spark/validator/SQLQueryValidator.java b/async-query-core/src/main/java/org/opensearch/sql/spark/validator/SQLQueryValidator.java
index f387cbad25..5d934411bf 100644
--- a/async-query-core/src/main/java/org/opensearch/sql/spark/validator/SQLQueryValidator.java
+++ b/async-query-core/src/main/java/org/opensearch/sql/spark/validator/SQLQueryValidator.java
@@ -36,4 +36,12 @@ public void validate(String sqlQuery, DataSourceType datasourceType) {
       throw e;
     }
   }
+
+  /**
+   * Validates a query from the Flint extension grammar. The method is currently a no-op.
+   *
+   * @param sqlQuery The Flint extension query to be validated
+   * @param dataSourceType The type of the datasource the query is being run on
+   */
+  public void validateFlintExtensionQuery(String sqlQuery, DataSourceType dataSourceType) {}
 }
diff --git a/async-query-core/src/test/java/org/opensearch/sql/spark/dispatcher/SparkQueryDispatcherTest.java b/async-query-core/src/test/java/org/opensearch/sql/spark/dispatcher/SparkQueryDispatcherTest.java
index 343dc3040a..405fdf511d 100644
--- a/async-query-core/src/test/java/org/opensearch/sql/spark/dispatcher/SparkQueryDispatcherTest.java
+++ b/async-query-core/src/test/java/org/opensearch/sql/spark/dispatcher/SparkQueryDispatcherTest.java
@@ -641,6 +641,11 @@ void testDispatchVacuumIndexQuery() {
 
   @Test
   void testDispatchRecoverIndexQuery() {
+    DataSourceMetadata dataSourceMetadata = constructMyGlueDataSourceMetadata();
+    when(dataSourceService.verifyDataSourceAccessAndGetRawMetadata(
+            MY_GLUE, asyncQueryRequestContext))
+        .thenReturn(dataSourceMetadata);
+
     String query = "RECOVER INDEX JOB `flint_spark_catalog_default_test_skipping_index`";
     Assertions.assertThrows(
         IllegalArgumentException.class,
diff --git a/async-query-core/src/test/java/org/opensearch/sql/spark/validator/SQLQueryValidatorTest.java b/async-query-core/src/test/java/org/opensearch/sql/spark/validator/SQLQueryValidatorTest.java
index 695a083809..fbc00109da 100644
--- a/async-query-core/src/test/java/org/opensearch/sql/spark/validator/SQLQueryValidatorTest.java
+++ b/async-query-core/src/test/java/org/opensearch/sql/spark/validator/SQLQueryValidatorTest.java
@@ -5,11 +5,13 @@
 
 package org.opensearch.sql.spark.validator;
 
+import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.when;
 
 import java.util.Arrays;
+import java.util.UUID;
 import lombok.AllArgsConstructor;
 import lombok.Getter;
 import org.antlr.v4.runtime.CommonTokenStream;
@@ -561,6 +563,14 @@ void testSecurityLakeQueries() {
     v.ng(TestElement.INTEGRATION_WITH_HIVE_UDFS_UDAFS_UDTFS);
   }
 
+  @Test
+  void testValidateFlintExtensionQuery() {
+    assertDoesNotThrow(
+        () ->
+            sqlQueryValidator.validateFlintExtensionQuery(
+                UUID.randomUUID().toString(), DataSourceType.SECURITY_LAKE));
+  }
+
   @AllArgsConstructor
   private static class VerifyValidator {
     private final SQLQueryValidator validator;