[Experimental] Support direct query data sources (#4375)

Co-authored-by: Megha Goyal <[email protected]>
Co-authored-by: Aaron Alvarez <[email protected]>
Co-authored-by: Kai Huang <[email protected]>
Co-authored-by: Songkan Tang <[email protected]>
Co-authored-by: Simeon Widdis <[email protected]>
Co-authored-by: Chen Dai <[email protected]>
Co-authored-by: Jialiang Liang <[email protected]>
Co-authored-by: ritvibhatt <[email protected]>
Co-authored-by: Aaron Alvarez <[email protected]>
Co-authored-by: Vamsi Manohar <[email protected]>
Co-authored-by: Tomoyuki MORITA <[email protected]>
Co-authored-by: Lantao Jin <[email protected]>
Co-authored-by: qianheng <[email protected]>
Co-authored-by: Shenoy Pratik <[email protected]>
Co-authored-by: Yuanchun Shen <[email protected]>
Co-authored-by: Peng Huo <[email protected]>
Co-authored-by: Xinyuan Lu <[email protected]>
Co-authored-by: Joshua Li <[email protected]>

Author: 19 people

async-query-core/src/main/java/org/opensearch/sql/spark/rest/model/LangType.java

@@ -8,7 +8,8 @@
 /** Language type accepted in async query apis. */
 public enum LangType {
   SQL("sql"),
-  PPL("ppl");
+  PPL("ppl"),
+  PROMQL("promql");
   private final String text;
 
   LangType(String text) {

direct-query-core/build.gradle

@@ -0,0 +1,105 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+plugins {
+    id 'java-library'
+    id "io.freefair.lombok"
+    id 'jacoco'
+    id 'java-test-fixtures'
+}
+
+repositories {
+    mavenCentral()
+}
+
+dependencies {
+    api project(':core')
+    implementation project(':datasources')
+    implementation project(':async-query-core')
+    
+    // Common dependencies
+    implementation group: 'org.opensearch', name: 'opensearch', version: "${opensearch_version}"
+    implementation group: 'org.json', name: 'json', version: '20231013'
+    implementation group: 'commons-io', name: 'commons-io', version: "${commons_io_version}"
+
+    // Test dependencies
+    testImplementation(platform("org.junit:junit-bom:5.9.3"))
+    testImplementation 'org.junit.jupiter:junit-jupiter-api:5.9.3'
+    testImplementation group: 'org.mockito', name: 'mockito-core', version: "${mockito_version}"
+    testImplementation group: 'org.mockito', name: 'mockito-junit-jupiter', version: "${mockito_version}"
+    testImplementation group: 'com.squareup.okhttp3', name: 'mockwebserver', version: '4.12.0'
+
+    testCompileOnly('junit:junit:4.13.1') {
+        exclude group: 'org.hamcrest', module: 'hamcrest-core'
+    }
+    testRuntimeOnly("org.junit.vintage:junit-vintage-engine") {
+        exclude group: 'org.hamcrest', module: 'hamcrest-core'
+    }
+    testRuntimeOnly("org.junit.jupiter:junit-jupiter-engine") {
+        exclude group: 'org.hamcrest', module: 'hamcrest-core'
+    }
+    testImplementation("org.opensearch.test:framework:${opensearch_version}")
+}
+
+test {
+    useJUnitPlatform()
+    testLogging {
+        events "failed"
+        exceptionFormat "full"
+    }
+}
+task junit4(type: Test) {
+    useJUnitPlatform {
+        includeEngines("junit-vintage")
+    }
+    systemProperty 'tests.security.manager', 'false'
+    testLogging {
+        events "failed"
+        exceptionFormat "full"
+    }
+}
+
+jacocoTestReport {
+    dependsOn test, junit4
+    executionData test, junit4
+    reports {
+        html.required = true
+        xml.required = true
+    }
+    afterEvaluate {
+        classDirectories.setFrom(files(classDirectories.files.collect {
+            fileTree(dir: it)
+        }))
+    }
+}
+
+jacocoTestCoverageVerification {
+    dependsOn test, junit4
+    executionData test, junit4
+    violationRules {
+        rule {
+            element = 'CLASS'
+            excludes = [
+                'org.opensearch.sql.prometheus.model.*',
+                'org.opensearch.sql.directquery.rest.model.*'
+            ]
+            limit {
+                counter = 'LINE'
+                minimum = 1.0
+            }
+            limit {
+                counter = 'BRANCH'
+                minimum = 1.0
+            }
+        }
+    }
+    afterEvaluate {
+        classDirectories.setFrom(files(classDirectories.files.collect {
+            fileTree(dir: it)
+        }))
+    }
+}
+check.dependsOn jacocoTestCoverageVerification
+jacocoTestCoverageVerification.dependsOn jacocoTestReport

direct-query-core/src/main/java/org/opensearch/sql/datasource/client/DataSourceClient.java

@@ -0,0 +1,14 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.datasource.client;
+
+/**
+ * Base interface for all data source clients. This interface serves as a marker interface for all
+ * client implementations.
+ *
+ * @opensearch.experimental
+ */
+public interface DataSourceClient {}

direct-query-core/src/main/java/org/opensearch/sql/datasource/client/DataSourceClientFactory.java

@@ -0,0 +1,91 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.datasource.client;
+
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
+import org.opensearch.common.inject.Inject;
+import org.opensearch.sql.common.setting.Settings;
+import org.opensearch.sql.datasource.DataSourceService;
+import org.opensearch.sql.datasource.client.exceptions.DataSourceClientException;
+import org.opensearch.sql.datasource.model.DataSourceMetadata;
+import org.opensearch.sql.datasource.model.DataSourceType;
+import org.opensearch.sql.prometheus.utils.PrometheusClientUtils;
+
+/**
+ * Factory for creating data source clients based on the data source type.
+ *
+ * @opensearch.experimental
+ */
+public class DataSourceClientFactory {
+
+  private static final Logger LOG = LogManager.getLogger();
+
+  private final Settings settings;
+  private final DataSourceService dataSourceService;
+
+  @Inject
+  public DataSourceClientFactory(DataSourceService dataSourceService, Settings settings) {
+    this.settings = settings;
+    this.dataSourceService = dataSourceService;
+  }
+
+  /**
+   * Creates a client for the specified data source with appropriate type.
+   *
+   * @param <T> The type of client to create, must implement DataSourceClient
+   * @param dataSourceName The name of the data source
+   * @return The appropriate client for the data source type
+   * @throws DataSourceClientException If client creation fails
+   */
+  @SuppressWarnings("unchecked")
+  public <T extends DataSourceClient> T createClient(String dataSourceName)
+      throws DataSourceClientException {
+    try {
+      if (!dataSourceService.dataSourceExists(dataSourceName)) {
+        throw new DataSourceClientException("Data source does not exist: " + dataSourceName);
+      }
+
+      DataSourceMetadata metadata =
+          dataSourceService.verifyDataSourceAccessAndGetRawMetadata(dataSourceName, null);
+      DataSourceType dataSourceType = metadata.getConnector();
+
+      return (T) createClientForType(dataSourceType.name(), metadata);
+    } catch (Exception e) {
+      if (e instanceof DataSourceClientException) {
+        throw e;
+      }
+      LOG.error("Failed to create client for data source: " + dataSourceName, e);
+      throw new DataSourceClientException(
+          "Failed to create client for data source: " + dataSourceName, e);
+    }
+  }
+
+  /**
+   * Gets the data source type for a given data source name.
+   *
+   * @param dataSourceName The name of the data source
+   * @return The type of the data source
+   * @throws DataSourceClientException If the data source doesn't exist
+   */
+  public DataSourceType getDataSourceType(String dataSourceName) throws DataSourceClientException {
+    if (!dataSourceService.dataSourceExists(dataSourceName)) {
+      throw new DataSourceClientException("Data source does not exist: " + dataSourceName);
+    }
+
+    return dataSourceService.getDataSourceMetadata(dataSourceName).getConnector();
+  }
+
+  private DataSourceClient createClientForType(String dataSourceType, DataSourceMetadata metadata)
+      throws DataSourceClientException {
+    switch (dataSourceType) {
+      case "PROMETHEUS":
+        return PrometheusClientUtils.createPrometheusClient(metadata, settings);
+      default:
+        throw new DataSourceClientException("Unsupported data source type: " + dataSourceType);
+    }
+  }
+}

direct-query-core/src/main/java/org/opensearch/sql/datasource/client/exceptions/DataSourceClientException.java

@@ -0,0 +1,22 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.datasource.client.exceptions;
+
+/**
+ * Exception thrown when there are issues with data source client operations.
+ *
+ * @opensearch.experimental
+ */
+public class DataSourceClientException extends RuntimeException {
+
+  public DataSourceClientException(String message) {
+    super(message);
+  }
+
+  public DataSourceClientException(String message, Throwable cause) {
+    super(message, cause);
+  }
+}

direct-query-core/src/main/java/org/opensearch/sql/datasource/query/QueryHandler.java

@@ -0,0 +1,79 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.datasource.query;
+
+import java.io.IOException;
+import org.opensearch.sql.datasource.client.DataSourceClient;
+import org.opensearch.sql.datasource.model.DataSourceType;
+import org.opensearch.sql.directquery.rest.model.ExecuteDirectQueryRequest;
+import org.opensearch.sql.directquery.rest.model.GetDirectQueryResourcesRequest;
+import org.opensearch.sql.directquery.rest.model.GetDirectQueryResourcesResponse;
+import org.opensearch.sql.directquery.rest.model.WriteDirectQueryResourcesRequest;
+import org.opensearch.sql.directquery.rest.model.WriteDirectQueryResourcesResponse;
+
+/**
+ * Interface for handling queries for specific data source types.
+ *
+ * @param <T> The client type this handler works with, extending DataSourceClient
+ *
+ * @opensearch.experimental
+ */
+public interface QueryHandler<T extends DataSourceClient> {
+
+  /**
+   * Returns the data source type this handler supports.
+   *
+   * @return The supported data source type
+   */
+  DataSourceType getSupportedDataSourceType();
+
+  /**
+   * Executes a query for the supported data source type.
+   *
+   * @param client The client instance to use
+   * @param request The query request
+   * @return JSON string result of the query
+   * @throws IOException If query execution fails
+   */
+  String executeQuery(T client, ExecuteDirectQueryRequest request) throws IOException;
+
+  /**
+   * Gets resources from the data source.
+   *
+   * @param client The client instance to use
+   * @param request The resources request
+   * @return Response containing the requested resources
+   * @throws IOException If resource retrieval fails
+   */
+  GetDirectQueryResourcesResponse<?> getResources(T client, GetDirectQueryResourcesRequest request)
+      throws IOException;
+
+  /**
+   * Writes resources to the data source.
+   *
+   * @param client The client instance to use
+   * @param request The resources request
+   * @return Response containing the requested resources
+   * @throws IOException If resource retrieval fails
+   */
+  WriteDirectQueryResourcesResponse<?> writeResources(T client, WriteDirectQueryResourcesRequest request)
+          throws IOException;
+
+  /**
+   * Checks if this handler can handle the given client type.
+   *
+   * @param client The client to check
+   * @return true if this handler can handle the client
+   */
+  boolean canHandle(DataSourceClient client);
+
+  /**
+   * Gets the client class this handler supports.
+   *
+   * @return The class of client this handler supports
+   */
+  Class<T> getClientClass();
+}

direct-query-core/src/main/java/org/opensearch/sql/datasource/query/QueryHandlerRegistry.java

@@ -0,0 +1,51 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.datasource.query;
+
+import java.util.List;
+import java.util.Optional;
+import org.opensearch.common.inject.Inject;
+import org.opensearch.sql.datasource.client.DataSourceClient;
+
+/**
+ * Registry for all query handlers.
+ *
+ * @opensearch.experimental
+ */
+public class QueryHandlerRegistry {
+
+  private final List<QueryHandler<?>> handlers;
+
+  @Inject
+  public QueryHandlerRegistry(List<QueryHandler<?>> handlers) {
+    this.handlers = handlers;
+  }
+
+  /**
+   * Finds a handler that can process the given client.
+   *
+   * @param client The client to find a handler for
+   * @param <T> The type of client, extending DataSourceClient
+   * @return An optional containing the handler if found
+   */
+  @SuppressWarnings("unchecked")
+  public <T extends DataSourceClient> Optional<QueryHandler<T>> getQueryHandler(T client) {
+    return handlers.stream()
+        .filter(
+            handler -> {
+              try {
+                // Get the handler's client class and check if it's compatible with our client
+                Class<?> handlerClientClass = handler.getClientClass();
+                return handlerClientClass.isInstance(client)
+                    && ((QueryHandler<T>) handler).canHandle(client);
+              } catch (ClassCastException e) {
+                return false;
+              }
+            })
+        .map(handler -> (QueryHandler<T>) handler)
+        .findFirst();
+  }
+}