Add customized result index in data source etc

This PR - Introduce `spark.flint.datasource.name` parameter for data source specification. - Enhance data source creation to allow custom result indices; fallback to default if unavailable. - Include error details in the async result response, sourced from the result index. - Migrate to `org.apache.spark.sql.FlintJob` following updates in OpenSearch-Spark. - Populate query status from result index over EMR-S job status to handle edge cases where jobs may succeed, but queries or mappings fail. Testing done: 1. manual testing including if with or without custom result index async query still works 2. added new unit tests Signed-off-by: Kaituo Li <kaituo@amazon.com>
opensearch-project · Oct 4, 2023 · cc6b5a6 · cc6b5a6
1 parent 9a2151b
commit cc6b5a6
Show file tree

Hide file tree

Showing 30 changed files with 287 additions and 94 deletions.
diff --git a/core/src/main/java/org/opensearch/sql/datasource/model/DataSourceMetadata.java b/core/src/main/java/org/opensearch/sql/datasource/model/DataSourceMetadata.java
@@ -42,16 +42,20 @@ public class DataSourceMetadata {
 
   @JsonProperty private Map<String, String> properties;
 
+  @JsonProperty private String resultIndex;
+
   public DataSourceMetadata(
       String name,
       DataSourceType connector,
       List<String> allowedRoles,
-      Map<String, String> properties) {
+      Map<String, String> properties,
+      String resultIndex) {
     this.name = name;
     this.connector = connector;
     this.description = StringUtils.EMPTY;
     this.properties = properties;
     this.allowedRoles = allowedRoles;
+    this.resultIndex = resultIndex;
   }
 
   public DataSourceMetadata() {
@@ -69,6 +73,7 @@ public static DataSourceMetadata defaultOpenSearchDataSourceMetadata() {
         DEFAULT_DATASOURCE_NAME,
         DataSourceType.OPENSEARCH,
         Collections.emptyList(),
-        ImmutableMap.of());
+        ImmutableMap.of(),
+        null);
   }
 }
diff --git a/core/src/test/java/org/opensearch/sql/analysis/AnalyzerTestBase.java b/core/src/test/java/org/opensearch/sql/analysis/AnalyzerTestBase.java
@@ -199,7 +199,8 @@ public Set<DataSourceMetadata> getDataSourceMetadata(boolean isDefaultDataSource
                       ds.getName(),
                       ds.getConnectorType(),
                       Collections.emptyList(),
-                      ImmutableMap.of()))
+                      ImmutableMap.of(),
+                      null))
           .collect(Collectors.toSet());
     }
 

diff --git a/...src/test/java/org/opensearch/sql/planner/physical/datasource/DataSourceTableScanTest.java b/...src/test/java/org/opensearch/sql/planner/physical/datasource/DataSourceTableScanTest.java
@@ -64,7 +64,8 @@ void testIterator() {
                         dataSource.getName(),
                         dataSource.getConnectorType(),
                         Collections.emptyList(),
-                        ImmutableMap.of()))
+                        ImmutableMap.of(),
+                        null))
             .collect(Collectors.toSet());
     when(dataSourceService.getDataSourceMetadata(true)).thenReturn(dataSourceMetadata);
 

