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 <[email protected]>
opensearch-project · Oct 5, 2023 · 10e2d07 · 10e2d07
1 parent 5df6105
commit 10e2d07
Show file tree

Hide file tree

Showing 33 changed files with 291 additions and 136 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(false)).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
@@ -66,7 +66,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());
@@ -96,7 +97,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.
@@ -108,7 +110,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());
@@ -141,7 +144,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.
@@ -179,7 +183,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.
@@ -215,7 +220,8 @@ public void issue2196() {
                 "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());

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;
@@ -74,11 +76,13 @@ public CreateAsyncQueryResponse createAsyncQuery(
                 createAsyncQueryRequest.getLang(),
                 sparkExecutionEngineConfig.getExecutionRoleARN(),
                 clusterName.value()));
+
     asyncQueryJobMetadataStorageService.storeJobMetadata(
         new AsyncQueryJobMetadata(
             sparkExecutionEngineConfig.getApplicationId(),
             dispatchQueryResponse.getJobId(),
-            dispatchQueryResponse.isDropIndexQuery()));
+            dispatchQueryResponse.isDropIndexQuery(),
+            dispatchQueryResponse.getResultIndex()));
     return new CreateAsyncQueryResponse(dispatchQueryResponse.getJobId());
   }
 
@@ -89,17 +93,21 @@ public AsyncQueryExecutionResponse getAsyncQueryResults(String queryId) {
         asyncQueryJobMetadataStorageService.getJobMetadata(queryId);
     if (jobMetadata.isPresent()) {
       JSONObject jsonObject = sparkQueryDispatcher.getQueryResponse(jobMetadata.get());
-      if (JobRunState.SUCCESS.toString().equals(jsonObject.getString("status"))) {
+      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
@@ -29,11 +29,13 @@ public class AsyncQueryJobMetadata {
   private String applicationId;
   private String jobId;
   private boolean isDropIndexQuery;
+  private String resultIndex;
 
-  public AsyncQueryJobMetadata(String applicationId, String jobId) {
+  public AsyncQueryJobMetadata(String applicationId, String jobId, String resultIndex) {
     this.applicationId = applicationId;
     this.jobId = jobId;
     this.isDropIndexQuery = false;
+    this.resultIndex = resultIndex;
   }
 
   @Override
@@ -54,6 +56,7 @@ public static XContentBuilder convertToXContent(AsyncQueryJobMetadata metadata)
     builder.field("jobId", metadata.getJobId());
     builder.field("applicationId", metadata.getApplicationId());
     builder.field("isDropIndexQuery", metadata.isDropIndexQuery());
+    builder.field("resultIndex", metadata.getResultIndex());
     builder.endObject();
     return builder;
   }
@@ -88,6 +91,7 @@ public static AsyncQueryJobMetadata toJobMetadata(XContentParser parser) throws
     String jobId = null;
     String applicationId = null;
     boolean isDropIndexQuery = false;
+    String resultIndex = null;
     ensureExpectedToken(XContentParser.Token.START_OBJECT, parser.nextToken(), parser);
     while (parser.nextToken() != XContentParser.Token.END_OBJECT) {
       String fieldName = parser.currentName();
@@ -102,13 +106,16 @@ public static AsyncQueryJobMetadata toJobMetadata(XContentParser parser) throws
         case "isDropIndexQuery":
           isDropIndexQuery = parser.booleanValue();
           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(applicationId, jobId, isDropIndexQuery);
+    return new AsyncQueryJobMetadata(applicationId, jobId, isDropIndexQuery, 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;
   }
 }