Add validation method for Flint extension queries and wire it into th…

…e 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>
opensearch-project · Oct 28, 2024 · 3b86612 · 3b86612
1 parent f00244e
commit 3b86612
Show file tree

Hide file tree

Showing 4 changed files with 24 additions and 0 deletions.
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
@@ -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
@@ -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/...uery-core/src/test/java/org/opensearch/sql/spark/dispatcher/SparkQueryDispatcherTest.java b/...uery-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
@@ -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;