diff --git a/datasources/src/main/java/org/opensearch/sql/datasources/utils/XContentParserUtils.java b/datasources/src/main/java/org/opensearch/sql/datasources/utils/XContentParserUtils.java
@@ -32,6 +32,8 @@ public class XContentParserUtils {
   public static final String PROPERTIES_FIELD = "properties";
   public static final String ALLOWED_ROLES_FIELD = "allowedRoles";
 
+  public static final String RESULT_INDEX_FIELD = "resultIndex";
+
   /**
    * Convert xcontent parser to DataSourceMetadata.
    *
@@ -45,6 +47,7 @@ public static DataSourceMetadata toDataSourceMetadata(XContentParser parser) thr
     DataSourceType connector = null;
     List<String> allowedRoles = new ArrayList<>();
     Map<String, String> properties = new HashMap<>();
+    String resultIndex = null;
     ensureExpectedToken(XContentParser.Token.START_OBJECT, parser.nextToken(), parser);
     while (parser.nextToken() != XContentParser.Token.END_OBJECT) {
       String fieldName = parser.currentName();
@@ -73,14 +76,18 @@ public static DataSourceMetadata toDataSourceMetadata(XContentParser parser) thr
             properties.put(key, value);
           }
           break;
+        case RESULT_INDEX_FIELD:
+          resultIndex = parser.textOrNull();
+          break;
         default:
           throw new IllegalArgumentException("Unknown field: " + fieldName);
       }
     }
     if (name == null || connector == null) {
       throw new IllegalArgumentException("name and connector are required fields.");
     }
-    return new DataSourceMetadata(name, description, connector, allowedRoles, properties);
+    return new DataSourceMetadata(
+        name, description, connector, allowedRoles, properties, resultIndex);
   }
 
   /**
@@ -122,6 +129,7 @@ public static XContentBuilder convertToXContent(DataSourceMetadata metadata) thr
       builder.field(entry.getKey(), entry.getValue());
     }
     builder.endObject();
+    builder.field(RESULT_INDEX_FIELD, metadata.getResultIndex());
     builder.endObject();
     return builder;
   }

diff --git a/datasources/src/main/resources/datasources-index-mapping.yml b/datasources/src/main/resources/datasources-index-mapping.yml
@@ -14,4 +14,6 @@ properties:
       keyword:
         type: keyword
   connector:
+    type: keyword
+  resultIndex:
     type: keyword
diff --git a/...urces/src/test/java/org/opensearch/sql/datasources/service/DataSourceServiceImplTest.java b/...urces/src/test/java/org/opensearch/sql/datasources/service/DataSourceServiceImplTest.java
@@ -344,7 +344,8 @@ void testRemovalOfAuthorizationInfo() {
             "testDS",
             DataSourceType.PROMETHEUS,
             Collections.singletonList("prometheus_access"),
-            properties);
+            properties,
+            null);
     when(dataSourceMetadataStorage.getDataSourceMetadata("testDS"))
         .thenReturn(Optional.of(dataSourceMetadata));
 
@@ -398,7 +399,8 @@ void testGetRawDataSourceMetadata() {
             "testDS",
             DataSourceType.PROMETHEUS,
             Collections.singletonList("prometheus_access"),
-            properties);
+            properties,
+            null);
     when(dataSourceMetadataStorage.getDataSourceMetadata("testDS"))
         .thenReturn(Optional.of(dataSourceMetadata));
 

diff --git a/datasources/src/test/java/org/opensearch/sql/datasources/utils/XContentParserUtilsTest.java b/datasources/src/test/java/org/opensearch/sql/datasources/utils/XContentParserUtilsTest.java
@@ -31,7 +31,7 @@ public void testConvertToXContent() {
     XContentBuilder contentBuilder = XContentParserUtils.convertToXContent(dataSourceMetadata);
     String contentString = BytesReference.bytes(contentBuilder).utf8ToString();
     Assertions.assertEquals(
-        "{\"name\":\"testDS\",\"description\":\"\",\"connector\":\"PROMETHEUS\",\"allowedRoles\":[\"prometheus_access\"],\"properties\":{\"prometheus.uri\":\"https://localhost:9090\"}}",
+        "{\"name\":\"testDS\",\"description\":\"\",\"connector\":\"PROMETHEUS\",\"allowedRoles\":[\"prometheus_access\"],\"properties\":{\"prometheus.uri\":\"https://localhost:9090\"},\"resultIndex\":null}",
         contentString);
   }
 

diff --git a/integ-test/src/test/java/org/opensearch/sql/datasource/DataSourceAPIsIT.java b/integ-test/src/test/java/org/opensearch/sql/datasource/DataSourceAPIsIT.java
@@ -62,7 +62,8 @@ public void createDataSourceAPITest() {
                 "prometheus.auth.username",
                 "username",
                 "prometheus.auth.password",
-                "password"));
+                "password"),
+            null);
     Request createRequest = getCreateDataSourceRequest(createDSM);
     Response response = client().performRequest(createRequest);
     Assert.assertEquals(201, response.getStatusLine().getStatusCode());
@@ -92,7 +93,8 @@ public void updateDataSourceAPITest() {
             "update_prometheus",
             DataSourceType.PROMETHEUS,
             ImmutableList.of(),
-            ImmutableMap.of("prometheus.uri", "https://localhost:9090"));
+            ImmutableMap.of("prometheus.uri", "https://localhost:9090"),
+            null);
     Request createRequest = getCreateDataSourceRequest(createDSM);
     client().performRequest(createRequest);
     // Datasource is not immediately created. so introducing a sleep of 2s.
@@ -104,7 +106,8 @@ public void updateDataSourceAPITest() {
             "update_prometheus",
             DataSourceType.PROMETHEUS,
             ImmutableList.of(),
-            ImmutableMap.of("prometheus.uri", "https://randomtest.com:9090"));
+            ImmutableMap.of("prometheus.uri", "https://randomtest.com:9090"),
+            null);
     Request updateRequest = getUpdateDataSourceRequest(updateDSM);
     Response updateResponse = client().performRequest(updateRequest);
     Assert.assertEquals(200, updateResponse.getStatusLine().getStatusCode());
@@ -137,7 +140,8 @@ public void deleteDataSourceTest() {
             "delete_prometheus",
             DataSourceType.PROMETHEUS,
             ImmutableList.of(),
-            ImmutableMap.of("prometheus.uri", "https://localhost:9090"));
+            ImmutableMap.of("prometheus.uri", "https://localhost:9090"),
+            null);
     Request createRequest = getCreateDataSourceRequest(createDSM);
     client().performRequest(createRequest);
     // Datasource is not immediately created. so introducing a sleep of 2s.
@@ -175,7 +179,8 @@ public void getAllDataSourceTest() {
             "get_all_prometheus",
             DataSourceType.PROMETHEUS,
             ImmutableList.of(),
-            ImmutableMap.of("prometheus.uri", "https://localhost:9090"));
+            ImmutableMap.of("prometheus.uri", "https://localhost:9090"),
+            null);
     Request createRequest = getCreateDataSourceRequest(createDSM);
     client().performRequest(createRequest);
     // Datasource is not immediately created. so introducing a sleep of 2s.

diff --git a/integ-test/src/test/java/org/opensearch/sql/ppl/InformationSchemaCommandIT.java b/integ-test/src/test/java/org/opensearch/sql/ppl/InformationSchemaCommandIT.java
@@ -46,7 +46,8 @@ protected void init() throws InterruptedException, IOException {
             "my_prometheus",
             DataSourceType.PROMETHEUS,
             ImmutableList.of(),
-            ImmutableMap.of("prometheus.uri", "http://localhost:9090"));
+            ImmutableMap.of("prometheus.uri", "http://localhost:9090"),
+            null);
     Request createRequest = getCreateDataSourceRequest(createDSM);
     Response response = client().performRequest(createRequest);
     Assert.assertEquals(201, response.getStatusLine().getStatusCode());

diff --git a/integ-test/src/test/java/org/opensearch/sql/ppl/PrometheusDataSourceCommandsIT.java b/integ-test/src/test/java/org/opensearch/sql/ppl/PrometheusDataSourceCommandsIT.java
@@ -58,7 +58,8 @@ protected void init() throws InterruptedException, IOException {
             "my_prometheus",
             DataSourceType.PROMETHEUS,
             ImmutableList.of(),
-            ImmutableMap.of("prometheus.uri", "http://localhost:9090"));
+            ImmutableMap.of("prometheus.uri", "http://localhost:9090"),
+            null);
     Request createRequest = getCreateDataSourceRequest(createDSM);
     Response response = client().performRequest(createRequest);
     Assert.assertEquals(201, response.getStatusLine().getStatusCode());

diff --git a/integ-test/src/test/java/org/opensearch/sql/ppl/ShowDataSourcesCommandIT.java b/integ-test/src/test/java/org/opensearch/sql/ppl/ShowDataSourcesCommandIT.java
@@ -46,7 +46,8 @@ protected void init() throws InterruptedException, IOException {
             "my_prometheus",
             DataSourceType.PROMETHEUS,
             ImmutableList.of(),
-            ImmutableMap.of("prometheus.uri", "http://localhost:9090"));
+            ImmutableMap.of("prometheus.uri", "http://localhost:9090"),
+            null);
     Request createRequest = getCreateDataSourceRequest(createDSM);
     Response response = client().performRequest(createRequest);
     Assert.assertEquals(201, response.getStatusLine().getStatusCode());

diff --git a/spark/src/main/java/org/opensearch/sql/spark/asyncquery/AsyncQueryExecutorServiceImpl.java b/spark/src/main/java/org/opensearch/sql/spark/asyncquery/AsyncQueryExecutorServiceImpl.java
@@ -7,6 +7,8 @@
 
 import static org.opensearch.sql.common.setting.Settings.Key.CLUSTER_NAME;
 import static org.opensearch.sql.common.setting.Settings.Key.SPARK_EXECUTION_ENGINE_CONFIG;
+import static org.opensearch.sql.spark.data.constants.SparkConstants.ERROR_FIELD;
+import static org.opensearch.sql.spark.data.constants.SparkConstants.STATUS_FIELD;
 
 import com.amazonaws.services.emrserverless.model.JobRunState;
 import java.security.AccessController;
@@ -15,6 +17,7 @@
 import java.util.List;
 import java.util.Optional;
 import lombok.AllArgsConstructor;
+import org.apache.commons.lang3.tuple.Pair;
 import org.json.JSONObject;
 import org.opensearch.cluster.ClusterName;
 import org.opensearch.sql.common.setting.Settings;
@@ -64,17 +67,22 @@ public CreateAsyncQueryResponse createAsyncQuery(
                     SparkExecutionEngineConfig.toSparkExecutionEngineConfig(
                         sparkExecutionEngineConfigString));
     ClusterName clusterName = settings.getSettingValue(CLUSTER_NAME);
-    String jobId =
+    String query = createAsyncQueryRequest.getQuery();
+    Pair<String, String> jobIdResultIndexPair =
         sparkQueryDispatcher.dispatch(
             new DispatchQueryRequest(
                 sparkExecutionEngineConfig.getApplicationId(),
-                createAsyncQueryRequest.getQuery(),
+                query,
                 createAsyncQueryRequest.getLang(),
                 sparkExecutionEngineConfig.getExecutionRoleARN(),
                 clusterName.value()));
+
     asyncQueryJobMetadataStorageService.storeJobMetadata(
-        new AsyncQueryJobMetadata(jobId, sparkExecutionEngineConfig.getApplicationId()));
-    return new CreateAsyncQueryResponse(jobId);
+        new AsyncQueryJobMetadata(
+            jobIdResultIndexPair.getLeft(),
+            sparkExecutionEngineConfig.getApplicationId(),
+            jobIdResultIndexPair.getRight()));
+    return new CreateAsyncQueryResponse(jobIdResultIndexPair.getLeft());
   }
 
   @Override
@@ -85,18 +93,24 @@ public AsyncQueryExecutionResponse getAsyncQueryResults(String queryId) {
     if (jobMetadata.isPresent()) {
       JSONObject jsonObject =
           sparkQueryDispatcher.getQueryResponse(
-              jobMetadata.get().getApplicationId(), jobMetadata.get().getJobId());
-      if (JobRunState.SUCCESS.toString().equals(jsonObject.getString("status"))) {
+              jobMetadata.get().getApplicationId(),
+              jobMetadata.get().getJobId(),
+              jobMetadata.get().getResultIndex());
+      if (JobRunState.SUCCESS.toString().equals(jsonObject.getString(STATUS_FIELD))) {
         DefaultSparkSqlFunctionResponseHandle sparkSqlFunctionResponseHandle =
             new DefaultSparkSqlFunctionResponseHandle(jsonObject);
         List<ExprValue> result = new ArrayList<>();
         while (sparkSqlFunctionResponseHandle.hasNext()) {
           result.add(sparkSqlFunctionResponseHandle.next());
         }
         return new AsyncQueryExecutionResponse(
-            JobRunState.SUCCESS.toString(), sparkSqlFunctionResponseHandle.schema(), result);
+            JobRunState.SUCCESS.toString(), sparkSqlFunctionResponseHandle.schema(), result, null);
       } else {
-        return new AsyncQueryExecutionResponse(jsonObject.getString("status"), null, null);
+        return new AsyncQueryExecutionResponse(
+            jsonObject.optString(STATUS_FIELD, JobRunState.FAILED.toString()),
+            null,
+            null,
+            jsonObject.optString(ERROR_FIELD, ""));
       }
     }
     throw new AsyncQueryNotFoundException(String.format("QueryId: %s not found", queryId));

diff --git a/.../src/main/java/org/opensearch/sql/spark/asyncquery/model/AsyncQueryExecutionResponse.java b/.../src/main/java/org/opensearch/sql/spark/asyncquery/model/AsyncQueryExecutionResponse.java
@@ -18,4 +18,5 @@ public class AsyncQueryExecutionResponse {
   private final String status;
   private final ExecutionEngine.Schema schema;
   private final List<ExprValue> results;
+  private final String error;
 }
diff --git a/spark/src/main/java/org/opensearch/sql/spark/asyncquery/model/AsyncQueryJobMetadata.java b/spark/src/main/java/org/opensearch/sql/spark/asyncquery/model/AsyncQueryJobMetadata.java
@@ -26,6 +26,7 @@
 public class AsyncQueryJobMetadata {
   private String jobId;
   private String applicationId;
+  private String resultIndex;
 
   @Override
   public String toString() {
@@ -44,6 +45,7 @@ public static XContentBuilder convertToXContent(AsyncQueryJobMetadata metadata)
     builder.startObject();
     builder.field("jobId", metadata.getJobId());
     builder.field("applicationId", metadata.getApplicationId());
+    builder.field("resultIndex", metadata.getResultIndex());
     builder.endObject();
     return builder;
   }
@@ -77,6 +79,7 @@ public static AsyncQueryJobMetadata toJobMetadata(String json) throws IOExceptio
   public static AsyncQueryJobMetadata toJobMetadata(XContentParser parser) throws IOException {
     String jobId = null;
     String applicationId = null;
+    String resultIndex = null;
     ensureExpectedToken(XContentParser.Token.START_OBJECT, parser.nextToken(), parser);
     while (parser.nextToken() != XContentParser.Token.END_OBJECT) {
       String fieldName = parser.currentName();
@@ -88,13 +91,16 @@ public static AsyncQueryJobMetadata toJobMetadata(XContentParser parser) throws
         case "applicationId":
           applicationId = parser.textOrNull();
           break;
+        case "resultIndex":
+          resultIndex = parser.textOrNull();
+          break;
         default:
           throw new IllegalArgumentException("Unknown field: " + fieldName);
       }
     }
     if (jobId == null || applicationId == null) {
       throw new IllegalArgumentException("jobId and applicationId are required fields.");
     }
-    return new AsyncQueryJobMetadata(jobId, applicationId);
+    return new AsyncQueryJobMetadata(jobId, applicationId, resultIndex);
   }
 }
diff --git a/spark/src/main/java/org/opensearch/sql/spark/asyncquery/model/AsyncQueryResult.java b/spark/src/main/java/org/opensearch/sql/spark/asyncquery/model/AsyncQueryResult.java
@@ -11,19 +11,26 @@
 public class AsyncQueryResult extends QueryResult {
 
   @Getter private final String status;
+  @Getter private final String error;
 
   public AsyncQueryResult(
       String status,
       ExecutionEngine.Schema schema,
       Collection<ExprValue> exprValues,
-      Cursor cursor) {
+      Cursor cursor,
+      String error) {
     super(schema, exprValues, cursor);
     this.status = status;
+    this.error = error;
   }
 
   public AsyncQueryResult(
-      String status, ExecutionEngine.Schema schema, Collection<ExprValue> exprValues) {
+      String status,
+      ExecutionEngine.Schema schema,
+      Collection<ExprValue> exprValues,
+      String error) {
     super(schema, exprValues);
     this.status = status;
+    this.error = error;
   }
 }
diff --git a/spark/src/main/java/org/opensearch/sql/spark/client/EmrServerlessClientImplEMR.java b/spark/src/main/java/org/opensearch/sql/spark/client/EmrServerlessClientImplEMR.java
@@ -34,6 +34,10 @@ public EmrServerlessClientImplEMR(AWSEMRServerless emrServerless) {
 
   @Override
   public String startJobRun(StartJobRequest startJobRequest) {
+    String resultIndex =
+        startJobRequest.getResultIndex() == null
+            ? SPARK_RESPONSE_BUFFER_INDEX_NAME
+            : startJobRequest.getResultIndex();
     StartJobRunRequest request =
         new StartJobRunRequest()
             .withName(startJobRequest.getJobName())
@@ -45,8 +49,7 @@ public String startJobRun(StartJobRequest startJobRequest) {
                     .withSparkSubmit(
                         new SparkSubmit()
                             .withEntryPoint(SPARK_SQL_APPLICATION_JAR)
-                            .withEntryPointArguments(
-                                startJobRequest.getQuery(), SPARK_RESPONSE_BUFFER_INDEX_NAME)
+                            .withEntryPointArguments(startJobRequest.getQuery(), resultIndex)
                             .withSparkSubmitParameters(startJobRequest.getSparkSubmitParams())));
     StartJobRunResult startJobRunResult =
         AccessController.doPrivileged(

diff --git a/spark/src/main/java/org/opensearch/sql/spark/client/StartJobRequest.java b/spark/src/main/java/org/opensearch/sql/spark/client/StartJobRequest.java
@@ -20,4 +20,5 @@ public class StartJobRequest {
   private final String executionRoleArn;
   private final String sparkSubmitParams;
   private final Map<String, String> tags;
+  private final String resultIndex;
 }