From a64eefeeb9bdca8bd6eeab6dbfd2ce6025aa7836 Mon Sep 17 00:00:00 2001 From: Joshua Li Date: Mon, 31 Mar 2025 19:50:56 +0000 Subject: [PATCH] Add direct-query-core module for prometheus integration Co-authored-by: Megha Goyal Signed-off-by: Joshua Li --- .../sql/spark/rest/model/LangType.java | 3 +- direct-query-core/build.gradle | 105 +++ .../datasource/client/DataSourceClient.java | 12 + .../client/DataSourceClientFactory.java | 87 +++ .../exceptions/DataSourceClientException.java | 18 + .../sql/datasource/query/QueryHandler.java | 64 ++ .../query/QueryHandlerRegistry.java | 47 ++ .../DirectQueryExecutorService.java | 31 + .../DirectQueryExecutorServiceImpl.java | 90 +++ .../directquery/model/DataSourceOptions.java | 9 + .../rest/model/DirectQueryResourceType.java | 40 ++ .../rest/model/ExecuteDirectQueryRequest.java | 55 ++ .../model/ExecuteDirectQueryResponse.java | 20 + .../model/GetDirectQueryResourcesRequest.java | 32 + .../GetDirectQueryResourcesResponse.java | 40 ++ .../DirectQueryRequestValidator.java | 86 +++ .../prometheus/client/PrometheusClient.java | 99 +++ .../client/PrometheusClientImpl.java | 397 +++++++++++ .../exception/PrometheusClientException.java | 15 + .../sql/prometheus}/model/MetricMetadata.java | 8 +- .../prometheus/model/PrometheusOptions.java | 21 + .../prometheus/model/PrometheusQueryType.java | 35 + .../query/PrometheusQueryHandler.java | 192 ++++++ .../utils/PrometheusClientUtils.java | 169 +++++ .../client/DataSourceClientFactoryTest.java | 177 +++++ .../query/QueryHandlerRegistryTest.java | 147 ++++ .../DirectQueryExecutorServiceImplTest.java | 285 ++++++++ .../DirectQueryRequestValidatorTest.java | 331 +++++++++ .../client/PrometheusClientImplTest.java | 578 ++++++++++++++++ .../query/PrometheusQueryHandlerTest.java | 644 ++++++++++++++++++ .../utils/PrometheusClientUtilsTest.java | 313 +++++++++ .../src/test/resources/non_json_response.json | 3 + .../prometheus/constant}/TestConstants.java | 2 +- .../sql/prometheus/utils/TestUtils.java | 0 prometheus/build.gradle | 3 +- .../prometheus/client/PrometheusClient.java | 24 - .../client/PrometheusClientImpl.java | 134 ---- .../exceptions/PrometheusClientException.java | 17 - .../PrometheusDescribeMetricRequest.java | 2 +- .../system/PrometheusListMetricsRequest.java | 2 +- .../storage/PrometheusStorageFactory.java | 73 +- .../client/PrometheusClientImplTest.java | 223 ------ ...ExemplarsFunctionTableScanBuilderTest.java | 6 +- ...xemplarsFunctionTableScanOperatorTest.java | 6 +- ...ueryRangeFunctionTableScanBuilderTest.java | 8 +- ...eryRangeFunctionTableScanOperatorTest.java | 8 +- .../PrometheusDescribeMetricRequestTest.java | 2 +- .../PrometheusListMetricsRequestTest.java | 2 +- .../storage/PrometheusMetricScanTest.java | 8 +- .../storage/PrometheusMetricTableTest.java | 2 +- .../storage/QueryExemplarsTableTest.java | 6 +- settings.gradle | 2 + 52 files changed, 4197 insertions(+), 486 deletions(-) create mode 100644 direct-query-core/build.gradle create mode 100644 direct-query-core/src/main/java/org/opensearch/sql/datasource/client/DataSourceClient.java create mode 100644 direct-query-core/src/main/java/org/opensearch/sql/datasource/client/DataSourceClientFactory.java create mode 100644 direct-query-core/src/main/java/org/opensearch/sql/datasource/client/exceptions/DataSourceClientException.java create mode 100644 direct-query-core/src/main/java/org/opensearch/sql/datasource/query/QueryHandler.java create mode 100644 direct-query-core/src/main/java/org/opensearch/sql/datasource/query/QueryHandlerRegistry.java create mode 100644 direct-query-core/src/main/java/org/opensearch/sql/directquery/DirectQueryExecutorService.java create mode 100644 direct-query-core/src/main/java/org/opensearch/sql/directquery/DirectQueryExecutorServiceImpl.java create mode 100644 direct-query-core/src/main/java/org/opensearch/sql/directquery/model/DataSourceOptions.java create mode 100644 direct-query-core/src/main/java/org/opensearch/sql/directquery/rest/model/DirectQueryResourceType.java create mode 100644 direct-query-core/src/main/java/org/opensearch/sql/directquery/rest/model/ExecuteDirectQueryRequest.java create mode 100644 direct-query-core/src/main/java/org/opensearch/sql/directquery/rest/model/ExecuteDirectQueryResponse.java create mode 100644 direct-query-core/src/main/java/org/opensearch/sql/directquery/rest/model/GetDirectQueryResourcesRequest.java create mode 100644 direct-query-core/src/main/java/org/opensearch/sql/directquery/rest/model/GetDirectQueryResourcesResponse.java create mode 100644 direct-query-core/src/main/java/org/opensearch/sql/directquery/validator/DirectQueryRequestValidator.java create mode 100644 direct-query-core/src/main/java/org/opensearch/sql/prometheus/client/PrometheusClient.java create mode 100644 direct-query-core/src/main/java/org/opensearch/sql/prometheus/client/PrometheusClientImpl.java create mode 100644 direct-query-core/src/main/java/org/opensearch/sql/prometheus/exception/PrometheusClientException.java rename {prometheus/src/main/java/org/opensearch/sql/prometheus/request/system => direct-query-core/src/main/java/org/opensearch/sql/prometheus}/model/MetricMetadata.java (74%) create mode 100644 direct-query-core/src/main/java/org/opensearch/sql/prometheus/model/PrometheusOptions.java create mode 100644 direct-query-core/src/main/java/org/opensearch/sql/prometheus/model/PrometheusQueryType.java create mode 100644 direct-query-core/src/main/java/org/opensearch/sql/prometheus/query/PrometheusQueryHandler.java create mode 100644 direct-query-core/src/main/java/org/opensearch/sql/prometheus/utils/PrometheusClientUtils.java create mode 100644 direct-query-core/src/test/java/org/opensearch/sql/datasource/client/DataSourceClientFactoryTest.java create mode 100644 direct-query-core/src/test/java/org/opensearch/sql/datasource/query/QueryHandlerRegistryTest.java create mode 100644 direct-query-core/src/test/java/org/opensearch/sql/directquery/DirectQueryExecutorServiceImplTest.java create mode 100644 direct-query-core/src/test/java/org/opensearch/sql/directquery/validator/DirectQueryRequestValidatorTest.java create mode 100644 direct-query-core/src/test/java/org/opensearch/sql/prometheus/client/PrometheusClientImplTest.java create mode 100644 direct-query-core/src/test/java/org/opensearch/sql/prometheus/query/PrometheusQueryHandlerTest.java create mode 100644 direct-query-core/src/test/java/org/opensearch/sql/prometheus/utils/PrometheusClientUtilsTest.java create mode 100644 direct-query-core/src/test/resources/non_json_response.json rename {prometheus/src/test/java/org/opensearch/sql/prometheus/constants => direct-query-core/src/testFixtures/java/org/opensearch/sql/prometheus/constant}/TestConstants.java (88%) rename {prometheus/src/test => direct-query-core/src/testFixtures}/java/org/opensearch/sql/prometheus/utils/TestUtils.java (100%) delete mode 100644 prometheus/src/main/java/org/opensearch/sql/prometheus/client/PrometheusClient.java delete mode 100644 prometheus/src/main/java/org/opensearch/sql/prometheus/client/PrometheusClientImpl.java delete mode 100644 prometheus/src/main/java/org/opensearch/sql/prometheus/exceptions/PrometheusClientException.java delete mode 100644 prometheus/src/test/java/org/opensearch/sql/prometheus/client/PrometheusClientImplTest.java diff --git a/async-query-core/src/main/java/org/opensearch/sql/spark/rest/model/LangType.java b/async-query-core/src/main/java/org/opensearch/sql/spark/rest/model/LangType.java index 51fa8d2b134..ea02ea74ffb 100644 --- a/async-query-core/src/main/java/org/opensearch/sql/spark/rest/model/LangType.java +++ b/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) { diff --git a/direct-query-core/build.gradle b/direct-query-core/build.gradle new file mode 100644 index 00000000000..25b0f7869c1 --- /dev/null +++ b/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 diff --git a/direct-query-core/src/main/java/org/opensearch/sql/datasource/client/DataSourceClient.java b/direct-query-core/src/main/java/org/opensearch/sql/datasource/client/DataSourceClient.java new file mode 100644 index 00000000000..3f37b5481b6 --- /dev/null +++ b/direct-query-core/src/main/java/org/opensearch/sql/datasource/client/DataSourceClient.java @@ -0,0 +1,12 @@ +/* + * 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. + */ +public interface DataSourceClient {} diff --git a/direct-query-core/src/main/java/org/opensearch/sql/datasource/client/DataSourceClientFactory.java b/direct-query-core/src/main/java/org/opensearch/sql/datasource/client/DataSourceClientFactory.java new file mode 100644 index 00000000000..23b04572175 --- /dev/null +++ b/direct-query-core/src/main/java/org/opensearch/sql/datasource/client/DataSourceClientFactory.java @@ -0,0 +1,87 @@ +/* + * 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. */ +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 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 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); + } + } +} diff --git a/direct-query-core/src/main/java/org/opensearch/sql/datasource/client/exceptions/DataSourceClientException.java b/direct-query-core/src/main/java/org/opensearch/sql/datasource/client/exceptions/DataSourceClientException.java new file mode 100644 index 00000000000..74906fc5b84 --- /dev/null +++ b/direct-query-core/src/main/java/org/opensearch/sql/datasource/client/exceptions/DataSourceClientException.java @@ -0,0 +1,18 @@ +/* + * 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. */ +public class DataSourceClientException extends RuntimeException { + + public DataSourceClientException(String message) { + super(message); + } + + public DataSourceClientException(String message, Throwable cause) { + super(message, cause); + } +} diff --git a/direct-query-core/src/main/java/org/opensearch/sql/datasource/query/QueryHandler.java b/direct-query-core/src/main/java/org/opensearch/sql/datasource/query/QueryHandler.java new file mode 100644 index 00000000000..0fd4311c325 --- /dev/null +++ b/direct-query-core/src/main/java/org/opensearch/sql/datasource/query/QueryHandler.java @@ -0,0 +1,64 @@ +/* + * 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; + +/** + * Interface for handling queries for specific data source types. + * + * @param The client type this handler works with, extending DataSourceClient + */ +public interface QueryHandler { + + /** + * 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; + + /** + * 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 getClientClass(); +} diff --git a/direct-query-core/src/main/java/org/opensearch/sql/datasource/query/QueryHandlerRegistry.java b/direct-query-core/src/main/java/org/opensearch/sql/datasource/query/QueryHandlerRegistry.java new file mode 100644 index 00000000000..ec583c96a82 --- /dev/null +++ b/direct-query-core/src/main/java/org/opensearch/sql/datasource/query/QueryHandlerRegistry.java @@ -0,0 +1,47 @@ +/* + * 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. */ +public class QueryHandlerRegistry { + + private final List> handlers; + + @Inject + public QueryHandlerRegistry(List> handlers) { + this.handlers = handlers; + } + + /** + * Finds a handler that can process the given client. + * + * @param client The client to find a handler for + * @param The type of client, extending DataSourceClient + * @return An optional containing the handler if found + */ + @SuppressWarnings("unchecked") + public Optional> 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) handler).canHandle(client); + } catch (ClassCastException e) { + return false; + } + }) + .map(handler -> (QueryHandler) handler) + .findFirst(); + } +} diff --git a/direct-query-core/src/main/java/org/opensearch/sql/directquery/DirectQueryExecutorService.java b/direct-query-core/src/main/java/org/opensearch/sql/directquery/DirectQueryExecutorService.java new file mode 100644 index 00000000000..828f1864f14 --- /dev/null +++ b/direct-query-core/src/main/java/org/opensearch/sql/directquery/DirectQueryExecutorService.java @@ -0,0 +1,31 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.directquery; + +import org.opensearch.sql.directquery.rest.model.ExecuteDirectQueryRequest; +import org.opensearch.sql.directquery.rest.model.ExecuteDirectQueryResponse; +import org.opensearch.sql.directquery.rest.model.GetDirectQueryResourcesRequest; +import org.opensearch.sql.directquery.rest.model.GetDirectQueryResourcesResponse; + +public interface DirectQueryExecutorService { + + /** + * Execute a direct query request. + * + * @param request The direct query request + * @return A response containing the result + */ + ExecuteDirectQueryResponse executeDirectQuery(ExecuteDirectQueryRequest request); + + /** + * Get resources from a data source. + * + * @param request The resources request + * @return A response containing the requested resources + */ + GetDirectQueryResourcesResponse getDirectQueryResources( + GetDirectQueryResourcesRequest request); +} diff --git a/direct-query-core/src/main/java/org/opensearch/sql/directquery/DirectQueryExecutorServiceImpl.java b/direct-query-core/src/main/java/org/opensearch/sql/directquery/DirectQueryExecutorServiceImpl.java new file mode 100644 index 00000000000..1a2e9e58cc0 --- /dev/null +++ b/direct-query-core/src/main/java/org/opensearch/sql/directquery/DirectQueryExecutorServiceImpl.java @@ -0,0 +1,90 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.directquery; + +import java.io.IOException; +import java.util.UUID; +import org.opensearch.common.inject.Inject; +import org.opensearch.sql.datasource.client.DataSourceClient; +import org.opensearch.sql.datasource.client.DataSourceClientFactory; +import org.opensearch.sql.datasource.client.exceptions.DataSourceClientException; +import org.opensearch.sql.datasource.query.QueryHandlerRegistry; +import org.opensearch.sql.directquery.rest.model.ExecuteDirectQueryRequest; +import org.opensearch.sql.directquery.rest.model.ExecuteDirectQueryResponse; +import org.opensearch.sql.directquery.rest.model.GetDirectQueryResourcesRequest; +import org.opensearch.sql.directquery.rest.model.GetDirectQueryResourcesResponse; + +public class DirectQueryExecutorServiceImpl implements DirectQueryExecutorService { + + private final DataSourceClientFactory dataSourceClientFactory; + private final QueryHandlerRegistry queryHandlerRegistry; + + @Inject + public DirectQueryExecutorServiceImpl( + DataSourceClientFactory dataSourceClientFactory, QueryHandlerRegistry queryHandlerRegistry) { + this.dataSourceClientFactory = dataSourceClientFactory; + this.queryHandlerRegistry = queryHandlerRegistry; + } + + @Override + public ExecuteDirectQueryResponse executeDirectQuery(ExecuteDirectQueryRequest request) { + // TODO: Replace with the data source query id. + String queryId = UUID.randomUUID().toString(); + String sessionId = request.getSessionId(); // Session ID is passed as is + String dataSourceName = request.getDataSources(); + String dataSourceType = null; + String result; + + try { + dataSourceType = + dataSourceClientFactory.getDataSourceType(dataSourceName).name().toLowerCase(); + + DataSourceClient client = dataSourceClientFactory.createClient(dataSourceName); + + result = + queryHandlerRegistry + .getQueryHandler(client) + .map( + handler -> { + try { + return handler.executeQuery(client, request); + } catch (IOException e) { + return "{\"error\": \"Error executing query: " + e.getMessage() + "\"}"; + } + }) + .orElse("{\"error\": \"Unsupported data source type\"}"); + + } catch (Exception e) { + result = "{\"error\": \"" + e.getMessage() + "\"}"; + } + + return new ExecuteDirectQueryResponse(queryId, result, sessionId, dataSourceType); + } + + @Override + public GetDirectQueryResourcesResponse getDirectQueryResources( + GetDirectQueryResourcesRequest request) { + DataSourceClient client = dataSourceClientFactory.createClient(request.getDataSource()); + return queryHandlerRegistry + .getQueryHandler(client) + .map( + handler -> { + try { + return handler.getResources(client, request); + } catch (IOException e) { + throw new DataSourceClientException( + String.format( + "Error retrieving resources for data source type: %s", + request.getDataSource()), + e); + } + }) + .orElseThrow( + () -> + new IllegalArgumentException( + "Unsupported data source type: " + request.getDataSource())); + } +} diff --git a/direct-query-core/src/main/java/org/opensearch/sql/directquery/model/DataSourceOptions.java b/direct-query-core/src/main/java/org/opensearch/sql/directquery/model/DataSourceOptions.java new file mode 100644 index 00000000000..8faa06609e1 --- /dev/null +++ b/direct-query-core/src/main/java/org/opensearch/sql/directquery/model/DataSourceOptions.java @@ -0,0 +1,9 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.directquery.model; + +/** Interface for data source specific options. */ +public interface DataSourceOptions {} diff --git a/direct-query-core/src/main/java/org/opensearch/sql/directquery/rest/model/DirectQueryResourceType.java b/direct-query-core/src/main/java/org/opensearch/sql/directquery/rest/model/DirectQueryResourceType.java new file mode 100644 index 00000000000..1a7070c6e83 --- /dev/null +++ b/direct-query-core/src/main/java/org/opensearch/sql/directquery/rest/model/DirectQueryResourceType.java @@ -0,0 +1,40 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.directquery.rest.model; + +/** Enum representing the types of resources that can be queried. */ +public enum DirectQueryResourceType { + UNKNOWN, + LABELS, + LABEL, + METADATA, + SERIES, + ALERTS, + RULES, + ALERTMANAGER_ALERTS, + ALERTMANAGER_ALERT_GROUPS, + ALERTMANAGER_RECEIVERS, + ALERTMANAGER_SILENCES; + + /** + * Convert a string to the corresponding enum value, case-insensitive. + * + * @param value The string value to convert + * @return The corresponding enum value + * @throws IllegalArgumentException if the value doesn't match any enum value + */ + public static DirectQueryResourceType fromString(String value) { + if (value == null) { + throw new IllegalArgumentException("Resource type cannot be null"); + } + + try { + return valueOf(value.toUpperCase()); + } catch (IllegalArgumentException e) { + throw new IllegalArgumentException("Invalid resource type: " + value); + } + } +} diff --git a/direct-query-core/src/main/java/org/opensearch/sql/directquery/rest/model/ExecuteDirectQueryRequest.java b/direct-query-core/src/main/java/org/opensearch/sql/directquery/rest/model/ExecuteDirectQueryRequest.java new file mode 100644 index 00000000000..6eae6e1ef4e --- /dev/null +++ b/direct-query-core/src/main/java/org/opensearch/sql/directquery/rest/model/ExecuteDirectQueryRequest.java @@ -0,0 +1,55 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.directquery.rest.model; + +import lombok.Data; +import lombok.NoArgsConstructor; +import lombok.NonNull; +import lombok.Setter; +import org.opensearch.sql.directquery.model.DataSourceOptions; +import org.opensearch.sql.prometheus.model.PrometheusOptions; +import org.opensearch.sql.spark.rest.model.LangType; + +@Data +@NoArgsConstructor +public class ExecuteDirectQueryRequest { + // Required fields + private String dataSources; // Required: From URI path parameter or request body + private String query; // Required: String for Prometheus, object for CloudWatch + @Setter private LangType language; // Required: SQL, PPL, or PROMQL + private String sourceVersion; // Required: API version + + // Optional fields + private Integer maxResults; // Optional: limit for Prometheus, maxDataPoints for CW + private Integer timeout; // Optional: number of seconds + private DataSourceOptions options; // Optional: Source specific arguments + private String sessionId; // For session management + + /** + * Helper method to get PrometheusOptions. If options is already PrometheusOptions, returns it. + * Otherwise, returns a new empty PrometheusOptions. + * + * @return PrometheusOptions object + */ + @NonNull + public PrometheusOptions getPrometheusOptions() { + if (options instanceof PrometheusOptions) { + return (PrometheusOptions) options; + } + + // Create new PrometheusOptions if options is null or not PrometheusOptions + return new PrometheusOptions(); + } + + /** + * Set Prometheus options. + * + * @param prometheusOptions The Prometheus options + */ + public void setPrometheusOptions(PrometheusOptions prometheusOptions) { + this.options = prometheusOptions; + } +} diff --git a/direct-query-core/src/main/java/org/opensearch/sql/directquery/rest/model/ExecuteDirectQueryResponse.java b/direct-query-core/src/main/java/org/opensearch/sql/directquery/rest/model/ExecuteDirectQueryResponse.java new file mode 100644 index 00000000000..8bad0693688 --- /dev/null +++ b/direct-query-core/src/main/java/org/opensearch/sql/directquery/rest/model/ExecuteDirectQueryResponse.java @@ -0,0 +1,20 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.directquery.rest.model; + +import lombok.AllArgsConstructor; +import lombok.Data; +import lombok.NoArgsConstructor; + +@Data +@AllArgsConstructor +@NoArgsConstructor +public class ExecuteDirectQueryResponse { + private String queryId; + private String result; + private String sessionId; + private String dataSourceType; +} diff --git a/direct-query-core/src/main/java/org/opensearch/sql/directquery/rest/model/GetDirectQueryResourcesRequest.java b/direct-query-core/src/main/java/org/opensearch/sql/directquery/rest/model/GetDirectQueryResourcesRequest.java new file mode 100644 index 00000000000..54a6b315ef7 --- /dev/null +++ b/direct-query-core/src/main/java/org/opensearch/sql/directquery/rest/model/GetDirectQueryResourcesRequest.java @@ -0,0 +1,32 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.directquery.rest.model; + +import java.util.Map; +import lombok.Data; +import lombok.NoArgsConstructor; + +@Data +@NoArgsConstructor +public class GetDirectQueryResourcesRequest { + private String dataSource; + private DirectQueryResourceType resourceType; + private String resourceName; + + // Optional fields + private Map queryParams; + + /** + * Sets the resource type from a string value. + * + * @param resourceTypeStr The resource type as a string + */ + public void setResourceTypeFromString(String resourceTypeStr) { + if (resourceTypeStr != null) { + this.resourceType = DirectQueryResourceType.fromString(resourceTypeStr); + } + } +} diff --git a/direct-query-core/src/main/java/org/opensearch/sql/directquery/rest/model/GetDirectQueryResourcesResponse.java b/direct-query-core/src/main/java/org/opensearch/sql/directquery/rest/model/GetDirectQueryResourcesResponse.java new file mode 100644 index 00000000000..076423533d3 --- /dev/null +++ b/direct-query-core/src/main/java/org/opensearch/sql/directquery/rest/model/GetDirectQueryResourcesResponse.java @@ -0,0 +1,40 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.directquery.rest.model; + +import com.fasterxml.jackson.annotation.JsonInclude; +import java.util.List; +import java.util.Map; +import lombok.Data; +import lombok.NoArgsConstructor; + +/** + * Response class for direct query resources. + * + * @param The type of data contained in the response + */ +@Data +@NoArgsConstructor +@JsonInclude(JsonInclude.Include.NON_NULL) +public class GetDirectQueryResourcesResponse { + private T data; + + private GetDirectQueryResourcesResponse(T data) { + this.data = data; + } + + public static GetDirectQueryResourcesResponse> withStringList(List data) { + return new GetDirectQueryResourcesResponse<>(data); + } + + public static GetDirectQueryResourcesResponse> withList(List data) { + return new GetDirectQueryResourcesResponse<>(data); + } + + public static GetDirectQueryResourcesResponse> withMap(Map data) { + return new GetDirectQueryResourcesResponse<>(data); + } +} diff --git a/direct-query-core/src/main/java/org/opensearch/sql/directquery/validator/DirectQueryRequestValidator.java b/direct-query-core/src/main/java/org/opensearch/sql/directquery/validator/DirectQueryRequestValidator.java new file mode 100644 index 00000000000..6a945574fd0 --- /dev/null +++ b/direct-query-core/src/main/java/org/opensearch/sql/directquery/validator/DirectQueryRequestValidator.java @@ -0,0 +1,86 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.directquery.validator; + +import org.opensearch.sql.directquery.rest.model.ExecuteDirectQueryRequest; +import org.opensearch.sql.prometheus.model.PrometheusOptions; +import org.opensearch.sql.prometheus.model.PrometheusQueryType; +import org.opensearch.sql.spark.rest.model.LangType; + +public class DirectQueryRequestValidator { + private DirectQueryRequestValidator() {} + + public static void validateRequest(ExecuteDirectQueryRequest request) { + if (request == null) { + throw new IllegalArgumentException("Request cannot be null"); + } + + if (request.getDataSources() == null || request.getDataSources().isEmpty()) { + throw new IllegalArgumentException("Datasource is required"); + } + + if (request.getQuery() == null || request.getQuery().isEmpty()) { + throw new IllegalArgumentException("Query is required"); + } + + if (request.getLanguage() == null) { + throw new IllegalArgumentException("Language type is required"); + } + + if (request.getLanguage() == LangType.PROMQL) { + PrometheusOptions prometheusOptions = request.getPrometheusOptions(); + if (prometheusOptions.getQueryType() == null) { + throw new IllegalArgumentException("Prometheus options are required for PROMQL queries"); + } + + // Validate based on query type + switch (prometheusOptions.getQueryType()) { + case PrometheusQueryType.RANGE: + String start = prometheusOptions.getStart(); + String end = prometheusOptions.getEnd(); + + if (start == null || end == null) { + throw new IllegalArgumentException( + "Start and end times are required for range queries"); + } + + // Validate step parameter + if (prometheusOptions.getStep() == null || prometheusOptions.getStep().isEmpty()) { + throw new IllegalArgumentException("Step parameter is required for range queries"); + } + + // Validate timestamps + try { + long startTimestamp = Long.parseLong(start); + long endTimestamp = Long.parseLong(end); + if (endTimestamp <= startTimestamp) { + throw new IllegalArgumentException("End time must be after start time"); + } + } catch (NumberFormatException e) { + throw new IllegalArgumentException( + "Invalid time format: start and end must be numeric timestamps"); + } + break; + + case PrometheusQueryType.INSTANT: + default: // should not happen. Replace with switch expression when dropping JDK11 support + // For instant queries, validate time parameter + if (prometheusOptions.getTime() == null) { + throw new IllegalArgumentException("Time parameter is required for instant queries"); + } + + // Validate time format + try { + Long.parseLong(prometheusOptions.getTime()); + } catch (NumberFormatException e) { + throw new IllegalArgumentException( + "Invalid time format: time must be a numeric timestamp"); + } + break; + } + } + } +} diff --git a/direct-query-core/src/main/java/org/opensearch/sql/prometheus/client/PrometheusClient.java b/direct-query-core/src/main/java/org/opensearch/sql/prometheus/client/PrometheusClient.java new file mode 100644 index 00000000000..92b54f071d3 --- /dev/null +++ b/direct-query-core/src/main/java/org/opensearch/sql/prometheus/client/PrometheusClient.java @@ -0,0 +1,99 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.prometheus.client; + +import java.io.IOException; +import java.util.List; +import java.util.Map; +import org.json.JSONArray; +import org.json.JSONObject; +import org.opensearch.sql.datasource.client.DataSourceClient; +import org.opensearch.sql.prometheus.model.MetricMetadata; + +public interface PrometheusClient extends DataSourceClient { + + JSONObject queryRange(String query, Long start, Long end, String step) throws IOException; + + JSONObject queryRange( + String query, Long start, Long end, String step, Integer limit, Integer timeout) + throws IOException; + + List getLabels(String metricName) throws IOException; + + List getLabels(Map queryParams) throws IOException; + + List getLabel(String labelName, Map queryParams) throws IOException; + + Map> getAllMetrics() throws IOException; + + Map> getAllMetrics(Map queryParams) + throws IOException; + + List> getSeries(Map queryParams) throws IOException; + + JSONArray queryExemplars(String query, Long start, Long end) throws IOException; + + /** + * Execute an instant query at a single point in time. + * + * @param query The Prometheus expression query string + * @param time Optional evaluation timestamp (Unix timestamp in seconds) + * @return JSONObject containing the query result data + * @throws IOException If there is an issue with the request + */ + JSONObject query(String query, Long time, Integer limit, Integer timeout) throws IOException; + + /** + * Get all alerting rules. + * + * @return JSONObject containing the alerting rules + * @throws IOException If there is an issue with the request + */ + JSONObject getAlerts() throws IOException; + + /** + * Get all recording and alerting rules. + * + * @param queryParams Map of query parameters to include in the request + * @return JSONObject containing all rules + * @throws IOException If there is an issue with the request + */ + JSONObject getRules(Map queryParams) throws IOException; + + /** + * Get all alerts from Alertmanager. + * + * @param queryParams Map of query parameters to include in the request + * @return JSONArray containing the alerts + * @throws IOException If there is an issue with the request + */ + JSONArray getAlertmanagerAlerts(Map queryParams) throws IOException; + + /** + * Get alerts grouped according to Alertmanager configuration. + * + * @param queryParams Map of query parameters to include in the request + * @return JSONArray containing the alert groups + * @throws IOException If there is an issue with the request + */ + JSONArray getAlertmanagerAlertGroups(Map queryParams) throws IOException; + + /** + * Get all receivers configured in Alertmanager. + * + * @return JSONArray containing the receivers + * @throws IOException If there is an issue with the request + */ + JSONArray getAlertmanagerReceivers() throws IOException; + + /** + * Get all silences configured in Alertmanager. + * + * @return JSONArray containing the silences + * @throws IOException If there is an issue with the request + */ + JSONArray getAlertmanagerSilences() throws IOException; +} diff --git a/direct-query-core/src/main/java/org/opensearch/sql/prometheus/client/PrometheusClientImpl.java b/direct-query-core/src/main/java/org/opensearch/sql/prometheus/client/PrometheusClientImpl.java new file mode 100644 index 00000000000..c9633e93211 --- /dev/null +++ b/direct-query-core/src/main/java/org/opensearch/sql/prometheus/client/PrometheusClientImpl.java @@ -0,0 +1,397 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.prometheus.client; + +import com.fasterxml.jackson.core.type.TypeReference; +import com.fasterxml.jackson.databind.ObjectMapper; +import java.io.IOException; +import java.net.URI; +import java.net.URLEncoder; +import java.nio.charset.StandardCharsets; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Objects; +import java.util.stream.Collectors; +import okhttp3.OkHttpClient; +import okhttp3.Request; +import okhttp3.Response; +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; +import org.json.JSONArray; +import org.json.JSONException; +import org.json.JSONObject; +import org.opensearch.sql.prometheus.model.MetricMetadata; + +public class PrometheusClientImpl implements PrometheusClient { + + private static final Logger logger = LogManager.getLogger(PrometheusClientImpl.class); + + private final OkHttpClient prometheusHttpClient; + private final OkHttpClient alertmanagerHttpClient; + + private final URI prometheusUri; + private final URI alertmanagerUri; + + public PrometheusClientImpl(OkHttpClient prometheusHttpClient, URI prometheusUri) { + this( + prometheusHttpClient, + prometheusUri, + prometheusHttpClient, + URI.create(prometheusUri.toString().replaceAll("/$", "") + "/alertmanager")); + } + + public PrometheusClientImpl( + OkHttpClient prometheusHttpClient, + URI prometheusUri, + OkHttpClient alertmanagerHttpClient, + URI alertmanagerUri) { + this.prometheusHttpClient = prometheusHttpClient; + this.prometheusUri = prometheusUri; + this.alertmanagerHttpClient = alertmanagerHttpClient; + this.alertmanagerUri = alertmanagerUri; + } + + private String paramsToQueryString(Map queryParams) { + String queryString = + queryParams.entrySet().stream() + .map( + entry -> + URLEncoder.encode(entry.getKey(), StandardCharsets.UTF_8) + + "=" + + URLEncoder.encode(entry.getValue(), StandardCharsets.UTF_8)) + .collect(Collectors.joining("&")); + return queryString.isEmpty() ? "" : "?" + queryString; + } + + @Override + public JSONObject queryRange(String query, Long start, Long end, String step) throws IOException { + return queryRange(query, start, end, step, null, null); + } + + @Override + public JSONObject queryRange( + String query, Long start, Long end, String step, Integer limit, Integer timeout) + throws IOException { + String queryString = buildQueryString(query, start, end, step, limit, timeout); + String queryUrl = + String.format( + "%s/api/v1/query_range%s", prometheusUri.toString().replaceAll("/$", ""), queryString); + + logger.debug("Making Prometheus query_range request: {}", queryUrl); + Request request = new Request.Builder().url(queryUrl).build(); + + logger.debug("Executing Prometheus request with headers: {}", request.headers().toString()); + Response response = this.prometheusHttpClient.newCall(request).execute(); + + logger.debug("Received Prometheus response for query_range: code={}", response); + + JSONObject jsonObject = readResponse(response); + return jsonObject.getJSONObject("data"); + } + + @Override + public JSONObject query(String query, Long time, Integer limit, Integer timeout) + throws IOException { + Map params = new HashMap<>(); + params.put("query", query); + + // Add optional parameters + if (time != null) { + params.put("time", time.toString()); + } + if (limit != null) { + params.put("limit", limit.toString()); + } + if (timeout != null) { + params.put("timeout", timeout.toString()); + } + + String queryString = this.paramsToQueryString(params); + + String queryUrl = + String.format( + "%s/api/v1/query%s", prometheusUri.toString().replaceAll("/$", ""), queryString); + + logger.info("Making Prometheus instant query request: {}", queryUrl); + Request request = new Request.Builder().url(queryUrl).build(); + + logger.info("Executing Prometheus request with headers: {}", request.headers().toString()); + Response response = this.prometheusHttpClient.newCall(request).execute(); + + logger.info("Received Prometheus response for instant query: code={}", response); + // Return the full response object, not just the data field + return readResponse(response); + } + + @Override + public List getLabels(String metricName) throws IOException { + return getLabels(Map.of("match[]", metricName)); + } + + @Override + public List getLabels(Map queryParams) throws IOException { + String queryString = this.paramsToQueryString(queryParams); + String queryUrl = + String.format( + "%s/api/v1/labels%s", prometheusUri.toString().replaceAll("/$", ""), queryString); + logger.debug("queryUrl: " + queryUrl); + Request request = new Request.Builder().url(queryUrl).build(); + Response response = this.prometheusHttpClient.newCall(request).execute(); + JSONObject jsonObject = readResponse(response); + return toListOfLabels(jsonObject.getJSONArray("data")); + } + + @Override + public List getLabel(String labelName, Map queryParams) + throws IOException { + String queryString = this.paramsToQueryString(queryParams); + String queryUrl = + String.format( + "%s/api/v1/label/%s/values%s", + prometheusUri.toString().replaceAll("/$", ""), labelName, queryString); + logger.debug("queryUrl: " + queryUrl); + Request request = new Request.Builder().url(queryUrl).build(); + Response response = this.prometheusHttpClient.newCall(request).execute(); + JSONObject jsonObject = readResponse(response); + return toListOfLabels(jsonObject.getJSONArray("data")); + } + + @Override + public Map> getAllMetrics(Map queryParams) + throws IOException { + String queryString = this.paramsToQueryString(queryParams); + String queryUrl = + String.format( + "%s/api/v1/metadata%s", prometheusUri.toString().replaceAll("/$", ""), queryString); + logger.debug("queryUrl: " + queryUrl); + Request request = new Request.Builder().url(queryUrl).build(); + Response response = this.prometheusHttpClient.newCall(request).execute(); + JSONObject jsonObject = readResponse(response); + TypeReference>> typeRef = new TypeReference<>() {}; + return new ObjectMapper().readValue(jsonObject.getJSONObject("data").toString(), typeRef); + } + + @Override + public Map> getAllMetrics() throws IOException { + return getAllMetrics(Map.of()); + } + + @Override + public List> getSeries(Map queryParams) throws IOException { + String queryString = this.paramsToQueryString(queryParams); + String queryUrl = + String.format( + "%s/api/v1/series%s", prometheusUri.toString().replaceAll("/$", ""), queryString); + logger.debug("queryUrl: " + queryUrl); + Request request = new Request.Builder().url(queryUrl).build(); + Response response = this.prometheusHttpClient.newCall(request).execute(); + JSONObject jsonObject = readResponse(response); + JSONArray dataArray = jsonObject.getJSONArray("data"); + return toListOfSeries(dataArray); + } + + @Override + public JSONArray queryExemplars(String query, Long start, Long end) throws IOException { + String queryUrl = + String.format( + "%s/api/v1/query_exemplars?query=%s&start=%s&end=%s", + prometheusUri.toString().replaceAll("/$", ""), + URLEncoder.encode(query, StandardCharsets.UTF_8), + start, + end); + logger.debug("queryUrl: " + queryUrl); + Request request = new Request.Builder().url(queryUrl).build(); + Response response = this.prometheusHttpClient.newCall(request).execute(); + JSONObject jsonObject = readResponse(response); + return jsonObject.getJSONArray("data"); + } + + @Override + public JSONObject getAlerts() throws IOException { + String queryUrl = + String.format("%s/api/v1/alerts", prometheusUri.toString().replaceAll("/$", "")); + logger.debug("Making Prometheus alerts request: {}", queryUrl); + Request request = new Request.Builder().url(queryUrl).build(); + Response response = this.prometheusHttpClient.newCall(request).execute(); + JSONObject jsonObject = readResponse(response); + return jsonObject.getJSONObject("data"); + } + + @Override + public JSONObject getRules(Map queryParams) throws IOException { + String queryString = this.paramsToQueryString(queryParams); + String queryUrl = + String.format( + "%s/api/v1/rules%s", prometheusUri.toString().replaceAll("/$", ""), queryString); + logger.debug("Making Prometheus rules request: {}", queryUrl); + Request request = new Request.Builder().url(queryUrl).build(); + Response response = this.prometheusHttpClient.newCall(request).execute(); + JSONObject jsonObject = readResponse(response); + return jsonObject.getJSONObject("data"); + } + + @Override + public JSONArray getAlertmanagerAlerts(Map queryParams) throws IOException { + String queryString = this.paramsToQueryString(queryParams); + String baseUrl = alertmanagerUri.toString().replaceAll("/$", ""); + String queryUrl = String.format("%s/api/v2/alerts%s", baseUrl, queryString); + + logger.debug("Making Alertmanager alerts request: {}", queryUrl); + Request request = new Request.Builder().url(queryUrl).build(); + Response response = this.alertmanagerHttpClient.newCall(request).execute(); + + return readAlertmanagerResponse(response); + } + + @Override + public JSONArray getAlertmanagerAlertGroups(Map queryParams) throws IOException { + String queryString = this.paramsToQueryString(queryParams); + String baseUrl = alertmanagerUri.toString().replaceAll("/$", ""); + String queryUrl = String.format("%s/api/v2/alerts/groups%s", baseUrl, queryString); + + logger.debug("Making Alertmanager alert groups request: {}", queryUrl); + Request request = new Request.Builder().url(queryUrl).build(); + Response response = this.alertmanagerHttpClient.newCall(request).execute(); + + return readAlertmanagerResponse(response); + } + + @Override + public JSONArray getAlertmanagerReceivers() throws IOException { + String baseUrl = alertmanagerUri.toString().replaceAll("/$", ""); + String queryUrl = String.format("%s/api/v2/receivers", baseUrl); + + logger.debug("Making Alertmanager receivers request: {}", queryUrl); + Request request = new Request.Builder().url(queryUrl).build(); + Response response = this.alertmanagerHttpClient.newCall(request).execute(); + + return readAlertmanagerResponse(response); + } + + @Override + public JSONArray getAlertmanagerSilences() throws IOException { + String baseUrl = alertmanagerUri.toString().replaceAll("/$", ""); + String queryUrl = String.format("%s/api/v2/silences", baseUrl); + + logger.debug("Making Alertmanager silences request: {}", queryUrl); + Request request = new Request.Builder().url(queryUrl).build(); + Response response = this.alertmanagerHttpClient.newCall(request).execute(); + + return readAlertmanagerResponse(response); + } + + /** + * Reads and processes an Alertmanager API response. + * + * @param response The HTTP response from Alertmanager + * @return A JSONArray with the processed response + * @throws IOException If there's an error reading the response + */ + private JSONArray readAlertmanagerResponse(Response response) throws IOException { + if (response.isSuccessful()) { + String bodyString = Objects.requireNonNull(response.body()).string(); + logger.debug("Alertmanager response body: {}", bodyString); + + // Parse the response body directly as JSON array + return new JSONArray(bodyString); + } else { + String errorBody = response.body() != null ? response.body().string() : "No response body"; + logger.error( + "Alertmanager request failed with code: {}, error body: {}", response.code(), errorBody); + throw new org.opensearch.sql.prometheus.exception.PrometheusClientException( + String.format( + "Alertmanager request failed with code: %s. Error details: %s", + response.code(), errorBody)); + } + } + + private List toListOfLabels(JSONArray array) { + List result = new ArrayList<>(); + for (int i = 0; i < array.length(); i++) { + // __name__ is internal label in prometheus representing the metric name. + // Exempting this from labels list as it is not required in any of the operations. + if (!"__name__".equals(array.optString(i))) { + result.add(array.optString(i)); + } + } + return result; + } + + private List> toListOfSeries(JSONArray array) { + List> result = new ArrayList<>(); + for (int i = 0; i < array.length(); i++) { + JSONObject obj = array.getJSONObject(i); + Map map = new HashMap<>(); + for (String key : obj.keySet()) { + map.put(key, obj.getString(key)); + } + result.add(map); + } + return result; + } + + private JSONObject readResponse(Response response) throws IOException { + // Log the response information + logger.debug("Prometheus response code: {}", response.code()); + + String requestId = response.header("X-Request-ID"); + if (requestId != null) { + logger.info("Prometheus request ID: {}", requestId); + } + if (response.isSuccessful()) { + JSONObject jsonObject; + try { + String bodyString = Objects.requireNonNull(response.body()).string(); + logger.debug("Prometheus response body: {}", bodyString); + jsonObject = new JSONObject(bodyString); + } catch (JSONException jsonException) { + logger.error("Failed to parse Prometheus response as JSON", jsonException); + throw new org.opensearch.sql.prometheus.exception.PrometheusClientException( + "Prometheus returned unexpected body, " + + "please verify your prometheus server setup."); + } + + String status = jsonObject.getString("status"); + logger.debug("Prometheus response status: {}", status); + + if ("success".equals(status)) { + return jsonObject; + } else { + String errorMessage = jsonObject.getString("error"); + logger.error("Prometheus returned error status: {}", errorMessage); + throw new org.opensearch.sql.prometheus.exception.PrometheusClientException(errorMessage); + } + } else { + String errorBody = response.body() != null ? response.body().string() : "No response body"; + logger.error( + "Prometheus request failed with code: {}, error body: {}", response.code(), errorBody); + throw new org.opensearch.sql.prometheus.exception.PrometheusClientException( + String.format( + "Request to Prometheus is Unsuccessful with code: %s. Error details: %s", + response.code(), errorBody)); + } + } + + private String buildQueryString( + String query, Long start, Long end, String step, Integer limit, Integer timeout) { + Map params = new HashMap<>(); + params.put("query", query); + params.put("start", start.toString()); + params.put("end", end.toString()); + params.put("step", step); + + if (limit != null) { + params.put("limit", limit.toString()); + } + if (timeout != null) { + params.put("timeout", timeout.toString()); + } + + return this.paramsToQueryString(params); + } +} diff --git a/direct-query-core/src/main/java/org/opensearch/sql/prometheus/exception/PrometheusClientException.java b/direct-query-core/src/main/java/org/opensearch/sql/prometheus/exception/PrometheusClientException.java new file mode 100644 index 00000000000..6a491243d8d --- /dev/null +++ b/direct-query-core/src/main/java/org/opensearch/sql/prometheus/exception/PrometheusClientException.java @@ -0,0 +1,15 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.prometheus.exception; + +import org.opensearch.sql.datasource.client.exceptions.DataSourceClientException; + +/** PrometheusClientException. */ +public class PrometheusClientException extends DataSourceClientException { + public PrometheusClientException(String message) { + super(message); + } +} diff --git a/prometheus/src/main/java/org/opensearch/sql/prometheus/request/system/model/MetricMetadata.java b/direct-query-core/src/main/java/org/opensearch/sql/prometheus/model/MetricMetadata.java similarity index 74% rename from prometheus/src/main/java/org/opensearch/sql/prometheus/request/system/model/MetricMetadata.java rename to direct-query-core/src/main/java/org/opensearch/sql/prometheus/model/MetricMetadata.java index 195d56f405a..f0f8d8cb220 100644 --- a/prometheus/src/main/java/org/opensearch/sql/prometheus/request/system/model/MetricMetadata.java +++ b/direct-query-core/src/main/java/org/opensearch/sql/prometheus/model/MetricMetadata.java @@ -1,11 +1,9 @@ /* - * - * * Copyright OpenSearch Contributors - * * SPDX-License-Identifier: Apache-2.0 - * + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 */ -package org.opensearch.sql.prometheus.request.system.model; +package org.opensearch.sql.prometheus.model; import com.fasterxml.jackson.annotation.JsonIgnoreProperties; import lombok.AllArgsConstructor; diff --git a/direct-query-core/src/main/java/org/opensearch/sql/prometheus/model/PrometheusOptions.java b/direct-query-core/src/main/java/org/opensearch/sql/prometheus/model/PrometheusOptions.java new file mode 100644 index 00000000000..ab4870557ae --- /dev/null +++ b/direct-query-core/src/main/java/org/opensearch/sql/prometheus/model/PrometheusOptions.java @@ -0,0 +1,21 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.prometheus.model; + +import lombok.Data; +import lombok.NoArgsConstructor; +import org.opensearch.sql.directquery.model.DataSourceOptions; + +/** Prometheus-specific options for direct queries. */ +@Data +@NoArgsConstructor +public class PrometheusOptions implements DataSourceOptions { + private PrometheusQueryType queryType; + private String step; // Duration string in seconds + private String time; // ISO timestamp for instant queries + private String start; // ISO timestamp for range queries + private String end; // ISO timestamp for range queries +} diff --git a/direct-query-core/src/main/java/org/opensearch/sql/prometheus/model/PrometheusQueryType.java b/direct-query-core/src/main/java/org/opensearch/sql/prometheus/model/PrometheusQueryType.java new file mode 100644 index 00000000000..5fb08e030fe --- /dev/null +++ b/direct-query-core/src/main/java/org/opensearch/sql/prometheus/model/PrometheusQueryType.java @@ -0,0 +1,35 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.prometheus.model; + +import lombok.Getter; + +/** Enum representing the types of Prometheus queries. */ +@Getter +public enum PrometheusQueryType { + INSTANT("instant"), + RANGE("range"); + + private final String value; + + PrometheusQueryType(String value) { + this.value = value; + } + + public static PrometheusQueryType fromString(String text) { + for (PrometheusQueryType type : PrometheusQueryType.values()) { + if (type.value.equalsIgnoreCase(text)) { + return type; + } + } + throw new IllegalArgumentException("Unknown query type: " + text); + } + + @Override + public String toString() { + return value; + } +} diff --git a/direct-query-core/src/main/java/org/opensearch/sql/prometheus/query/PrometheusQueryHandler.java b/direct-query-core/src/main/java/org/opensearch/sql/prometheus/query/PrometheusQueryHandler.java new file mode 100644 index 00000000000..102b39c727d --- /dev/null +++ b/direct-query-core/src/main/java/org/opensearch/sql/prometheus/query/PrometheusQueryHandler.java @@ -0,0 +1,192 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.prometheus.query; + +import java.io.IOException; +import java.security.AccessController; +import java.security.PrivilegedAction; +import java.util.List; +import java.util.Map; +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; +import org.json.JSONArray; +import org.json.JSONObject; +import org.opensearch.sql.datasource.client.DataSourceClient; +import org.opensearch.sql.datasource.model.DataSourceType; +import org.opensearch.sql.datasource.query.QueryHandler; +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.prometheus.client.PrometheusClient; +import org.opensearch.sql.prometheus.exception.PrometheusClientException; +import org.opensearch.sql.prometheus.model.MetricMetadata; +import org.opensearch.sql.prometheus.model.PrometheusOptions; +import org.opensearch.sql.prometheus.model.PrometheusQueryType; + +public class PrometheusQueryHandler implements QueryHandler { + private static final Logger LOG = LogManager.getLogger(PrometheusQueryHandler.class); + + @Override + public DataSourceType getSupportedDataSourceType() { + return DataSourceType.PROMETHEUS; + } + + @Override + public boolean canHandle(DataSourceClient client) { + return client instanceof PrometheusClient; + } + + @Override + public Class getClientClass() { + return PrometheusClient.class; + } + + @Override + public String executeQuery(PrometheusClient client, ExecuteDirectQueryRequest request) + throws IOException { + return AccessController.doPrivileged( + (PrivilegedAction) + () -> { + try { + PrometheusOptions options = request.getPrometheusOptions(); + PrometheusQueryType queryType = options.getQueryType(); + if (queryType == null) { + return createErrorJson("Query type is required for Prometheus queries"); + } + + String startTimeStr = options.getStart(); + String endTimeStr = options.getEnd(); + Integer limit = request.getMaxResults(); + Integer timeout = request.getTimeout(); + + if (queryType == PrometheusQueryType.RANGE + && (startTimeStr == null || endTimeStr == null)) { + return createErrorJson("Start and end times are required for Prometheus queries"); + } else if (queryType == PrometheusQueryType.INSTANT && options.getTime() == null) { + return createErrorJson("Time is required for instant Prometheus queries"); + } + + switch (queryType) { + case RANGE: + { + JSONObject metricData = + client.queryRange( + request.getQuery(), + Long.parseLong(startTimeStr), + Long.parseLong(endTimeStr), + options.getStep(), + limit, + timeout); + return metricData.toString(); + } + + case INSTANT: + default: + { + JSONObject metricData = + client.query( + request.getQuery(), + Long.parseLong(options.getTime()), + limit, + timeout); + return metricData.toString(); + } + } + } catch (NumberFormatException e) { + return createErrorJson("Invalid time format: " + e.getMessage()); + } catch (PrometheusClientException e) { + LOG.error("Prometheus client error executing query", e); + return createErrorJson(e.getMessage()); + } catch (IOException e) { + LOG.error("Error executing query", e); + return createErrorJson(e.getMessage()); + } + }); + } + + private String createErrorJson(String message) { + return new JSONObject().put("error", message).toString(); + } + + @Override + public GetDirectQueryResourcesResponse getResources( + PrometheusClient client, GetDirectQueryResourcesRequest request) { + return AccessController.doPrivileged( + (PrivilegedAction>) + () -> { + try { + if (request.getResourceType() == null) { + throw new IllegalArgumentException("Resource type cannot be null"); + } + + switch (request.getResourceType()) { + case LABELS: + { + List labels = client.getLabels(request.getQueryParams()); + return GetDirectQueryResourcesResponse.withStringList(labels); + } + case LABEL: + { + List labelValues = + client.getLabel(request.getResourceName(), request.getQueryParams()); + return GetDirectQueryResourcesResponse.withStringList(labelValues); + } + case METADATA: + { + Map> metadata = + client.getAllMetrics(request.getQueryParams()); + return GetDirectQueryResourcesResponse.withMap(metadata); + } + case SERIES: + { + List> series = client.getSeries(request.getQueryParams()); + return GetDirectQueryResourcesResponse.withList(series); + } + case ALERTS: + { + JSONObject alerts = client.getAlerts(); + return GetDirectQueryResourcesResponse.withMap(alerts.toMap()); + } + case RULES: + { + JSONObject rules = client.getRules(request.getQueryParams()); + return GetDirectQueryResourcesResponse.withMap(rules.toMap()); + } + case ALERTMANAGER_ALERTS: + { + JSONArray alerts = client.getAlertmanagerAlerts(request.getQueryParams()); + return GetDirectQueryResourcesResponse.withList(alerts.toList()); + } + case ALERTMANAGER_ALERT_GROUPS: + { + JSONArray alertGroups = + client.getAlertmanagerAlertGroups(request.getQueryParams()); + return GetDirectQueryResourcesResponse.withList(alertGroups.toList()); + } + case ALERTMANAGER_RECEIVERS: + { + JSONArray receivers = client.getAlertmanagerReceivers(); + return GetDirectQueryResourcesResponse.withList(receivers.toList()); + } + case ALERTMANAGER_SILENCES: + { + JSONArray silences = client.getAlertmanagerSilences(); + return GetDirectQueryResourcesResponse.withList(silences.toList()); + } + default: + throw new IllegalArgumentException( + "Invalid prometheus resource type: " + request.getResourceType()); + } + } catch (IOException e) { + LOG.error("Error getting resources", e); + throw new org.opensearch.sql.prometheus.exception.PrometheusClientException( + String.format( + "Error while getting resources for %s: %s", + request.getResourceType(), e.getMessage())); + } + }); + } +} diff --git a/direct-query-core/src/main/java/org/opensearch/sql/prometheus/utils/PrometheusClientUtils.java b/direct-query-core/src/main/java/org/opensearch/sql/prometheus/utils/PrometheusClientUtils.java new file mode 100644 index 00000000000..fd47da9b2b9 --- /dev/null +++ b/direct-query-core/src/main/java/org/opensearch/sql/prometheus/utils/PrometheusClientUtils.java @@ -0,0 +1,169 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.prometheus.utils; + +import com.amazonaws.auth.AWSStaticCredentialsProvider; +import com.amazonaws.auth.BasicAWSCredentials; +import java.net.URI; +import java.security.AccessController; +import java.security.PrivilegedAction; +import java.util.HashMap; +import java.util.Map; +import java.util.Objects; +import java.util.concurrent.TimeUnit; +import okhttp3.OkHttpClient; +import org.opensearch.sql.common.interceptors.AwsSigningInterceptor; +import org.opensearch.sql.common.interceptors.BasicAuthenticationInterceptor; +import org.opensearch.sql.common.interceptors.URIValidatorInterceptor; +import org.opensearch.sql.common.setting.Settings; +import org.opensearch.sql.datasource.client.exceptions.DataSourceClientException; +import org.opensearch.sql.datasource.model.DataSourceMetadata; +import org.opensearch.sql.datasources.auth.AuthenticationType; +import org.opensearch.sql.prometheus.client.PrometheusClient; +import org.opensearch.sql.prometheus.client.PrometheusClientImpl; + +public class PrometheusClientUtils { + private PrometheusClientUtils() {} + + // Prometheus auth constants + public static final String AUTH_TYPE = "prometheus.auth.type"; + public static final String USERNAME = "prometheus.auth.username"; + public static final String PASSWORD = "prometheus.auth.password"; + public static final String REGION = "prometheus.auth.region"; + public static final String ACCESS_KEY = "prometheus.auth.access_key"; + public static final String SECRET_KEY = "prometheus.auth.secret_key"; + + // Prometheus URI constant + public static final String PROMETHEUS_URI = "prometheus.uri"; + + // AlertManager constants + public static final String ALERTMANAGER_URI = "alertmanager.uri"; + public static final String ALERTMANAGER_AUTH_TYPE = "alertmanager.auth.type"; + public static final String ALERTMANAGER_USERNAME = "alertmanager.auth.username"; + public static final String ALERTMANAGER_PASSWORD = "alertmanager.auth.password"; + public static final String ALERTMANAGER_REGION = "alertmanager.auth.region"; + public static final String ALERTMANAGER_ACCESS_KEY = "alertmanager.auth.access_key"; + public static final String ALERTMANAGER_SECRET_KEY = "alertmanager.auth.secret_key"; + + public static OkHttpClient getHttpClient(Map config, Settings settings) { + return AccessController.doPrivileged( + (PrivilegedAction) + () -> { + OkHttpClient.Builder okHttpClient = new OkHttpClient.Builder(); + okHttpClient.callTimeout(1, TimeUnit.MINUTES); + okHttpClient.connectTimeout(30, TimeUnit.SECONDS); + okHttpClient.followRedirects(false); + okHttpClient.addInterceptor( + new URIValidatorInterceptor( + settings.getSettingValue(Settings.Key.DATASOURCES_URI_HOSTS_DENY_LIST))); + if (config.get(AUTH_TYPE) != null) { + AuthenticationType authenticationType = + AuthenticationType.get(config.get(AUTH_TYPE)); + if (AuthenticationType.BASICAUTH.equals(authenticationType)) { + okHttpClient.addInterceptor( + new BasicAuthenticationInterceptor( + config.get(USERNAME), config.get(PASSWORD))); + } else if (AuthenticationType.AWSSIGV4AUTH.equals(authenticationType)) { + okHttpClient.addInterceptor( + new AwsSigningInterceptor( + new AWSStaticCredentialsProvider( + new BasicAWSCredentials( + config.get(ACCESS_KEY), config.get(SECRET_KEY))), + config.get(REGION), + "aps")); + } else { + throw new IllegalArgumentException( + String.format( + "AUTH Type : %s is not supported with Prometheus Connector", + config.get(AUTH_TYPE))); + } + } + return okHttpClient.build(); + }); + } + + /** + * Creates a properties map for Alertmanager authentication based on the data source properties. + * + * @param properties The data source properties + * @return A map containing Alertmanager authentication properties + */ + public static Map createAlertmanagerProperties(Map properties) { + Map alertmanagerProperties = new HashMap<>(); + + if (properties.containsKey(ALERTMANAGER_AUTH_TYPE)) { + alertmanagerProperties.put(AUTH_TYPE, properties.get(ALERTMANAGER_AUTH_TYPE)); + + String authType = properties.get(ALERTMANAGER_AUTH_TYPE); + if (Objects.nonNull(authType)) { + if (authType.equalsIgnoreCase("basicauth")) { + alertmanagerProperties.put(USERNAME, properties.get(ALERTMANAGER_USERNAME)); + alertmanagerProperties.put(PASSWORD, properties.get(ALERTMANAGER_PASSWORD)); + } else if (authType.equalsIgnoreCase("awssigv4auth")) { + alertmanagerProperties.put(ACCESS_KEY, properties.get(ALERTMANAGER_ACCESS_KEY)); + alertmanagerProperties.put(SECRET_KEY, properties.get(ALERTMANAGER_SECRET_KEY)); + alertmanagerProperties.put(REGION, properties.get(ALERTMANAGER_REGION)); + } + } + } + + return alertmanagerProperties; + } + + /** + * Checks if AlertManager configuration is present in the properties. + * + * @param properties The data source properties + * @return true if Alertmanager URI is present, false otherwise + */ + public static boolean hasAlertmanagerConfig(Map properties) { + return Objects.nonNull(properties.get(ALERTMANAGER_URI)); + } + + /** + * Creates a PrometheusClient instance based on the provided data source metadata and settings. If + * alertmanager settings are present, it creates an alertmanager client as well. Otherwise, it + * reuses the prometheus http client for alertmanager calls with URI to be + * {prometheus.url}/alertmanager. + * + * @param metadata The data source metadata + * @param settings The application settings + * @return A PrometheusClient instance + * @throws DataSourceClientException if the host is not provided + */ + public static PrometheusClient createPrometheusClient( + DataSourceMetadata metadata, Settings settings) { + String host = metadata.getProperties().get(PrometheusClientUtils.PROMETHEUS_URI); + if (Objects.isNull(host)) { + throw new DataSourceClientException("Host is required for Prometheus data source"); + } + URI uri = URI.create(host); + + Map properties = metadata.getProperties(); + OkHttpClient prometheusHttpClient = PrometheusClientUtils.getHttpClient(properties, settings); + + URI alertmanagerUri; + OkHttpClient alertmanagerHttpClient = prometheusHttpClient; + + if (PrometheusClientUtils.hasAlertmanagerConfig(properties)) { + String alertmanagerHost = properties.get(PrometheusClientUtils.ALERTMANAGER_URI); + alertmanagerUri = URI.create(alertmanagerHost); + + Map alertmanagerProperties = + PrometheusClientUtils.createAlertmanagerProperties(properties); + + if (!alertmanagerProperties.isEmpty()) { + alertmanagerHttpClient = + PrometheusClientUtils.getHttpClient(alertmanagerProperties, settings); + } + } else { + alertmanagerUri = URI.create(host.replaceAll("/$", "") + "/alertmanager"); + } + + return new PrometheusClientImpl( + prometheusHttpClient, uri, alertmanagerHttpClient, alertmanagerUri); + } +} diff --git a/direct-query-core/src/test/java/org/opensearch/sql/datasource/client/DataSourceClientFactoryTest.java b/direct-query-core/src/test/java/org/opensearch/sql/datasource/client/DataSourceClientFactoryTest.java new file mode 100644 index 00000000000..3db0e8c0fcf --- /dev/null +++ b/direct-query-core/src/test/java/org/opensearch/sql/datasource/client/DataSourceClientFactoryTest.java @@ -0,0 +1,177 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.datasource.client; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertTrue; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import com.google.common.collect.ImmutableList; +import java.util.HashMap; +import java.util.Map; +import org.junit.Before; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.mockito.Mock; +import org.mockito.junit.MockitoJUnitRunner; +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.client.PrometheusClient; +import org.opensearch.sql.prometheus.utils.PrometheusClientUtils; + +@RunWith(MockitoJUnitRunner.class) +public class DataSourceClientFactoryTest { + + @Mock private DataSourceService dataSourceService; + + @Mock private Settings settings; + + private DataSourceClientFactory dataSourceClientFactory; + + @Before + public void setUp() { + when(settings.getSettingValue(Settings.Key.DATASOURCES_URI_HOSTS_DENY_LIST)) + .thenReturn(ImmutableList.of("http://localhost:9200")); + dataSourceClientFactory = new DataSourceClientFactory(dataSourceService, settings); + } + + @Test + public void testCreatePrometheusClientSuccessful() { + // Setup + String dataSourceName = "prometheusDataSource"; + Map properties = new HashMap<>(); + properties.put(PrometheusClientUtils.PROMETHEUS_URI, "http://prometheus:9090"); + + DataSourceMetadata metadata = + new DataSourceMetadata.Builder() + .setName(dataSourceName) + .setConnector(DataSourceType.PROMETHEUS) + .setProperties(properties) + .build(); + + when(dataSourceService.dataSourceExists(dataSourceName)).thenReturn(true); + when(dataSourceService.verifyDataSourceAccessAndGetRawMetadata(dataSourceName, null)) + .thenReturn(metadata); + + // Test + PrometheusClient client = dataSourceClientFactory.createClient(dataSourceName); + + // Verify + assertNotNull("Client should not be null", client); + verify(dataSourceService).dataSourceExists(dataSourceName); + verify(dataSourceService).verifyDataSourceAccessAndGetRawMetadata(dataSourceName, null); + } + + @Test(expected = DataSourceClientException.class) + public void testCreateClientForNonexistentDataSource() { + // Setup + String dataSourceName = "nonExistent"; + when(dataSourceService.dataSourceExists(dataSourceName)).thenReturn(false); + + // Test - should throw exception + dataSourceClientFactory.createClient(dataSourceName); + } + + @Test(expected = DataSourceClientException.class) + public void testCreateClientForUnsupportedDataSourceType() { + // Setup + String dataSourceName = "unsupportedType"; + DataSourceMetadata metadata = + new DataSourceMetadata.Builder() + .setName(dataSourceName) + .setConnector(DataSourceType.OPENSEARCH) // Unsupported type in current implementation + .setProperties(new HashMap<>()) + .build(); + + when(dataSourceService.dataSourceExists(dataSourceName)).thenReturn(true); + when(dataSourceService.verifyDataSourceAccessAndGetRawMetadata(dataSourceName, null)) + .thenReturn(metadata); + + // Test - should throw exception + dataSourceClientFactory.createClient(dataSourceName); + } + + @Test(expected = DataSourceClientException.class) + public void testCreateClientWrapsNonDataSourceClientException() { + // Setup + String dataSourceName = "exceptionSource"; + RuntimeException genericException = new RuntimeException("Generic error"); + + when(dataSourceService.dataSourceExists(dataSourceName)).thenReturn(true); + when(dataSourceService.verifyDataSourceAccessAndGetRawMetadata(dataSourceName, null)) + .thenThrow(genericException); + + // Test - should wrap the generic exception in a DataSourceClientException + dataSourceClientFactory.createClient(dataSourceName); + } + + @Test + public void testSuppressedWarningOnGenericTypeUsage() { + // This test verifies the @SuppressWarnings("unchecked") annotation is properly used + // by checking that the generic method works correctly with different return types + + // Setup for Prometheus client + String prometheusDs = "prometheusSource"; + Map properties = new HashMap<>(); + properties.put(PrometheusClientUtils.PROMETHEUS_URI, "http://prometheus:9090"); + + DataSourceMetadata metadata = + new DataSourceMetadata.Builder() + .setName(prometheusDs) + .setConnector(DataSourceType.PROMETHEUS) + .setProperties(properties) + .build(); + + when(dataSourceService.dataSourceExists(prometheusDs)).thenReturn(true); + when(dataSourceService.verifyDataSourceAccessAndGetRawMetadata(prometheusDs, null)) + .thenReturn(metadata); + + // Test that generic type inference works for explicit type parameter + PrometheusClient prometheusClient = dataSourceClientFactory.createClient(prometheusDs); + assertNotNull(prometheusClient); + + // Test with Object return type + Object genericClient = dataSourceClientFactory.createClient(prometheusDs); + assertTrue(genericClient instanceof PrometheusClient); + } + + @Test + public void testGetDataSourceTypeSuccessful() { + // Setup + String dataSourceName = "prometheusDataSource"; + DataSourceMetadata metadata = + new DataSourceMetadata.Builder() + .setName(dataSourceName) + .setConnector(DataSourceType.PROMETHEUS) + .build(); + + when(dataSourceService.dataSourceExists(dataSourceName)).thenReturn(true); + when(dataSourceService.getDataSourceMetadata(dataSourceName)).thenReturn(metadata); + + // Test + DataSourceType dataSourceType = dataSourceClientFactory.getDataSourceType(dataSourceName); + + // Verify + assertEquals(DataSourceType.PROMETHEUS, dataSourceType); + verify(dataSourceService).dataSourceExists(dataSourceName); + verify(dataSourceService).getDataSourceMetadata(dataSourceName); + } + + @Test(expected = DataSourceClientException.class) + public void testGetDataSourceTypeForNonexistentDataSource() { + // Setup + String dataSourceName = "nonExistent"; + when(dataSourceService.dataSourceExists(dataSourceName)).thenReturn(false); + + // Test - should throw exception + dataSourceClientFactory.getDataSourceType(dataSourceName); + } +} diff --git a/direct-query-core/src/test/java/org/opensearch/sql/datasource/query/QueryHandlerRegistryTest.java b/direct-query-core/src/test/java/org/opensearch/sql/datasource/query/QueryHandlerRegistryTest.java new file mode 100644 index 00000000000..aeaa2a0e139 --- /dev/null +++ b/direct-query-core/src/test/java/org/opensearch/sql/datasource/query/QueryHandlerRegistryTest.java @@ -0,0 +1,147 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.datasource.query; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +import java.util.Arrays; +import java.util.Collections; +import java.util.Optional; +import org.junit.jupiter.api.DisplayNameGeneration; +import org.junit.jupiter.api.DisplayNameGenerator; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.Mock; +import org.mockito.junit.jupiter.MockitoExtension; +import org.opensearch.sql.datasource.client.DataSourceClient; + +@DisplayNameGeneration(DisplayNameGenerator.ReplaceUnderscores.class) +@ExtendWith(MockitoExtension.class) +public class QueryHandlerRegistryTest { + + @Mock private QueryHandler mockHandler1; + @Mock private QueryHandler mockHandler2; + + @Test + void should_return_empty_when_registry_is_empty() { + QueryHandlerRegistry registry = new QueryHandlerRegistry(Collections.emptyList()); + Optional> result = registry.getQueryHandler(new TestClient()); + assertFalse(result.isPresent()); + } + + @Test + void should_return_handler_when_client_matches() { + // Setup + when(mockHandler1.getClientClass()).thenReturn(TestClient.class); + when(mockHandler1.canHandle(any(TestClient.class))).thenReturn(true); + QueryHandlerRegistry registry = + new QueryHandlerRegistry(Collections.singletonList(mockHandler1)); + + // Execute + Optional> result = registry.getQueryHandler(new TestClient()); + + // Verify + assertTrue(result.isPresent()); + assertEquals(mockHandler1, result.get()); + } + + @Test + void should_return_empty_when_handler_cannot_handle_client() { + // Setup + when(mockHandler1.getClientClass()).thenReturn(TestClient.class); + when(mockHandler1.canHandle(any(TestClient.class))).thenReturn(false); + QueryHandlerRegistry registry = + new QueryHandlerRegistry(Collections.singletonList(mockHandler1)); + + // Execute + Optional> result = registry.getQueryHandler(new TestClient()); + + // Verify + assertFalse(result.isPresent()); + } + + @Test + void should_return_correct_handler_when_multiple_handlers_exist() { + // Setup + when(mockHandler1.getClientClass()).thenReturn(TestClient.class); + when(mockHandler2.getClientClass()).thenReturn(AnotherTestClient.class); + when(mockHandler1.canHandle(any(TestClient.class))).thenReturn(true); + when(mockHandler2.canHandle(any(AnotherTestClient.class))).thenReturn(true); + QueryHandlerRegistry registry = + new QueryHandlerRegistry(Arrays.asList(mockHandler1, mockHandler2)); + + // Execute + Optional> result1 = registry.getQueryHandler(new TestClient()); + Optional> result2 = + registry.getQueryHandler(new AnotherTestClient()); + + // Verify + assertTrue(result1.isPresent()); + assertEquals(mockHandler1, result1.get()); + + assertTrue(result2.isPresent()); + assertEquals(mockHandler2, result2.get()); + } + + @Test + void should_handle_class_cast_exception_gracefully() { + // Setup - create a handler with incompatible client class + @SuppressWarnings("unchecked") + QueryHandler badHandler = mock(QueryHandler.class); + when(badHandler.getClientClass()).thenReturn((Class) IncompatibleDataSourceClient.class); + + // Setup regular handler + when(mockHandler1.getClientClass()).thenReturn(TestClient.class); + when(mockHandler1.canHandle(any(TestClient.class))).thenReturn(true); + + // Create registry with both handlers + QueryHandlerRegistry registry = + new QueryHandlerRegistry(Arrays.asList(badHandler, mockHandler1)); + + // Execute + Optional> result = registry.getQueryHandler(new TestClient()); + + // Verify - should skip the incompatible handler and find the compatible one + assertTrue(result.isPresent()); + assertEquals(mockHandler1, result.get()); + } + + @Test + void should_catch_class_cast_exception_during_can_handle_check() { + // Setup - create a handler that throws ClassCastException when canHandle is called + @SuppressWarnings("unchecked") + QueryHandler problematicHandler = mock(QueryHandler.class); + when(problematicHandler.getClientClass()).thenReturn(TestClient.class); + when(problematicHandler.canHandle(any(TestClient.class))).thenThrow(ClassCastException.class); + + // Setup regular handler + when(mockHandler1.getClientClass()).thenReturn(TestClient.class); + when(mockHandler1.canHandle(any(TestClient.class))).thenReturn(true); + + // Create registry with both handlers + QueryHandlerRegistry registry = + new QueryHandlerRegistry(Arrays.asList(problematicHandler, mockHandler1)); + + // Execute + Optional> result = registry.getQueryHandler(new TestClient()); + + // Verify - should skip the handler that throws exception and find the working one + assertTrue(result.isPresent()); + assertEquals(mockHandler1, result.get()); + } + + // Test client classes + private static class TestClient implements DataSourceClient {} + + private static class AnotherTestClient implements DataSourceClient {} + + private static class IncompatibleDataSourceClient implements DataSourceClient {} +} diff --git a/direct-query-core/src/test/java/org/opensearch/sql/directquery/DirectQueryExecutorServiceImplTest.java b/direct-query-core/src/test/java/org/opensearch/sql/directquery/DirectQueryExecutorServiceImplTest.java new file mode 100644 index 00000000000..2df9b4f9de5 --- /dev/null +++ b/direct-query-core/src/test/java/org/opensearch/sql/directquery/DirectQueryExecutorServiceImplTest.java @@ -0,0 +1,285 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.directquery; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.when; + +import java.io.IOException; +import java.util.HashMap; +import java.util.Map; +import java.util.Optional; +import org.json.JSONObject; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.Mock; +import org.mockito.junit.jupiter.MockitoExtension; +import org.opensearch.sql.datasource.client.DataSourceClientFactory; +import org.opensearch.sql.datasource.client.exceptions.DataSourceClientException; +import org.opensearch.sql.datasource.model.DataSourceType; +import org.opensearch.sql.datasource.query.QueryHandler; +import org.opensearch.sql.datasource.query.QueryHandlerRegistry; +import org.opensearch.sql.directquery.rest.model.DirectQueryResourceType; +import org.opensearch.sql.directquery.rest.model.ExecuteDirectQueryRequest; +import org.opensearch.sql.directquery.rest.model.ExecuteDirectQueryResponse; +import org.opensearch.sql.directquery.rest.model.GetDirectQueryResourcesRequest; +import org.opensearch.sql.directquery.rest.model.GetDirectQueryResourcesResponse; +import org.opensearch.sql.prometheus.client.PrometheusClient; +import org.opensearch.sql.prometheus.model.PrometheusOptions; +import org.opensearch.sql.prometheus.model.PrometheusQueryType; +import org.opensearch.sql.spark.rest.model.LangType; + +@ExtendWith(MockitoExtension.class) +public class DirectQueryExecutorServiceImplTest { + + @Mock private DataSourceClientFactory dataSourceClientFactory; + + @Mock private QueryHandlerRegistry queryHandlerRegistry; + + @Mock private QueryHandler queryHandler; + + @Mock private PrometheusClient prometheusClient; + + private DirectQueryExecutorServiceImpl executorService; + + @BeforeEach + public void setUp() { + executorService = + new DirectQueryExecutorServiceImpl(dataSourceClientFactory, queryHandlerRegistry); + } + + @Test + public void testExecuteDirectQuerySuccessful() throws IOException { + // Setup + String dataSource = "prometheusDataSource"; + String query = "up"; + String sessionId = "test-session-id"; + + ExecuteDirectQueryRequest request = new ExecuteDirectQueryRequest(); + request.setDataSources(dataSource); + request.setQuery(query); + request.setLanguage(LangType.PROMQL); + request.setSessionId(sessionId); + + PrometheusOptions options = new PrometheusOptions(); + options.setQueryType(PrometheusQueryType.INSTANT); + options.setTime("1609459200"); + request.setPrometheusOptions(options); + + when(dataSourceClientFactory.createClient(dataSource)).thenReturn(prometheusClient); + when(dataSourceClientFactory.getDataSourceType(dataSource)) + .thenReturn(DataSourceType.PROMETHEUS); + when(queryHandlerRegistry.getQueryHandler(prometheusClient)) + .thenReturn(Optional.of(queryHandler)); + + String expectedResult = "{\"status\":\"success\",\"data\":{\"resultType\":\"vector\"}}"; + when(queryHandler.executeQuery(eq(prometheusClient), eq(request))).thenReturn(expectedResult); + + // Test + ExecuteDirectQueryResponse response = executorService.executeDirectQuery(request); + + // Verify + assertNotNull(response); + assertNotNull(response.getQueryId()); + assertEquals(expectedResult, response.getResult()); + assertEquals(sessionId, response.getSessionId()); + } + + @Test + public void testExecuteDirectQueryWithUnregisteredHandler() { + // Setup + String dataSource = "unsupportedDataSource"; + String query = "up"; + + ExecuteDirectQueryRequest request = new ExecuteDirectQueryRequest(); + request.setDataSources(dataSource); + request.setQuery(query); + + when(dataSourceClientFactory.createClient(dataSource)).thenReturn(prometheusClient); + when(dataSourceClientFactory.getDataSourceType(dataSource)) + .thenReturn(DataSourceType.PROMETHEUS); + when(queryHandlerRegistry.getQueryHandler(prometheusClient)).thenReturn(Optional.empty()); + + // Test + ExecuteDirectQueryResponse response = executorService.executeDirectQuery(request); + + // Verify + assertNotNull(response); + assertNotNull(response.getQueryId()); + JSONObject result = new JSONObject(response.getResult()); + assertTrue(result.has("error")); + assertEquals("Unsupported data source type", result.getString("error")); + } + + @Test + public void testExecuteDirectQueryWithClientError() throws IOException { + // Setup + String dataSource = "errorDataSource"; + String query = "up"; + + ExecuteDirectQueryRequest request = new ExecuteDirectQueryRequest(); + request.setDataSources(dataSource); + request.setQuery(query); + + when(dataSourceClientFactory.getDataSourceType(dataSource)) + .thenReturn(DataSourceType.PROMETHEUS); + when(dataSourceClientFactory.createClient(dataSource)) + .thenThrow(new DataSourceClientException("Failed to create client")); + + // Test + ExecuteDirectQueryResponse response = executorService.executeDirectQuery(request); + + // Verify + assertNotNull(response); + assertNotNull(response.getQueryId()); + JSONObject result = new JSONObject(response.getResult()); + assertTrue(result.has("error")); + assertEquals("Failed to create client", result.getString("error")); + } + + @Test + public void testExecuteDirectQueryWithExecutionError() throws IOException { + // Setup + String dataSource = "prometheusDataSource"; + String query = "up"; + + ExecuteDirectQueryRequest request = new ExecuteDirectQueryRequest(); + request.setDataSources(dataSource); + request.setQuery(query); + + when(dataSourceClientFactory.createClient(dataSource)).thenReturn(prometheusClient); + when(dataSourceClientFactory.getDataSourceType(dataSource)) + .thenReturn(DataSourceType.PROMETHEUS); + when(queryHandlerRegistry.getQueryHandler(prometheusClient)) + .thenReturn(Optional.of(queryHandler)); + when(queryHandler.executeQuery(eq(prometheusClient), eq(request))) + .thenThrow(new IOException("Query execution failed")); + + // Test + ExecuteDirectQueryResponse response = executorService.executeDirectQuery(request); + + // Verify + assertNotNull(response); + assertNotNull(response.getQueryId()); + JSONObject result = new JSONObject(response.getResult()); + assertTrue(result.has("error")); + assertTrue(result.getString("error").contains("Error executing query")); + } + + @Test + public void testGetDirectQueryResourcesSuccessful() throws IOException { + // Setup + String dataSource = "prometheusDataSource"; + + GetDirectQueryResourcesRequest request = new GetDirectQueryResourcesRequest(); + request.setDataSource(dataSource); + request.setResourceType(DirectQueryResourceType.LABELS); + + when(dataSourceClientFactory.createClient(dataSource)).thenReturn(prometheusClient); + when(queryHandlerRegistry.getQueryHandler(prometheusClient)) + .thenReturn(Optional.of(queryHandler)); + + Map resourcesMap = new HashMap<>(); + resourcesMap.put("labels", new String[] {"job", "instance"}); + + // Use the factory method from GetDirectQueryResourcesResponse + @SuppressWarnings("unchecked") + GetDirectQueryResourcesResponse> expectedResponse = + GetDirectQueryResourcesResponse.withMap(resourcesMap); + + // Use a raw type for the mock setup to avoid generic type issues + when(queryHandler.getResources(eq(prometheusClient), eq(request))) + .thenReturn((GetDirectQueryResourcesResponse) expectedResponse); + + // Test + GetDirectQueryResourcesResponse response = executorService.getDirectQueryResources(request); + + // Verify + assertNotNull(response); + assertEquals( + expectedResponse.getData(), ((GetDirectQueryResourcesResponse) response).getData()); + } + + @Test + public void testGetDirectQueryResourcesWithUnregisteredHandler() { + // Setup + String dataSource = "unsupportedDataSource"; + + GetDirectQueryResourcesRequest request = new GetDirectQueryResourcesRequest(); + request.setDataSource(dataSource); + request.setResourceType(DirectQueryResourceType.LABELS); + + when(dataSourceClientFactory.createClient(dataSource)).thenReturn(prometheusClient); + when(queryHandlerRegistry.getQueryHandler(prometheusClient)).thenReturn(Optional.empty()); + + // Test - should throw exception + IllegalArgumentException exception = + assertThrows( + IllegalArgumentException.class, () -> executorService.getDirectQueryResources(request)); + + // Verify exception message + assertEquals("Unsupported data source type: " + dataSource, exception.getMessage()); + } + + @Test + public void testGetDirectQueryResourcesWithIOError() throws IOException { + // Setup + String dataSource = "prometheusDataSource"; + + GetDirectQueryResourcesRequest request = new GetDirectQueryResourcesRequest(); + request.setDataSource(dataSource); + request.setResourceType(DirectQueryResourceType.LABELS); + + when(dataSourceClientFactory.createClient(dataSource)).thenReturn(prometheusClient); + when(queryHandlerRegistry.getQueryHandler(prometheusClient)) + .thenReturn(Optional.of(queryHandler)); + when(queryHandler.getResources(eq(prometheusClient), eq(request))) + .thenThrow(new IOException("Failed to get resources")); + + // Test + DataSourceClientException exception = + assertThrows( + DataSourceClientException.class, + () -> executorService.getDirectQueryResources(request)); + + // Verify + assertNotNull(exception); + assertEquals( + "Error retrieving resources for data source type: " + dataSource, exception.getMessage()); + assertTrue(exception.getCause() instanceof IOException); + assertEquals("Failed to get resources", exception.getCause().getMessage()); + } + + @Test + public void testExecuteDirectQueryWithNullClient() { + // Setup + String dataSource = "nullClientDataSource"; + String query = "up"; + + ExecuteDirectQueryRequest request = new ExecuteDirectQueryRequest(); + request.setDataSources(dataSource); + request.setQuery(query); + + when(dataSourceClientFactory.createClient(dataSource)).thenReturn(null); + when(dataSourceClientFactory.getDataSourceType(dataSource)) + .thenReturn(DataSourceType.PROMETHEUS); + + // Test + ExecuteDirectQueryResponse response = executorService.executeDirectQuery(request); + + // Verify + assertNotNull(response); + assertNotNull(response.getQueryId()); + JSONObject result = new JSONObject(response.getResult()); + assertTrue(result.has("error")); + assertEquals("Unsupported data source type", result.getString("error")); + } +} diff --git a/direct-query-core/src/test/java/org/opensearch/sql/directquery/validator/DirectQueryRequestValidatorTest.java b/direct-query-core/src/test/java/org/opensearch/sql/directquery/validator/DirectQueryRequestValidatorTest.java new file mode 100644 index 00000000000..0c6fec31dcc --- /dev/null +++ b/direct-query-core/src/test/java/org/opensearch/sql/directquery/validator/DirectQueryRequestValidatorTest.java @@ -0,0 +1,331 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.directquery.validator; + +import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; + +import org.junit.jupiter.api.Test; +import org.opensearch.sql.directquery.rest.model.ExecuteDirectQueryRequest; +import org.opensearch.sql.prometheus.model.PrometheusOptions; +import org.opensearch.sql.prometheus.model.PrometheusQueryType; +import org.opensearch.sql.spark.rest.model.LangType; + +public class DirectQueryRequestValidatorTest { + + @Test + public void testValidateNullRequest() { + IllegalArgumentException exception = + assertThrows( + IllegalArgumentException.class, + () -> DirectQueryRequestValidator.validateRequest(null)); + assertEquals("Request cannot be null", exception.getMessage()); + } + + @Test + public void testValidateRequestWithNullDataSource() { + ExecuteDirectQueryRequest request = new ExecuteDirectQueryRequest(); + request.setQuery("up"); + request.setLanguage(LangType.PROMQL); + + IllegalArgumentException exception = + assertThrows( + IllegalArgumentException.class, + () -> DirectQueryRequestValidator.validateRequest(request)); + assertEquals("Datasource is required", exception.getMessage()); + } + + @Test + public void testValidateRequestWithEmptyDataSource() { + ExecuteDirectQueryRequest request = new ExecuteDirectQueryRequest(); + request.setDataSources(""); + request.setQuery("up"); + request.setLanguage(LangType.PROMQL); + + IllegalArgumentException exception = + assertThrows( + IllegalArgumentException.class, + () -> DirectQueryRequestValidator.validateRequest(request)); + assertEquals("Datasource is required", exception.getMessage()); + } + + @Test + public void testValidateRequestWithNullQuery() { + ExecuteDirectQueryRequest request = new ExecuteDirectQueryRequest(); + request.setDataSources("prometheus"); + request.setLanguage(LangType.PROMQL); + + IllegalArgumentException exception = + assertThrows( + IllegalArgumentException.class, + () -> DirectQueryRequestValidator.validateRequest(request)); + assertEquals("Query is required", exception.getMessage()); + } + + @Test + public void testValidateRequestWithEmptyQuery() { + ExecuteDirectQueryRequest request = new ExecuteDirectQueryRequest(); + request.setDataSources("prometheus"); + request.setQuery(""); + request.setLanguage(LangType.PROMQL); + + IllegalArgumentException exception = + assertThrows( + IllegalArgumentException.class, + () -> DirectQueryRequestValidator.validateRequest(request)); + assertEquals("Query is required", exception.getMessage()); + } + + @Test + public void testValidateRequestWithNullLanguage() { + ExecuteDirectQueryRequest request = new ExecuteDirectQueryRequest(); + request.setDataSources("prometheus"); + request.setQuery("up"); + + IllegalArgumentException exception = + assertThrows( + IllegalArgumentException.class, + () -> DirectQueryRequestValidator.validateRequest(request)); + assertEquals("Language type is required", exception.getMessage()); + } + + @Test + public void testValidateRequestWithSqlLanguage() { + ExecuteDirectQueryRequest request = new ExecuteDirectQueryRequest(); + request.setDataSources("prometheus"); + request.setLanguage(LangType.SQL); + request.setQuery("select 1"); + + assertDoesNotThrow(() -> DirectQueryRequestValidator.validateRequest(request)); + } + + @Test + public void testValidatePromQLRequestWithoutOptions() { + ExecuteDirectQueryRequest request = new ExecuteDirectQueryRequest(); + request.setDataSources("prometheus"); + request.setQuery("up"); + request.setLanguage(LangType.PROMQL); + + IllegalArgumentException exception = + assertThrows( + IllegalArgumentException.class, + () -> DirectQueryRequestValidator.validateRequest(request)); + assertEquals("Prometheus options are required for PROMQL queries", exception.getMessage()); + } + + @Test + public void testValidatePromQLRangeQueryMissingStartEnd() { + ExecuteDirectQueryRequest request = new ExecuteDirectQueryRequest(); + request.setDataSources("prometheus"); + request.setQuery("up"); + request.setLanguage(LangType.PROMQL); + + PrometheusOptions options = new PrometheusOptions(); + options.setQueryType(PrometheusQueryType.RANGE); + options.setStep("15s"); + request.setPrometheusOptions(options); + + IllegalArgumentException exception = + assertThrows( + IllegalArgumentException.class, + () -> DirectQueryRequestValidator.validateRequest(request)); + assertEquals("Start and end times are required for range queries", exception.getMessage()); + } + + @Test + public void testValidatePromQLRangeQueryMissingEnd() { + ExecuteDirectQueryRequest request = new ExecuteDirectQueryRequest(); + request.setDataSources("prometheus"); + request.setQuery("up"); + request.setLanguage(LangType.PROMQL); + + PrometheusOptions options = new PrometheusOptions(); + options.setQueryType(PrometheusQueryType.RANGE); + options.setStep("15s"); + options.setStart("now"); + request.setPrometheusOptions(options); + + IllegalArgumentException exception = + assertThrows( + IllegalArgumentException.class, + () -> DirectQueryRequestValidator.validateRequest(request)); + assertEquals("Start and end times are required for range queries", exception.getMessage()); + } + + @Test + public void testValidatePromQLRangeQueryMissingStep() { + ExecuteDirectQueryRequest request = new ExecuteDirectQueryRequest(); + request.setDataSources("prometheus"); + request.setQuery("up"); + request.setLanguage(LangType.PROMQL); + + PrometheusOptions options = new PrometheusOptions(); + options.setQueryType(PrometheusQueryType.RANGE); + options.setStart("1609459200"); + options.setEnd("1609545600"); + // Missing step + request.setPrometheusOptions(options); + + IllegalArgumentException exception = + assertThrows( + IllegalArgumentException.class, + () -> DirectQueryRequestValidator.validateRequest(request)); + assertEquals("Step parameter is required for range queries", exception.getMessage()); + } + + @Test + public void testValidatePromQLRangeQueryEmptyStep() { + ExecuteDirectQueryRequest request = new ExecuteDirectQueryRequest(); + request.setDataSources("prometheus"); + request.setQuery("up"); + request.setLanguage(LangType.PROMQL); + + PrometheusOptions options = new PrometheusOptions(); + options.setQueryType(PrometheusQueryType.RANGE); + options.setStart("1609459200"); + options.setEnd("1609545600"); + options.setStep(""); + request.setPrometheusOptions(options); + + IllegalArgumentException exception = + assertThrows( + IllegalArgumentException.class, + () -> DirectQueryRequestValidator.validateRequest(request)); + assertEquals("Step parameter is required for range queries", exception.getMessage()); + } + + @Test + public void testValidatePromQLRangeQueryInvalidTimeFormat() { + ExecuteDirectQueryRequest request = new ExecuteDirectQueryRequest(); + request.setDataSources("prometheus"); + request.setQuery("up"); + request.setLanguage(LangType.PROMQL); + + PrometheusOptions options = new PrometheusOptions(); + options.setQueryType(PrometheusQueryType.RANGE); + options.setStart("invalid-time"); + options.setEnd("1609545600"); + options.setStep("15s"); + request.setPrometheusOptions(options); + + IllegalArgumentException exception = + assertThrows( + IllegalArgumentException.class, + () -> DirectQueryRequestValidator.validateRequest(request)); + assertEquals( + "Invalid time format: start and end must be numeric timestamps", exception.getMessage()); + } + + @Test + public void testValidatePromQLRangeQueryEndBeforeStart() { + ExecuteDirectQueryRequest request = new ExecuteDirectQueryRequest(); + request.setDataSources("prometheus"); + request.setQuery("up"); + request.setLanguage(LangType.PROMQL); + + PrometheusOptions options = new PrometheusOptions(); + options.setQueryType(PrometheusQueryType.RANGE); + options.setStart("1609545600"); // End time + options.setEnd("1609459200"); // Start time (which is earlier) + options.setStep("15s"); + request.setPrometheusOptions(options); + + IllegalArgumentException exception = + assertThrows( + IllegalArgumentException.class, + () -> DirectQueryRequestValidator.validateRequest(request)); + assertEquals("End time must be after start time", exception.getMessage()); + } + + @Test + public void testValidatePromQLInstantQueryMissingTime() { + ExecuteDirectQueryRequest request = new ExecuteDirectQueryRequest(); + request.setDataSources("prometheus"); + request.setQuery("up"); + request.setLanguage(LangType.PROMQL); + + PrometheusOptions options = new PrometheusOptions(); + options.setQueryType(PrometheusQueryType.INSTANT); + // Missing time + request.setPrometheusOptions(options); + + IllegalArgumentException exception = + assertThrows( + IllegalArgumentException.class, + () -> DirectQueryRequestValidator.validateRequest(request)); + assertEquals("Time parameter is required for instant queries", exception.getMessage()); + } + + @Test + public void testValidatePromQLInstantQueryInvalidTimeFormat() { + ExecuteDirectQueryRequest request = new ExecuteDirectQueryRequest(); + request.setDataSources("prometheus"); + request.setQuery("up"); + request.setLanguage(LangType.PROMQL); + + PrometheusOptions options = new PrometheusOptions(); + options.setQueryType(PrometheusQueryType.INSTANT); + options.setTime("invalid-time"); + request.setPrometheusOptions(options); + + IllegalArgumentException exception = + assertThrows( + IllegalArgumentException.class, + () -> DirectQueryRequestValidator.validateRequest(request)); + assertEquals("Invalid time format: time must be a numeric timestamp", exception.getMessage()); + } + + @Test + public void testValidatePromQLQueryNullQueryType() { + ExecuteDirectQueryRequest request = new ExecuteDirectQueryRequest(); + request.setDataSources("prometheus"); + request.setQuery("up"); + request.setLanguage(LangType.PROMQL); + + PrometheusOptions options = new PrometheusOptions(); + options.setQueryType(null); + request.setPrometheusOptions(options); + + IllegalArgumentException exception = + assertThrows( + IllegalArgumentException.class, + () -> DirectQueryRequestValidator.validateRequest(request)); + assertEquals("Prometheus options are required for PROMQL queries", exception.getMessage()); + } + + @Test + public void testValidateValidPromQLRangeQuery() { + ExecuteDirectQueryRequest request = new ExecuteDirectQueryRequest(); + request.setDataSources("prometheus"); + request.setQuery("up"); + request.setLanguage(LangType.PROMQL); + + PrometheusOptions options = new PrometheusOptions(); + options.setQueryType(PrometheusQueryType.RANGE); + options.setStart("1609459200"); // 2021-01-01 + options.setEnd("1609545600"); // 2021-01-02 + options.setStep("15s"); + request.setPrometheusOptions(options); + + assertDoesNotThrow(() -> DirectQueryRequestValidator.validateRequest(request)); + } + + @Test + public void testValidateValidPromQLInstantQuery() { + ExecuteDirectQueryRequest request = new ExecuteDirectQueryRequest(); + request.setDataSources("prometheus"); + request.setQuery("up"); + request.setLanguage(LangType.PROMQL); + + PrometheusOptions options = new PrometheusOptions(); + options.setQueryType(PrometheusQueryType.INSTANT); + options.setTime("1609459200"); // 2021-01-01 + request.setPrometheusOptions(options); + + assertDoesNotThrow(() -> DirectQueryRequestValidator.validateRequest(request)); + } +} diff --git a/direct-query-core/src/test/java/org/opensearch/sql/prometheus/client/PrometheusClientImplTest.java b/direct-query-core/src/test/java/org/opensearch/sql/prometheus/client/PrometheusClientImplTest.java new file mode 100644 index 00000000000..8c2444c6590 --- /dev/null +++ b/direct-query-core/src/test/java/org/opensearch/sql/prometheus/client/PrometheusClientImplTest.java @@ -0,0 +1,578 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.prometheus.client; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.junit.jupiter.api.Assertions.fail; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.doAnswer; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.spy; +import static org.mockito.Mockito.when; + +import java.io.IOException; +import java.net.URI; +import java.util.HashMap; +import java.util.List; +import okhttp3.Call; +import okhttp3.OkHttpClient; +import okhttp3.Protocol; +import okhttp3.Request; +import okhttp3.Response; +import okhttp3.mockwebserver.MockResponse; +import okhttp3.mockwebserver.MockWebServer; +import org.json.JSONArray; +import org.json.JSONObject; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.opensearch.sql.prometheus.exception.PrometheusClientException; + +public class PrometheusClientImplTest { + + private MockWebServer mockWebServer; + private PrometheusClientImpl client; + + @BeforeEach + public void setUp() throws IOException { + mockWebServer = new MockWebServer(); + mockWebServer.start(); + OkHttpClient httpClient = new OkHttpClient.Builder().build(); + client = + new PrometheusClientImpl( + httpClient, + URI.create(String.format("http://%s:%s", "localhost", mockWebServer.getPort()))); + } + + @AfterEach + public void tearDown() throws IOException { + mockWebServer.shutdown(); + } + + @Test + public void testQueryRange() throws IOException { + // Setup + String successResponse = + "{\"status\":\"success\",\"data\":{\"resultType\":\"matrix\",\"result\":[{\"metric\":{\"__name__\":\"up\",\"job\":\"prometheus\",\"instance\":\"localhost:9090\"},\"values\":[[1435781430.781,\"1\"],[1435781445.781,\"1\"],[1435781460.781,\"1\"]]}]}}"; + mockWebServer.enqueue(new MockResponse().setBody(successResponse)); + + // Test + JSONObject result = client.queryRange("up", 1435781430L, 1435781460L, "15s", 100, 30); + + // Verify + assertNotNull(result); + // assertEquals("success", result.getString("status")); + // JSONObject data = result.getJSONObject("data"); + assertEquals("matrix", result.getString("resultType")); + JSONArray resultArray = result.getJSONArray("result"); + assertEquals(1, resultArray.length()); + JSONObject metric = resultArray.getJSONObject(0).getJSONObject("metric"); + assertEquals("up", metric.getString("__name__")); + } + + @Test + public void testQueryRangeSimpleOverload() throws IOException { + // Setup + String successResponse = + "{\"status\":\"success\",\"data\":{\"resultType\":\"matrix\",\"result\":[{\"metric\":{\"__name__\":\"up\",\"job\":\"prometheus\",\"instance\":\"localhost:9090\"},\"values\":[[1435781430.781,\"1\"],[1435781445.781,\"1\"],[1435781460.781,\"1\"]]}]}}"; + mockWebServer.enqueue(new MockResponse().setBody(successResponse)); + + // Test + JSONObject result = client.queryRange("up", 1435781430L, 1435781460L, "15s"); + + // Verify + assertNotNull(result); + // assertEquals("success", result.getString("status")); + // JSONObject data = result.getJSONObject("data"); + assertEquals("matrix", result.getString("resultType")); + JSONArray resultArray = result.getJSONArray("result"); + assertEquals(1, resultArray.length()); + JSONObject metric = resultArray.getJSONObject(0).getJSONObject("metric"); + assertEquals("up", metric.getString("__name__")); + } + + @Test + public void testQueryRangeWith2xxStatusAndError() { + // Setup + String errorResponse = "{\"status\":\"error\",\"error\":\"Error\"}"; + mockWebServer.enqueue(new MockResponse().setBody(errorResponse).setResponseCode(200)); + + // Test & Verify + PrometheusClientException exception = + assertThrows( + PrometheusClientException.class, + () -> client.queryRange("up", 1435781430L, 1435781460L, "15s")); + assertEquals("Error", exception.getMessage()); + } + + @Test + public void testQueryRangeWithNonJsonResponse() { + // Setup + String nonJsonResponse = "Not a JSON response"; + mockWebServer.enqueue(new MockResponse().setBody(nonJsonResponse).setResponseCode(200)); + + // Test & Verify + PrometheusClientException exception = + assertThrows( + PrometheusClientException.class, + () -> client.queryRange("up", 1435781430L, 1435781460L, "15s")); + assertTrue( + exception + .getMessage() + .contains( + "Prometheus returned unexpected body, please verify your prometheus server" + + " setup.")); + } + + @Test + public void testQueryRangeWithNon2xxError() { + // Setup + mockWebServer.enqueue(new MockResponse().setResponseCode(400).setBody("Mock Error")); + + // Test & Verify + PrometheusClientException exception = + assertThrows( + PrometheusClientException.class, + () -> client.queryRange("up", 1435781430L, 1435781460L, "15s")); + assertTrue( + exception + .getMessage() + .contains( + "Request to Prometheus is Unsuccessful with code: 400. Error details: Mock Error")); + } + + /** response.body() is @Nullable, to test the null path we need to create a spy client. */ + @Test + public void testQueryRangeWithNon2xxErrorNullBody() { + Request dummyRequest = new Request.Builder().url(mockWebServer.url("/")).build(); + Response nullBodyResponse = + new Response.Builder() + .request(dummyRequest) + .protocol(Protocol.HTTP_1_1) + .code(400) + .message("Bad Request") + .body(null) + .build(); + OkHttpClient spyClient = spy(new OkHttpClient()); + Call mockCall = mock(Call.class); + try { + when(mockCall.execute()).thenReturn(nullBodyResponse); + } catch (IOException e) { + fail("Unexpected IOException"); + } + doAnswer(invocation -> mockCall).when(spyClient).newCall(any(Request.class)); + + PrometheusClientImpl nullBodyClient = + new PrometheusClientImpl( + spyClient, + URI.create(String.format("http://%s:%s", "localhost", mockWebServer.getPort()))); + + PrometheusClientException exception = + assertThrows( + PrometheusClientException.class, + () -> nullBodyClient.queryRange("up", 1435781430L, 1435781460L, "15s")); + assertTrue( + exception + .getMessage() + .contains( + "Request to Prometheus is Unsuccessful with code: 400. Error details: No response" + + " body")); + } + + @Test + public void testQuery() throws IOException { + // Setup + String successResponse = + "{\"status\":\"success\",\"data\":{\"resultType\":\"vector\",\"result\":[{\"metric\":{\"__name__\":\"up\",\"job\":\"prometheus\",\"instance\":\"localhost:9090\"},\"value\":[1435781460.781,\"1\"]}]}}"; + mockWebServer.enqueue(new MockResponse().setBody(successResponse)); + + // Test + JSONObject result = client.query("up", 1435781460L, 100, 30); + + // Verify + assertNotNull(result); + assertEquals("success", result.getString("status")); + JSONObject data = result.getJSONObject("data"); + assertEquals("vector", data.getString("resultType")); + JSONArray resultArray = data.getJSONArray("result"); + assertEquals(1, resultArray.length()); + JSONObject metric = resultArray.getJSONObject(0).getJSONObject("metric"); + assertEquals("up", metric.getString("__name__")); + } + + @Test + public void testQueryWithNullParams() throws IOException { + // Setup + String successResponse = + "{\"status\":\"success\",\"data\":{\"resultType\":\"vector\",\"result\":[{\"metric\":{\"__name__\":\"up\",\"job\":\"prometheus\",\"instance\":\"localhost:9090\"},\"value\":[1435781460.781,\"1\"]}]}}"; + mockWebServer.enqueue(new MockResponse().setBody(successResponse)); + + // Test + JSONObject result = client.query("up", null, null, null); + + // Verify + assertNotNull(result); + assertEquals("success", result.getString("status")); + JSONObject data = result.getJSONObject("data"); + assertEquals("vector", data.getString("resultType")); + JSONArray resultArray = data.getJSONArray("result"); + assertEquals(1, resultArray.length()); + JSONObject metric = resultArray.getJSONObject(0).getJSONObject("metric"); + assertEquals("up", metric.getString("__name__")); + } + + @Test + public void testGetLabels() throws IOException { + // Setup + String successResponse = "{\"status\":\"success\",\"data\":[\"job\",\"instance\",\"version\"]}"; + mockWebServer.enqueue(new MockResponse().setBody(successResponse)); + + // Test + List labels = client.getLabels(new HashMap<>()); + + // Verify + assertNotNull(labels); + assertEquals(3, labels.size()); + assertTrue(labels.contains("job")); + assertTrue(labels.contains("instance")); + assertTrue(labels.contains("version")); + } + + @Test + public void testGetLabelsByMetricName() throws IOException { + // Setup + String successResponse = + "{\"status\":\"success\",\"data\":[\"job\",\"instance\",\"__name__\"]}"; + mockWebServer.enqueue(new MockResponse().setBody(successResponse)); + + // Test + List labels = client.getLabels("http_requests_total"); + + // Verify + assertNotNull(labels); + assertEquals(2, labels.size()); + assertTrue(labels.contains("job")); + assertTrue(labels.contains("instance")); + // __name__ should be filtered out by the implementation + assertFalse(labels.contains("__name__")); + } + + @Test + public void testGetLabel() throws IOException { + // Setup + String successResponse = "{\"status\":\"success\",\"data\":[\"prometheus\",\"node-exporter\"]}"; + mockWebServer.enqueue(new MockResponse().setBody(successResponse)); + + // Test + List labelValues = client.getLabel("job", new HashMap<>()); + + // Verify + assertNotNull(labelValues); + assertEquals(2, labelValues.size()); + assertTrue(labelValues.contains("prometheus")); + assertTrue(labelValues.contains("node-exporter")); + } + + @Test + public void testQueryExemplars() throws IOException { + // Setup + String successResponse = + "{\"status\":\"success\",\"data\":[{\"seriesLabels\":{\"__name__\":\"http_request_duration_seconds_bucket\",\"handler\":\"/api/v1/query_range\",\"le\":\"1\"},\"exemplars\":[{\"labels\":{\"traceID\":\"19a801c37fb022d6\"},\"value\":0.207396059,\"timestamp\":1659284721.762}]}]}"; + mockWebServer.enqueue(new MockResponse().setBody(successResponse)); + + // Test + JSONArray result = + client.queryExemplars("http_request_duration_seconds_bucket", 1659284721L, 1659284722L); + + // Verify + assertNotNull(result); + assertEquals(1, result.length()); + JSONObject exemplarData = result.getJSONObject(0); + assertTrue(exemplarData.has("seriesLabels")); + assertTrue(exemplarData.has("exemplars")); + } + + @Test + public void testGetAllMetricsWithParams() throws IOException { + // Setup + String successResponse = + "{\"status\":\"success\",\"data\":{\"go_gc_duration_seconds\":[{\"type\":\"histogram\",\"help\":\"A" + + " summary of the pause duration of garbage collection cycles.\",\"unit\":\"\"}]}}"; + mockWebServer.enqueue(new MockResponse().setBody(successResponse)); + + // Test + HashMap params = new HashMap<>(); + params.put("limit", "100"); + var result = client.getAllMetrics(params); + + // Verify + assertNotNull(result); + assertEquals(1, result.size()); + assertTrue(result.containsKey("go_gc_duration_seconds")); + assertEquals(1, result.get("go_gc_duration_seconds").size()); + assertEquals("histogram", result.get("go_gc_duration_seconds").getFirst().getType()); + assertEquals( + "A summary of the pause duration of garbage collection cycles.", + result.get("go_gc_duration_seconds").getFirst().getHelp()); + } + + @Test + public void testGetAllMetrics() throws IOException { + // Setup + String successResponse = + "{\"status\":\"success\",\"data\":{\"http_requests_total\":[{\"type\":\"counter\",\"help\":\"Total" + + " number of HTTP requests\",\"unit\":\"requests\"}]}}"; + mockWebServer.enqueue(new MockResponse().setBody(successResponse)); + + // Test + var result = client.getAllMetrics(); + + // Verify + assertNotNull(result); + assertEquals(1, result.size()); + assertTrue(result.containsKey("http_requests_total")); + assertEquals(1, result.get("http_requests_total").size()); + assertEquals("counter", result.get("http_requests_total").getFirst().getType()); + assertEquals( + "Total number of HTTP requests", result.get("http_requests_total").getFirst().getHelp()); + assertEquals("requests", result.get("http_requests_total").getFirst().getUnit()); + } + + @Test + public void testGetSeries() throws IOException { + // Setup + String successResponse = + "{\"status\":\"success\",\"data\":[{\"__name__\":\"up\",\"job\":\"prometheus\",\"instance\":\"localhost:9090\"},{\"__name__\":\"up\",\"job\":\"node\",\"instance\":\"localhost:9100\"}]}"; + mockWebServer.enqueue(new MockResponse().setBody(successResponse)); + + // Test + HashMap params = new HashMap<>(); + params.put("match[]", "up"); + var result = client.getSeries(params); + + // Verify + assertNotNull(result); + assertEquals(2, result.size()); + + // First series + assertEquals("up", result.getFirst().get("__name__")); + assertEquals("prometheus", result.getFirst().get("job")); + assertEquals("localhost:9090", result.getFirst().get("instance")); + + // Second series + assertEquals("up", result.get(1).get("__name__")); + assertEquals("node", result.get(1).get("job")); + assertEquals("localhost:9100", result.get(1).get("instance")); + } + + @Test + public void testGetAlerts() throws IOException { + // Setup + String successResponse = + "{\"status\":\"success\",\"data\":{\"alerts\":[{\"labels\":{\"alertname\":\"HighErrorRate\",\"severity\":\"critical\"},\"annotations\":{\"summary\":\"High" + + " request error" + + " rate\"},\"state\":\"firing\",\"activeAt\":\"2023-01-01T00:00:00Z\",\"value\":\"0.15\"}]}}"; + mockWebServer.enqueue(new MockResponse().setBody(successResponse)); + + // Test + JSONObject result = client.getAlerts(); + + // Verify + assertNotNull(result); + assertTrue(result.has("alerts")); + JSONArray alerts = result.getJSONArray("alerts"); + assertEquals(1, alerts.length()); + JSONObject alert = alerts.getJSONObject(0); + assertEquals("HighErrorRate", alert.getJSONObject("labels").getString("alertname")); + assertEquals("critical", alert.getJSONObject("labels").getString("severity")); + assertEquals("firing", alert.getString("state")); + } + + @Test + public void testGetRules() throws IOException { + // Setup + String successResponse = + "{\"status\":\"success\",\"data\":{\"groups\":[{\"name\":\"example\",\"file\":\"rules.yml\",\"rules\":[{\"name\":\"HighErrorRate\",\"query\":\"rate(http_requests_total{status=~\\\"5..\\\"}[5m])" + + " / rate(http_requests_total[5m]) >" + + " 0.1\",\"type\":\"alerting\",\"health\":\"ok\",\"state\":\"inactive\"}]}]}}"; + mockWebServer.enqueue(new MockResponse().setBody(successResponse)); + + // Test + HashMap params = new HashMap<>(); + params.put("type", "alert"); + JSONObject result = client.getRules(params); + + // Verify + assertNotNull(result); + assertTrue(result.has("groups")); + JSONArray groups = result.getJSONArray("groups"); + assertEquals(1, groups.length()); + JSONObject group = groups.getJSONObject(0); + assertEquals("example", group.getString("name")); + JSONArray rules = group.getJSONArray("rules"); + assertEquals(1, rules.length()); + JSONObject rule = rules.getJSONObject(0); + assertEquals("HighErrorRate", rule.getString("name")); + assertEquals("alerting", rule.getString("type")); + } + + @Test + public void testGetAlertmanagerAlerts() throws IOException { + // Setup + String successResponse = + "[{\"labels\":{\"alertname\":\"HighErrorRate\",\"severity\":\"critical\"},\"annotations\":{\"summary\":\"High" + + " request error" + + " rate\"},\"state\":\"active\",\"activeAt\":\"2023-01-01T00:00:00Z\",\"value\":\"0.15\"}]"; + mockWebServer.enqueue(new MockResponse().setBody(successResponse)); + + // Test + HashMap params = new HashMap<>(); + params.put("active", "true"); + JSONArray result = client.getAlertmanagerAlerts(params); + + // Verify + assertNotNull(result); + assertEquals(1, result.length()); + JSONObject alert = result.getJSONObject(0); + assertEquals("HighErrorRate", alert.getJSONObject("labels").getString("alertname")); + assertEquals("critical", alert.getJSONObject("labels").getString("severity")); + assertEquals("active", alert.getString("state")); + } + + @Test + public void testGetAlertmanagerAlertGroups() throws IOException { + // Setup + String successResponse = + "[{\"labels\":{\"severity\":\"critical\"},\"alerts\":[{\"labels\":{\"alertname\":\"HighErrorRate\",\"severity\":\"critical\"},\"annotations\":{\"summary\":\"High" + + " request error rate\"},\"state\":\"active\"}]}]"; + mockWebServer.enqueue(new MockResponse().setBody(successResponse)); + + // Test + HashMap params = new HashMap<>(); + params.put("active", "true"); + JSONArray result = client.getAlertmanagerAlertGroups(params); + + // Verify + assertNotNull(result); + assertEquals(1, result.length()); + JSONObject group = result.getJSONObject(0); + assertEquals("critical", group.getJSONObject("labels").getString("severity")); + JSONArray alerts = group.getJSONArray("alerts"); + assertEquals(1, alerts.length()); + JSONObject alert = alerts.getJSONObject(0); + assertEquals("HighErrorRate", alert.getJSONObject("labels").getString("alertname")); + } + + @Test + public void testGetAlertmanagerReceivers() throws IOException { + // Setup + String successResponse = + "[{\"name\":\"email\",\"email_configs\":[{\"to\":\"admin@example.com\"}]},{\"name\":\"slack\",\"slack_configs\":[{\"channel\":\"#alerts\"}]}]"; + mockWebServer.enqueue(new MockResponse().setBody(successResponse)); + + // Test + JSONArray result = client.getAlertmanagerReceivers(); + + // Verify + assertNotNull(result); + assertEquals(2, result.length()); + assertEquals("email", result.getJSONObject(0).getString("name")); + assertEquals("slack", result.getJSONObject(1).getString("name")); + } + + @Test + public void testGetAlertmanagerSilences() throws IOException { + // Setup + String successResponse = + "[{\"id\":\"silence-123\",\"status\":{\"state\":\"active\"},\"createdBy\":\"admin\",\"comment\":\"Maintenance" + + " window\",\"startsAt\":\"2023-01-01T00:00:00Z\",\"endsAt\":\"2023-01-02T00:00:00Z\",\"matchers\":[{\"name\":\"severity\",\"value\":\"critical\",\"isRegex\":false}]}]"; + mockWebServer.enqueue(new MockResponse().setBody(successResponse)); + + // Test + JSONArray result = client.getAlertmanagerSilences(); + + // Verify + assertNotNull(result); + assertEquals(1, result.length()); + JSONObject silence = result.getJSONObject(0); + assertEquals("silence-123", silence.getString("id")); + assertEquals("active", silence.getJSONObject("status").getString("state")); + assertEquals("admin", silence.getString("createdBy")); + JSONArray matchers = silence.getJSONArray("matchers"); + assertEquals(1, matchers.length()); + assertEquals("severity", matchers.getJSONObject(0).getString("name")); + assertEquals("critical", matchers.getJSONObject(0).getString("value")); + } + + @Test + public void testReadResponseWithRequestId() throws Exception { + // Setup + String successResponse = "{\"status\":\"success\",\"data\":[\"job\",\"instance\"]}"; + MockResponse mockResponse = + new MockResponse().setBody(successResponse).addHeader("X-Request-ID", "test-request-id"); + mockWebServer.enqueue(mockResponse); + + // Test - the request ID will be logged but we can verify the method completes successfully + List labels = client.getLabels(new HashMap<>()); + + // Verify the method executed successfully + assertNotNull(labels); + assertEquals(2, labels.size()); + } + + @Test + public void testAlertmanagerResponseError() { + // Setup + mockWebServer.enqueue(new MockResponse().setResponseCode(500).setBody("Internal Server Error")); + + // Test & Verify + PrometheusClientException exception = + assertThrows( + PrometheusClientException.class, () -> client.getAlertmanagerAlerts(new HashMap<>())); + assertTrue(exception.getMessage().contains("Alertmanager request failed with code: 500")); + } + + @Test + public void testAlertmanagerResponseErrorWithNullBody() throws IOException { + // Setup - Create a mock response with null body + Request dummyRequest = new Request.Builder().url(mockWebServer.url("/")).build(); + Response nullBodyResponse = + new Response.Builder() + .request(dummyRequest) + .protocol(Protocol.HTTP_1_1) + .code(500) + .message("Internal Server Error") + .body(null) + .build(); + + // Create spy client that returns our custom response + OkHttpClient spyClient = spy(new OkHttpClient()); + Call mockCall = mock(Call.class); + when(mockCall.execute()).thenReturn(nullBodyResponse); + doAnswer(invocation -> mockCall).when(spyClient).newCall(any(Request.class)); + + // Create client with our spy + PrometheusClientImpl nullBodyClient = + new PrometheusClientImpl( + new OkHttpClient(), + URI.create(String.format("http://%s:%s", "localhost", mockWebServer.getPort())), + spyClient, + URI.create( + String.format("http://%s:%s/alertmanager", "localhost", mockWebServer.getPort()))); + + // Test & Verify + PrometheusClientException exception = + assertThrows( + PrometheusClientException.class, + () -> nullBodyClient.getAlertmanagerAlerts(new HashMap<>())); + assertTrue(exception.getMessage().contains("Alertmanager request failed with code: 500")); + assertTrue(exception.getMessage().contains("No response body")); + } +} diff --git a/direct-query-core/src/test/java/org/opensearch/sql/prometheus/query/PrometheusQueryHandlerTest.java b/direct-query-core/src/test/java/org/opensearch/sql/prometheus/query/PrometheusQueryHandlerTest.java new file mode 100644 index 00000000000..557e6425440 --- /dev/null +++ b/direct-query-core/src/test/java/org/opensearch/sql/prometheus/query/PrometheusQueryHandlerTest.java @@ -0,0 +1,644 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.prometheus.query; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertTrue; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.when; + +import java.io.IOException; +import java.util.Arrays; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import org.json.JSONArray; +import org.json.JSONObject; +import org.junit.Before; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.mockito.Mock; +import org.mockito.junit.MockitoJUnitRunner; +import org.opensearch.sql.datasource.model.DataSourceType; +import org.opensearch.sql.directquery.rest.model.DirectQueryResourceType; +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.prometheus.client.PrometheusClient; +import org.opensearch.sql.prometheus.exception.PrometheusClientException; +import org.opensearch.sql.prometheus.model.MetricMetadata; +import org.opensearch.sql.prometheus.model.PrometheusOptions; +import org.opensearch.sql.prometheus.model.PrometheusQueryType; + +@RunWith(MockitoJUnitRunner.class) +public class PrometheusQueryHandlerTest { + + private PrometheusQueryHandler handler; + + @Mock private PrometheusClient prometheusClient; + + @Before + public void setUp() { + handler = new PrometheusQueryHandler(); + } + + @Test + public void testGetSupportedDataSourceType() { + assertEquals(DataSourceType.PROMETHEUS, handler.getSupportedDataSourceType()); + } + + @Test + public void testCanHandle() { + assertTrue(handler.canHandle(prometheusClient)); + assertFalse(handler.canHandle(null)); + } + + @Test + public void testGetClientClass() { + assertEquals(PrometheusClient.class, handler.getClientClass()); + } + + @Test + public void testExecuteQueryWithRangeQuery() throws IOException { + // Setup + ExecuteDirectQueryRequest request = new ExecuteDirectQueryRequest(); + request.setQuery("up"); + request.setMaxResults(100); + request.setTimeout(30); + + PrometheusOptions options = new PrometheusOptions(); + options.setQueryType(PrometheusQueryType.RANGE); + options.setStart("1609459200"); // 2021-01-01 + options.setEnd("1609545600"); // 2021-01-02 + options.setStep("15s"); + request.setPrometheusOptions(options); + + JSONObject responseJson = + new JSONObject("{\"status\":\"success\",\"data\":{\"resultType\":\"matrix\"}}"); + when(prometheusClient.queryRange( + eq("up"), eq(1609459200L), eq(1609545600L), eq("15s"), eq(100), eq(30))) + .thenReturn(responseJson); + + // Test + String result = handler.executeQuery(prometheusClient, request); + + // Verify + assertNotNull(result); + JSONObject resultJson = new JSONObject(result); + assertEquals("success", resultJson.getString("status")); + assertEquals("matrix", resultJson.getJSONObject("data").getString("resultType")); + } + + @Test + public void testExecuteQueryWithInstantQuery() throws IOException { + // Setup + ExecuteDirectQueryRequest request = new ExecuteDirectQueryRequest(); + request.setQuery("up"); + request.setMaxResults(100); + request.setTimeout(30); + + PrometheusOptions options = new PrometheusOptions(); + options.setQueryType(PrometheusQueryType.INSTANT); + options.setTime("1609459200"); // 2021-01-01 + request.setPrometheusOptions(options); + + JSONObject responseJson = + new JSONObject("{\"status\":\"success\",\"data\":{\"resultType\":\"vector\"}}"); + when(prometheusClient.query(eq("up"), eq(1609459200L), eq(100), eq(30))) + .thenReturn(responseJson); + + // Test + String result = handler.executeQuery(prometheusClient, request); + + // Verify + assertNotNull(result); + JSONObject resultJson = new JSONObject(result); + assertEquals("success", resultJson.getString("status")); + assertEquals("vector", resultJson.getJSONObject("data").getString("resultType")); + } + + @Test + public void testExecuteQueryWithMissingStartEndTimes() throws IOException { + // Setup + ExecuteDirectQueryRequest request = new ExecuteDirectQueryRequest(); + request.setQuery("up"); + + PrometheusOptions options = new PrometheusOptions(); + options.setQueryType(PrometheusQueryType.RANGE); + // Missing start and end times + options.setStep("15s"); + request.setPrometheusOptions(options); + + // Test + String result = handler.executeQuery(prometheusClient, request); + + // Verify + assertNotNull(result); + JSONObject resultJson = new JSONObject(result); + assertTrue(resultJson.has("error")); + assertEquals( + "Start and end times are required for Prometheus queries", resultJson.getString("error")); + } + + @Test + public void testExecuteQueryWithMissingEndTime() throws IOException { + // Setup + ExecuteDirectQueryRequest request = new ExecuteDirectQueryRequest(); + request.setQuery("up"); + + PrometheusOptions options = new PrometheusOptions(); + options.setQueryType(PrometheusQueryType.RANGE); + options.setStep("15s"); + options.setStart("now"); + request.setPrometheusOptions(options); + + // Test + String result = handler.executeQuery(prometheusClient, request); + + // Verify + assertNotNull(result); + JSONObject resultJson = new JSONObject(result); + assertTrue(resultJson.has("error")); + assertEquals( + "Start and end times are required for Prometheus queries", resultJson.getString("error")); + } + + @Test + public void testExecuteQueryWithMissingTimeForInstant() throws IOException { + // Setup + ExecuteDirectQueryRequest request = new ExecuteDirectQueryRequest(); + request.setQuery("up"); + + PrometheusOptions options = new PrometheusOptions(); + options.setQueryType(PrometheusQueryType.INSTANT); + // Missing time + request.setPrometheusOptions(options); + + // Test + String result = handler.executeQuery(prometheusClient, request); + + // Verify + assertNotNull(result); + JSONObject resultJson = new JSONObject(result); + assertTrue(resultJson.has("error")); + assertEquals("Time is required for instant Prometheus queries", resultJson.getString("error")); + } + + @Test + public void testExecuteQueryWithInvalidTimeFormat() throws IOException { + // Setup + ExecuteDirectQueryRequest request = new ExecuteDirectQueryRequest(); + request.setQuery("up"); + + PrometheusOptions options = new PrometheusOptions(); + options.setQueryType(PrometheusQueryType.INSTANT); + options.setTime("invalid-time"); + request.setPrometheusOptions(options); + + // Test + String result = handler.executeQuery(prometheusClient, request); + + // Verify + assertNotNull(result); + JSONObject resultJson = new JSONObject(result); + assertTrue(resultJson.has("error")); + assertTrue(resultJson.getString("error").contains("Invalid time format")); + } + + @Test + public void testExecuteQueryWithNullQueryType() throws IOException { + // Setup + ExecuteDirectQueryRequest request = new ExecuteDirectQueryRequest(); + request.setQuery("up"); + + PrometheusOptions options = new PrometheusOptions(); + options.setQueryType(null); // Invalid query type + request.setPrometheusOptions(options); + + // Test + String result = handler.executeQuery(prometheusClient, request); + + // Verify + assertNotNull(result); + JSONObject resultJson = new JSONObject(result); + assertTrue(resultJson.has("error")); + assertEquals("Query type is required for Prometheus queries", resultJson.getString("error")); + } + + @Test + public void testGetResourcesLabels() throws IOException { + // Setup + GetDirectQueryResourcesRequest request = new GetDirectQueryResourcesRequest(); + request.setResourceType(DirectQueryResourceType.LABELS); + Map queryParams = new HashMap<>(); + request.setQueryParams(queryParams); + + List labels = Arrays.asList("job", "instance", "env"); + when(prometheusClient.getLabels(queryParams)).thenReturn(labels); + + // Test + GetDirectQueryResourcesResponse response = handler.getResources(prometheusClient, request); + + // Verify + assertNotNull(response); + assertEquals(labels, response.getData()); + } + + @Test + public void testGetResourcesLabel() throws IOException { + // Setup + GetDirectQueryResourcesRequest request = new GetDirectQueryResourcesRequest(); + request.setResourceType(DirectQueryResourceType.LABEL); + request.setResourceName("job"); + Map queryParams = new HashMap<>(); + request.setQueryParams(queryParams); + + List labelValues = Arrays.asList("prometheus", "node-exporter", "cadvisor"); + when(prometheusClient.getLabel("job", queryParams)).thenReturn(labelValues); + + // Test + GetDirectQueryResourcesResponse response = handler.getResources(prometheusClient, request); + + // Verify + assertNotNull(response); + assertEquals(labelValues, response.getData()); + } + + @Test + public void testGetResourcesMetadata() throws IOException { + // Setup + GetDirectQueryResourcesRequest request = new GetDirectQueryResourcesRequest(); + request.setResourceType(DirectQueryResourceType.METADATA); + Map queryParams = new HashMap<>(); + request.setQueryParams(queryParams); + + Map> metadata = new HashMap<>(); + metadata.put( + "up", Arrays.asList(new MetricMetadata("up", "gauge", "Whether the target is up"))); + when(prometheusClient.getAllMetrics(queryParams)).thenReturn(metadata); + + // Test + GetDirectQueryResourcesResponse response = handler.getResources(prometheusClient, request); + + // Verify + assertNotNull(response); + assertEquals(metadata, response.getData()); + } + + @Test + public void testGetResourcesSeries() throws IOException { + // Setup + GetDirectQueryResourcesRequest request = new GetDirectQueryResourcesRequest(); + request.setResourceType(DirectQueryResourceType.SERIES); + Map queryParams = new HashMap<>(); + request.setQueryParams(queryParams); + + List> series = + Arrays.asList( + Map.of("__name__", "up", "job", "prometheus"), + Map.of("__name__", "up", "job", "node-exporter")); + when(prometheusClient.getSeries(queryParams)).thenReturn(series); + + // Test + GetDirectQueryResourcesResponse response = handler.getResources(prometheusClient, request); + + // Verify + assertNotNull(response); + assertEquals(series, response.getData()); + } + + @Test(expected = IllegalArgumentException.class) + public void testGetResourcesInvalidType() { + // Setup + GetDirectQueryResourcesRequest request = new GetDirectQueryResourcesRequest(); + request.setResourceTypeFromString("INVALID_TYPE"); + + // Test - should throw exception + handler.getResources(prometheusClient, request); + } + + @Test(expected = IllegalArgumentException.class) + public void testGetResourcesUnknownType() { + // Setup + GetDirectQueryResourcesRequest request = new GetDirectQueryResourcesRequest(); + request.setResourceType(DirectQueryResourceType.UNKNOWN); + + // Test - should throw exception + handler.getResources(prometheusClient, request); + } + + @Test(expected = PrometheusClientException.class) + public void testGetResourcesIOException() throws IOException { + // Setup + GetDirectQueryResourcesRequest request = new GetDirectQueryResourcesRequest(); + request.setResourceType(DirectQueryResourceType.LABELS); + Map queryParams = new HashMap<>(); + request.setQueryParams(queryParams); + + when(prometheusClient.getLabels(queryParams)).thenThrow(new IOException("Connection failed")); + + // Test - should throw exception + handler.getResources(prometheusClient, request); + } + + @Test + public void testExecuteQueryWithPrometheusClientException() throws IOException { + // Setup + ExecuteDirectQueryRequest request = new ExecuteDirectQueryRequest(); + request.setQuery("up"); + + PrometheusOptions options = new PrometheusOptions(); + options.setQueryType(PrometheusQueryType.INSTANT); + options.setTime("1609459200"); // 2021-01-01 + request.setPrometheusOptions(options); + + String errorMessage = "Prometheus server error"; + when(prometheusClient.query(eq("up"), eq(1609459200L), eq(null), eq(null))) + .thenThrow( + new org.opensearch.sql.prometheus.exception.PrometheusClientException(errorMessage)); + + // Test + String result = handler.executeQuery(prometheusClient, request); + + // Verify + assertNotNull(result); + JSONObject resultJson = new JSONObject(result); + assertTrue(resultJson.has("error")); + assertEquals(errorMessage, resultJson.getString("error")); + } + + @Test + public void testExecuteQueryWithIOException() throws IOException { + // Setup + ExecuteDirectQueryRequest request = new ExecuteDirectQueryRequest(); + request.setQuery("up"); + + PrometheusOptions options = new PrometheusOptions(); + options.setQueryType(PrometheusQueryType.INSTANT); + options.setTime("1609459200"); // 2021-01-01 + request.setPrometheusOptions(options); + + String errorMessage = "Network connection error"; + when(prometheusClient.query(eq("up"), eq(1609459200L), eq(null), eq(null))) + .thenThrow(new IOException(errorMessage)); + + // Test + String result = handler.executeQuery(prometheusClient, request); + + // Verify + assertNotNull(result); + JSONObject resultJson = new JSONObject(result); + assertTrue(resultJson.has("error")); + assertEquals(errorMessage, resultJson.getString("error")); + } + + @Test(expected = IllegalArgumentException.class) + public void testGetResourcesWithNullResourceType() { + // Setup + GetDirectQueryResourcesRequest request = new GetDirectQueryResourcesRequest(); + request.setResourceType(null); // Null resource type + + // Test - should throw exception + handler.getResources(prometheusClient, request); + } + + @Test + public void testGetResourcesAlerts() throws IOException { + // Setup + GetDirectQueryResourcesRequest request = new GetDirectQueryResourcesRequest(); + request.setResourceType(DirectQueryResourceType.ALERTS); + + JSONObject alertsJson = new JSONObject(); + alertsJson.put("status", "success"); + alertsJson.put( + "data", + new JSONObject() + .put( + "alerts", + new JSONObject() + .put( + "alerts", + Arrays.asList( + new JSONObject() + .put("name", "HighCPULoad") + .put("state", "firing") + .put("activeAt", "2023-01-01T00:00:00Z"), + new JSONObject() + .put("name", "InstanceDown") + .put("state", "pending") + .put("activeAt", "2023-01-01T00:05:00Z"))))); + + when(prometheusClient.getAlerts()).thenReturn(alertsJson); + + // Test + GetDirectQueryResourcesResponse response = handler.getResources(prometheusClient, request); + + // Verify + assertNotNull(response); + Map data = (Map) response.getData(); + assertEquals("success", data.get("status")); + assertTrue(data.containsKey("data")); + } + + @Test + public void testGetResourcesRules() throws IOException { + // Setup + GetDirectQueryResourcesRequest request = new GetDirectQueryResourcesRequest(); + request.setResourceType(DirectQueryResourceType.RULES); + Map queryParams = new HashMap<>(); + request.setQueryParams(queryParams); + + JSONObject rulesJson = new JSONObject(); + rulesJson.put("status", "success"); + rulesJson.put( + "data", + new JSONObject() + .put( + "groups", + Arrays.asList( + new JSONObject() + .put("name", "example") + .put( + "rules", + Arrays.asList( + new JSONObject() + .put("name", "HighErrorRate") + .put( + "query", + "rate(http_requests_total{status=~\"5..\"}[5m]) > 0.5") + .put("type", "alerting"), + new JSONObject() + .put("name", "RequestRate") + .put("query", "rate(http_requests_total[5m])") + .put("type", "recording")))))); + + when(prometheusClient.getRules(queryParams)).thenReturn(rulesJson); + + // Test + GetDirectQueryResourcesResponse response = handler.getResources(prometheusClient, request); + + // Verify + assertNotNull(response); + Map data = (Map) response.getData(); + assertEquals("success", data.get("status")); + assertTrue(data.containsKey("data")); + } + + @Test + public void testGetResourcesAlertmanagerAlerts() throws IOException { + // Setup + GetDirectQueryResourcesRequest request = new GetDirectQueryResourcesRequest(); + request.setResourceType(DirectQueryResourceType.ALERTMANAGER_ALERTS); + Map queryParams = new HashMap<>(); + request.setQueryParams(queryParams); + + JSONObject alert1 = + new JSONObject() + .put("status", "firing") + .put( + "labels", + new JSONObject().put("alertname", "HighCPULoad").put("severity", "critical")) + .put( + "annotations", + new JSONObject() + .put("summary", "High CPU load on instance") + .put("description", "CPU load is above 90%")); + + JSONObject alert2 = + new JSONObject() + .put("status", "resolved") + .put( + "labels", + new JSONObject().put("alertname", "InstanceDown").put("severity", "critical")) + .put( + "annotations", + new JSONObject() + .put("summary", "Instance is down") + .put("description", "Instance has been down for more than 5 minutes")); + + JSONArray alertsArray = new JSONArray(); + alertsArray.put(alert1); + alertsArray.put(alert2); + + when(prometheusClient.getAlertmanagerAlerts(queryParams)).thenReturn(alertsArray); + + // Test + GetDirectQueryResourcesResponse response = handler.getResources(prometheusClient, request); + + // Verify + assertNotNull(response); + List data = (List) response.getData(); + assertEquals(2, data.size()); + } + + @Test + public void testGetResourcesAlertmanagerAlertGroups() throws IOException { + // Setup + GetDirectQueryResourcesRequest request = new GetDirectQueryResourcesRequest(); + request.setResourceType(DirectQueryResourceType.ALERTMANAGER_ALERT_GROUPS); + Map queryParams = new HashMap<>(); + request.setQueryParams(queryParams); + + JSONObject group1 = + new JSONObject() + .put("labels", new JSONObject().put("severity", "critical")) + .put( + "alerts", + new JSONArray() + .put( + new JSONObject() + .put("status", "firing") + .put( + "labels", + new JSONObject() + .put("alertname", "HighCPULoad") + .put("severity", "critical")))); + + JSONObject group2 = + new JSONObject() + .put("labels", new JSONObject().put("severity", "warning")) + .put( + "alerts", + new JSONArray() + .put( + new JSONObject() + .put("status", "firing") + .put( + "labels", + new JSONObject() + .put("alertname", "HighMemoryUsage") + .put("severity", "warning")))); + + JSONArray groupsArray = new JSONArray(); + groupsArray.put(group1); + groupsArray.put(group2); + + when(prometheusClient.getAlertmanagerAlertGroups(queryParams)).thenReturn(groupsArray); + + // Test + GetDirectQueryResourcesResponse response = handler.getResources(prometheusClient, request); + + // Verify + assertNotNull(response); + List data = (List) response.getData(); + assertEquals(2, data.size()); + } + + @Test + public void testGetResourcesAlertmanagerReceivers() throws IOException { + // Setup + GetDirectQueryResourcesRequest request = new GetDirectQueryResourcesRequest(); + request.setResourceType(DirectQueryResourceType.ALERTMANAGER_RECEIVERS); + + JSONArray receiversArray = new JSONArray(); + receiversArray.put(new JSONObject().put("name", "email-notifications")); + receiversArray.put(new JSONObject().put("name", "slack-alerts")); + receiversArray.put(new JSONObject().put("name", "pagerduty")); + + when(prometheusClient.getAlertmanagerReceivers()).thenReturn(receiversArray); + + // Test + GetDirectQueryResourcesResponse response = handler.getResources(prometheusClient, request); + + // Verify + assertNotNull(response); + List data = (List) response.getData(); + assertEquals(3, data.size()); + } + + @Test + public void testGetResourcesAlertmanagerSilences() throws IOException { + // Setup + GetDirectQueryResourcesRequest request = new GetDirectQueryResourcesRequest(); + request.setResourceType(DirectQueryResourceType.ALERTMANAGER_SILENCES); + + JSONArray silencesArray = new JSONArray(); + silencesArray.put( + new JSONObject() + .put("id", "silence-1") + .put("status", "active") + .put("createdBy", "admin") + .put("comment", "Maintenance window")); + silencesArray.put( + new JSONObject() + .put("id", "silence-2") + .put("status", "expired") + .put("createdBy", "admin") + .put("comment", "Weekend maintenance")); + + when(prometheusClient.getAlertmanagerSilences()).thenReturn(silencesArray); + + // Test + GetDirectQueryResourcesResponse response = handler.getResources(prometheusClient, request); + + // Verify + assertNotNull(response); + List data = (List) response.getData(); + assertEquals(2, data.size()); + } +} diff --git a/direct-query-core/src/test/java/org/opensearch/sql/prometheus/utils/PrometheusClientUtilsTest.java b/direct-query-core/src/test/java/org/opensearch/sql/prometheus/utils/PrometheusClientUtilsTest.java new file mode 100644 index 00000000000..4b7959d5428 --- /dev/null +++ b/direct-query-core/src/test/java/org/opensearch/sql/prometheus/utils/PrometheusClientUtilsTest.java @@ -0,0 +1,313 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.prometheus.utils; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertTrue; +import static org.junit.Assert.fail; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +import com.google.common.collect.ImmutableList; +import java.util.HashMap; +import java.util.Map; +import org.junit.Before; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.mockito.Mock; +import org.mockito.junit.MockitoJUnitRunner; +import org.opensearch.sql.common.setting.Settings; +import org.opensearch.sql.datasource.client.exceptions.DataSourceClientException; +import org.opensearch.sql.datasource.model.DataSourceMetadata; +import org.opensearch.sql.prometheus.client.PrometheusClient; + +@RunWith(MockitoJUnitRunner.class) +public class PrometheusClientUtilsTest { + + @Mock private Settings settings; + + @Before + public void setUp() { + when(settings.getSettingValue(Settings.Key.DATASOURCES_URI_HOSTS_DENY_LIST)) + .thenReturn(ImmutableList.of("http://localhost:9200")); + } + + @Test + public void testHasAlertmanagerConfigWithAlertmanagerUri() { + // Setup + Map properties = new HashMap<>(); + properties.put(PrometheusClientUtils.ALERTMANAGER_URI, "http://alertmanager:9093"); + + // Test + boolean result = PrometheusClientUtils.hasAlertmanagerConfig(properties); + + // Verify + assertTrue("Should return true when Alertmanager URI is present", result); + } + + @Test + public void testHasAlertmanagerConfigWithoutAlertmanagerUri() { + // Setup + Map properties = new HashMap<>(); + properties.put(PrometheusClientUtils.PROMETHEUS_URI, "http://prometheus:9090"); + // No Alertmanager URI + + // Test + boolean result = PrometheusClientUtils.hasAlertmanagerConfig(properties); + + // Verify + assertFalse("Should return false when Alertmanager URI is not present", result); + } + + @Test + public void testHasAlertmanagerConfigWithEmptyProperties() { + // Setup + Map properties = new HashMap<>(); + + // Test + boolean result = PrometheusClientUtils.hasAlertmanagerConfig(properties); + + // Verify + assertFalse("Should return false when properties are empty", result); + } + + @Test + public void testCreateAlertmanagerPropertiesWithBasicAuth() { + // Setup + Map properties = new HashMap<>(); + properties.put(PrometheusClientUtils.ALERTMANAGER_URI, "http://alertmanager:9093"); + properties.put(PrometheusClientUtils.ALERTMANAGER_AUTH_TYPE, "basicauth"); + properties.put(PrometheusClientUtils.ALERTMANAGER_USERNAME, "admin"); + properties.put(PrometheusClientUtils.ALERTMANAGER_PASSWORD, "password"); + + // Test + Map result = PrometheusClientUtils.createAlertmanagerProperties(properties); + + // Verify + assertNotNull("Result should not be null", result); + assertEquals("basicauth", result.get(PrometheusClientUtils.AUTH_TYPE)); + assertEquals("admin", result.get(PrometheusClientUtils.USERNAME)); + assertEquals("password", result.get(PrometheusClientUtils.PASSWORD)); + } + + @Test + public void testCreateAlertmanagerPropertiesWithAwsAuth() { + // Setup + Map properties = new HashMap<>(); + properties.put(PrometheusClientUtils.ALERTMANAGER_URI, "http://alertmanager:9093"); + properties.put(PrometheusClientUtils.ALERTMANAGER_AUTH_TYPE, "awssigv4auth"); + properties.put(PrometheusClientUtils.ALERTMANAGER_ACCESS_KEY, "access-key"); + properties.put(PrometheusClientUtils.ALERTMANAGER_SECRET_KEY, "secret-key"); + properties.put(PrometheusClientUtils.ALERTMANAGER_REGION, "us-west-1"); + + // Test + Map result = PrometheusClientUtils.createAlertmanagerProperties(properties); + + // Verify + assertNotNull("Result should not be null", result); + assertEquals("awssigv4auth", result.get(PrometheusClientUtils.AUTH_TYPE)); + assertEquals("access-key", result.get(PrometheusClientUtils.ACCESS_KEY)); + assertEquals("secret-key", result.get(PrometheusClientUtils.SECRET_KEY)); + assertEquals("us-west-1", result.get(PrometheusClientUtils.REGION)); + } + + @Test + public void testCreateAlertmanagerPropertiesWithNoAuth() { + // Setup + Map properties = new HashMap<>(); + properties.put(PrometheusClientUtils.ALERTMANAGER_URI, "http://alertmanager:9093"); + // No auth properties + + // Test + Map result = PrometheusClientUtils.createAlertmanagerProperties(properties); + + // Verify + assertTrue("Result should be empty when no auth is provided", result.isEmpty()); + } + + @Test + public void testCreateAlertmanagerPropertiesWithNullAuthType() { + // Setup + Map properties = new HashMap<>(); + properties.put(PrometheusClientUtils.ALERTMANAGER_URI, "http://alertmanager:9093"); + properties.put(PrometheusClientUtils.ALERTMANAGER_AUTH_TYPE, null); + + // Test + Map result = PrometheusClientUtils.createAlertmanagerProperties(properties); + + // Verify + assertNotNull("Result should not be null", result); + assertEquals(null, result.get(PrometheusClientUtils.AUTH_TYPE)); + } + + @Test + public void testCreateAlertmanagerPropertiesWithUnsupportedAuthType() { + // Setup + Map properties = new HashMap<>(); + properties.put(PrometheusClientUtils.ALERTMANAGER_URI, "http://alertmanager:9093"); + properties.put(PrometheusClientUtils.ALERTMANAGER_AUTH_TYPE, "unsupportedauth"); + + // Test + Map result = PrometheusClientUtils.createAlertmanagerProperties(properties); + + // Verify + assertNotNull("Result should not be null", result); + assertEquals("unsupportedauth", result.get(PrometheusClientUtils.AUTH_TYPE)); + // Should not contain any auth credentials since auth type is not recognized + assertFalse(result.containsKey(PrometheusClientUtils.USERNAME)); + assertFalse(result.containsKey(PrometheusClientUtils.PASSWORD)); + assertFalse(result.containsKey(PrometheusClientUtils.ACCESS_KEY)); + assertFalse(result.containsKey(PrometheusClientUtils.SECRET_KEY)); + assertFalse(result.containsKey(PrometheusClientUtils.REGION)); + } + + @Test + public void testCreatePrometheusClientWithAlertmanagerConfig() { + // Setup + Map properties = new HashMap<>(); + properties.put(PrometheusClientUtils.PROMETHEUS_URI, "http://prometheus:9090"); + properties.put(PrometheusClientUtils.ALERTMANAGER_URI, "http://alertmanager:9093"); + properties.put(PrometheusClientUtils.ALERTMANAGER_AUTH_TYPE, "basicauth"); + properties.put(PrometheusClientUtils.ALERTMANAGER_USERNAME, "admin"); + properties.put(PrometheusClientUtils.ALERTMANAGER_PASSWORD, "password"); + + DataSourceMetadata metadata = mock(DataSourceMetadata.class); + when(metadata.getProperties()).thenReturn(properties); + // Removed unnecessary stubbing for metadata.getConnector() + + // Test + PrometheusClient client = PrometheusClientUtils.createPrometheusClient(metadata, settings); + + // Verify + assertNotNull("Client should not be null", client); + } + + @Test + public void testCreatePrometheusClientWithoutAlertmanagerConfig() { + // Setup + Map properties = new HashMap<>(); + properties.put(PrometheusClientUtils.PROMETHEUS_URI, "http://prometheus:9090"); + // No Alertmanager URI + + DataSourceMetadata metadata = mock(DataSourceMetadata.class); + when(metadata.getProperties()).thenReturn(properties); + // Removed unnecessary stubbing for metadata.getConnector() + + // Test + PrometheusClient client = PrometheusClientUtils.createPrometheusClient(metadata, settings); + + // Verify + assertNotNull("Client should not be null", client); + } + + @Test + public void testCreatePrometheusClientWithBasicAuth() { + // Setup + Map properties = new HashMap<>(); + properties.put(PrometheusClientUtils.PROMETHEUS_URI, "http://prometheus:9090"); + properties.put(PrometheusClientUtils.AUTH_TYPE, "basicauth"); + properties.put(PrometheusClientUtils.USERNAME, "user"); + properties.put(PrometheusClientUtils.PASSWORD, "pass"); + + DataSourceMetadata metadata = mock(DataSourceMetadata.class); + when(metadata.getProperties()).thenReturn(properties); + + // Test + PrometheusClient client = PrometheusClientUtils.createPrometheusClient(metadata, settings); + + // Verify + assertNotNull("Client should not be null", client); + } + + @Test + public void testCreatePrometheusClientWithAwsAuth() { + // Setup + Map properties = new HashMap<>(); + properties.put(PrometheusClientUtils.PROMETHEUS_URI, "http://prometheus:9090"); + properties.put(PrometheusClientUtils.AUTH_TYPE, "awssigv4"); + properties.put(PrometheusClientUtils.ACCESS_KEY, "access-key"); + properties.put(PrometheusClientUtils.SECRET_KEY, "secret-key"); + properties.put(PrometheusClientUtils.REGION, "us-west-1"); + + DataSourceMetadata metadata = mock(DataSourceMetadata.class); + when(metadata.getProperties()).thenReturn(properties); + + // Test + PrometheusClient client = PrometheusClientUtils.createPrometheusClient(metadata, settings); + + // Verify + assertNotNull("Client should not be null", client); + } + + @Test + public void testCreatePrometheusClientWithAlertmanagerAndAwsAuth() { + // Setup + Map properties = new HashMap<>(); + properties.put(PrometheusClientUtils.PROMETHEUS_URI, "http://prometheus:9090"); + properties.put(PrometheusClientUtils.AUTH_TYPE, "awssigv4"); + properties.put(PrometheusClientUtils.ACCESS_KEY, "access-key"); + properties.put(PrometheusClientUtils.SECRET_KEY, "secret-key"); + properties.put(PrometheusClientUtils.REGION, "us-west-1"); + properties.put(PrometheusClientUtils.ALERTMANAGER_URI, "http://alertmanager:9093"); + + DataSourceMetadata metadata = mock(DataSourceMetadata.class); + when(metadata.getProperties()).thenReturn(properties); + + // Test + PrometheusClient client = PrometheusClientUtils.createPrometheusClient(metadata, settings); + + // Verify + assertNotNull("Client should not be null", client); + } + + @Test(expected = DataSourceClientException.class) + public void testCreatePrometheusClientWithMissingUri() { + // Setup + Map properties = new HashMap<>(); + // Missing URI property + + DataSourceMetadata metadata = mock(DataSourceMetadata.class); + when(metadata.getProperties()).thenReturn(properties); + + // Test - should throw exception + PrometheusClientUtils.createPrometheusClient(metadata, settings); + } + + @Test(expected = IllegalArgumentException.class) + public void testCreatePrometheusClientWithInvalidUri() { + // Setup + Map properties = new HashMap<>(); + // Using malformed URI that will definitely be rejected + properties.put(PrometheusClientUtils.PROMETHEUS_URI, "ht tp://invalid:9090"); + + DataSourceMetadata metadata = mock(DataSourceMetadata.class); + when(metadata.getProperties()).thenReturn(properties); + + // Test - should throw exception + PrometheusClientUtils.createPrometheusClient(metadata, settings); + } + + @Test + public void testCreatePrometheusClientWithUnsupportedAuthType() { + // Setup + Map properties = new HashMap<>(); + properties.put(PrometheusClientUtils.PROMETHEUS_URI, "http://prometheus:9090"); + properties.put(PrometheusClientUtils.AUTH_TYPE, "unsupported"); + + DataSourceMetadata metadata = mock(DataSourceMetadata.class); + when(metadata.getProperties()).thenReturn(properties); + + // Test - should throw IllegalArgumentException + try { + PrometheusClientUtils.createPrometheusClient(metadata, settings); + fail("Expected IllegalArgumentException to be thrown"); + } catch (IllegalArgumentException e) { + assertTrue(e.getMessage().contains("AUTH Type : unsupported is not supported")); + } + } +} diff --git a/direct-query-core/src/test/resources/non_json_response.json b/direct-query-core/src/test/resources/non_json_response.json new file mode 100644 index 00000000000..b77abc5c5fe --- /dev/null +++ b/direct-query-core/src/test/resources/non_json_response.json @@ -0,0 +1,3 @@ +{ + "invalid json content +} \ No newline at end of file diff --git a/prometheus/src/test/java/org/opensearch/sql/prometheus/constants/TestConstants.java b/direct-query-core/src/testFixtures/java/org/opensearch/sql/prometheus/constant/TestConstants.java similarity index 88% rename from prometheus/src/test/java/org/opensearch/sql/prometheus/constants/TestConstants.java rename to direct-query-core/src/testFixtures/java/org/opensearch/sql/prometheus/constant/TestConstants.java index 758e5b0cafd..9f1342ebe45 100644 --- a/prometheus/src/test/java/org/opensearch/sql/prometheus/constants/TestConstants.java +++ b/direct-query-core/src/testFixtures/java/org/opensearch/sql/prometheus/constant/TestConstants.java @@ -3,7 +3,7 @@ * SPDX-License-Identifier: Apache-2.0 */ -package org.opensearch.sql.prometheus.constants; +package org.opensearch.sql.prometheus.constant; public class TestConstants { public static final String QUERY = "test_query"; diff --git a/prometheus/src/test/java/org/opensearch/sql/prometheus/utils/TestUtils.java b/direct-query-core/src/testFixtures/java/org/opensearch/sql/prometheus/utils/TestUtils.java similarity index 100% rename from prometheus/src/test/java/org/opensearch/sql/prometheus/utils/TestUtils.java rename to direct-query-core/src/testFixtures/java/org/opensearch/sql/prometheus/utils/TestUtils.java diff --git a/prometheus/build.gradle b/prometheus/build.gradle index 5060ed6f10a..7523703a472 100644 --- a/prometheus/build.gradle +++ b/prometheus/build.gradle @@ -16,6 +16,8 @@ repositories { dependencies { api project(':core') implementation project(':datasources') + implementation project(':direct-query-core') + testImplementation(testFixtures(project(":direct-query-core"))) implementation group: 'org.opensearch', name: 'opensearch', version: "${opensearch_version}" implementation "io.github.resilience4j:resilience4j-retry:${resilience4j_version}" @@ -28,7 +30,6 @@ dependencies { testImplementation group: 'org.hamcrest', name: 'hamcrest-library', version: "${hamcrest_version}" 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' } test { diff --git a/prometheus/src/main/java/org/opensearch/sql/prometheus/client/PrometheusClient.java b/prometheus/src/main/java/org/opensearch/sql/prometheus/client/PrometheusClient.java deleted file mode 100644 index f58ca56a8c1..00000000000 --- a/prometheus/src/main/java/org/opensearch/sql/prometheus/client/PrometheusClient.java +++ /dev/null @@ -1,24 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.sql.prometheus.client; - -import java.io.IOException; -import java.util.List; -import java.util.Map; -import org.json.JSONArray; -import org.json.JSONObject; -import org.opensearch.sql.prometheus.request.system.model.MetricMetadata; - -public interface PrometheusClient { - - JSONObject queryRange(String query, Long start, Long end, String step) throws IOException; - - List getLabels(String metricName) throws IOException; - - Map> getAllMetrics() throws IOException; - - JSONArray queryExemplars(String query, Long start, Long end) throws IOException; -} diff --git a/prometheus/src/main/java/org/opensearch/sql/prometheus/client/PrometheusClientImpl.java b/prometheus/src/main/java/org/opensearch/sql/prometheus/client/PrometheusClientImpl.java deleted file mode 100644 index 48154964eb7..00000000000 --- a/prometheus/src/main/java/org/opensearch/sql/prometheus/client/PrometheusClientImpl.java +++ /dev/null @@ -1,134 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.sql.prometheus.client; - -import com.fasterxml.jackson.core.type.TypeReference; -import com.fasterxml.jackson.databind.ObjectMapper; -import java.io.IOException; -import java.net.URI; -import java.net.URLEncoder; -import java.nio.charset.StandardCharsets; -import java.util.ArrayList; -import java.util.HashMap; -import java.util.List; -import java.util.Map; -import java.util.Objects; -import okhttp3.OkHttpClient; -import okhttp3.Request; -import okhttp3.Response; -import org.apache.logging.log4j.LogManager; -import org.apache.logging.log4j.Logger; -import org.json.JSONArray; -import org.json.JSONException; -import org.json.JSONObject; -import org.opensearch.sql.prometheus.exceptions.PrometheusClientException; -import org.opensearch.sql.prometheus.request.system.model.MetricMetadata; - -public class PrometheusClientImpl implements PrometheusClient { - - private static final Logger logger = LogManager.getLogger(PrometheusClientImpl.class); - - private final OkHttpClient okHttpClient; - - private final URI uri; - - public PrometheusClientImpl(OkHttpClient okHttpClient, URI uri) { - this.okHttpClient = okHttpClient; - this.uri = uri; - } - - @Override - public JSONObject queryRange(String query, Long start, Long end, String step) throws IOException { - String queryUrl = - String.format( - "%s/api/v1/query_range?query=%s&start=%s&end=%s&step=%s", - uri.toString().replaceAll("/$", ""), - URLEncoder.encode(query, StandardCharsets.UTF_8), - start, - end, - step); - logger.debug("queryUrl: " + queryUrl); - Request request = new Request.Builder().url(queryUrl).build(); - Response response = this.okHttpClient.newCall(request).execute(); - JSONObject jsonObject = readResponse(response); - return jsonObject.getJSONObject("data"); - } - - @Override - public List getLabels(String metricName) throws IOException { - String queryUrl = - String.format( - "%s/api/v1/labels?%s=%s", - uri.toString().replaceAll("/$", ""), - URLEncoder.encode("match[]", StandardCharsets.UTF_8), - URLEncoder.encode(metricName, StandardCharsets.UTF_8)); - logger.debug("queryUrl: " + queryUrl); - Request request = new Request.Builder().url(queryUrl).build(); - Response response = this.okHttpClient.newCall(request).execute(); - JSONObject jsonObject = readResponse(response); - return toListOfLabels(jsonObject.getJSONArray("data")); - } - - @Override - public Map> getAllMetrics() throws IOException { - String queryUrl = String.format("%s/api/v1/metadata", uri.toString().replaceAll("/$", "")); - logger.debug("queryUrl: " + queryUrl); - Request request = new Request.Builder().url(queryUrl).build(); - Response response = this.okHttpClient.newCall(request).execute(); - JSONObject jsonObject = readResponse(response); - TypeReference>> typeRef = new TypeReference<>() {}; - return new ObjectMapper().readValue(jsonObject.getJSONObject("data").toString(), typeRef); - } - - @Override - public JSONArray queryExemplars(String query, Long start, Long end) throws IOException { - String queryUrl = - String.format( - "%s/api/v1/query_exemplars?query=%s&start=%s&end=%s", - uri.toString().replaceAll("/$", ""), - URLEncoder.encode(query, StandardCharsets.UTF_8), - start, - end); - logger.debug("queryUrl: " + queryUrl); - Request request = new Request.Builder().url(queryUrl).build(); - Response response = this.okHttpClient.newCall(request).execute(); - JSONObject jsonObject = readResponse(response); - return jsonObject.getJSONArray("data"); - } - - private List toListOfLabels(JSONArray array) { - List result = new ArrayList<>(); - for (int i = 0; i < array.length(); i++) { - // __name__ is internal label in prometheus representing the metric name. - // Exempting this from labels list as it is not required in any of the operations. - if (!"__name__".equals(array.optString(i))) { - result.add(array.optString(i)); - } - } - return result; - } - - private JSONObject readResponse(Response response) throws IOException { - if (response.isSuccessful()) { - JSONObject jsonObject; - try { - jsonObject = new JSONObject(Objects.requireNonNull(response.body()).string()); - } catch (JSONException jsonException) { - throw new PrometheusClientException( - "Prometheus returned unexpected body, " - + "please verify your prometheus server setup."); - } - if ("success".equals(jsonObject.getString("status"))) { - return jsonObject; - } else { - throw new PrometheusClientException(jsonObject.getString("error")); - } - } else { - throw new PrometheusClientException( - String.format("Request to Prometheus is Unsuccessful with code : %s", response.code())); - } - } -} diff --git a/prometheus/src/main/java/org/opensearch/sql/prometheus/exceptions/PrometheusClientException.java b/prometheus/src/main/java/org/opensearch/sql/prometheus/exceptions/PrometheusClientException.java deleted file mode 100644 index 4f429e00e41..00000000000 --- a/prometheus/src/main/java/org/opensearch/sql/prometheus/exceptions/PrometheusClientException.java +++ /dev/null @@ -1,17 +0,0 @@ -/* - * - * * Copyright OpenSearch Contributors - * * SPDX-License-Identifier: Apache-2.0 - * - */ - -package org.opensearch.sql.prometheus.exceptions; - -import org.opensearch.sql.datasources.exceptions.DataSourceClientException; - -/** PrometheusClientException. */ -public class PrometheusClientException extends DataSourceClientException { - public PrometheusClientException(String message) { - super(message); - } -} diff --git a/prometheus/src/main/java/org/opensearch/sql/prometheus/request/system/PrometheusDescribeMetricRequest.java b/prometheus/src/main/java/org/opensearch/sql/prometheus/request/system/PrometheusDescribeMetricRequest.java index 17e5f21bd26..6049f8c0c33 100644 --- a/prometheus/src/main/java/org/opensearch/sql/prometheus/request/system/PrometheusDescribeMetricRequest.java +++ b/prometheus/src/main/java/org/opensearch/sql/prometheus/request/system/PrometheusDescribeMetricRequest.java @@ -27,7 +27,7 @@ import org.opensearch.sql.data.type.ExprCoreType; import org.opensearch.sql.data.type.ExprType; import org.opensearch.sql.prometheus.client.PrometheusClient; -import org.opensearch.sql.prometheus.exceptions.PrometheusClientException; +import org.opensearch.sql.prometheus.exception.PrometheusClientException; import org.opensearch.sql.prometheus.storage.PrometheusMetricDefaultSchema; /** diff --git a/prometheus/src/main/java/org/opensearch/sql/prometheus/request/system/PrometheusListMetricsRequest.java b/prometheus/src/main/java/org/opensearch/sql/prometheus/request/system/PrometheusListMetricsRequest.java index 0e6c2bb2c68..c45f25fad4f 100644 --- a/prometheus/src/main/java/org/opensearch/sql/prometheus/request/system/PrometheusListMetricsRequest.java +++ b/prometheus/src/main/java/org/opensearch/sql/prometheus/request/system/PrometheusListMetricsRequest.java @@ -23,7 +23,7 @@ import org.opensearch.sql.data.model.ExprTupleValue; import org.opensearch.sql.data.model.ExprValue; import org.opensearch.sql.prometheus.client.PrometheusClient; -import org.opensearch.sql.prometheus.request.system.model.MetricMetadata; +import org.opensearch.sql.prometheus.model.MetricMetadata; @RequiredArgsConstructor public class PrometheusListMetricsRequest implements PrometheusSystemRequest { diff --git a/prometheus/src/main/java/org/opensearch/sql/prometheus/storage/PrometheusStorageFactory.java b/prometheus/src/main/java/org/opensearch/sql/prometheus/storage/PrometheusStorageFactory.java index b9c1b09d601..484b897c784 100644 --- a/prometheus/src/main/java/org/opensearch/sql/prometheus/storage/PrometheusStorageFactory.java +++ b/prometheus/src/main/java/org/opensearch/sql/prometheus/storage/PrometheusStorageFactory.java @@ -7,8 +7,6 @@ package org.opensearch.sql.prometheus.storage; -import com.amazonaws.auth.AWSStaticCredentialsProvider; -import com.amazonaws.auth.BasicAWSCredentials; import java.net.URI; import java.net.URISyntaxException; import java.net.UnknownHostException; @@ -16,12 +14,7 @@ import java.security.PrivilegedAction; import java.util.Map; import java.util.Set; -import java.util.concurrent.TimeUnit; import lombok.RequiredArgsConstructor; -import okhttp3.OkHttpClient; -import org.opensearch.sql.common.interceptors.AwsSigningInterceptor; -import org.opensearch.sql.common.interceptors.BasicAuthenticationInterceptor; -import org.opensearch.sql.common.interceptors.URIValidatorInterceptor; import org.opensearch.sql.common.setting.Settings; import org.opensearch.sql.datasource.model.DataSource; import org.opensearch.sql.datasource.model.DataSourceMetadata; @@ -30,6 +23,7 @@ import org.opensearch.sql.datasources.utils.DatasourceValidationUtils; import org.opensearch.sql.prometheus.client.PrometheusClient; import org.opensearch.sql.prometheus.client.PrometheusClientImpl; +import org.opensearch.sql.prometheus.utils.PrometheusClientUtils; import org.opensearch.sql.storage.DataSourceFactory; import org.opensearch.sql.storage.StorageEngine; @@ -57,28 +51,6 @@ public DataSource createDataSource(DataSourceMetadata metadata) { metadata.getName(), DataSourceType.PROMETHEUS, getStorageEngine(metadata.getProperties())); } - // Need to refactor to a separate Validator class. - private void validateDataSourceConfigProperties(Map dataSourceMetadataConfig) - throws URISyntaxException, UnknownHostException { - if (dataSourceMetadataConfig.get(AUTH_TYPE) != null) { - AuthenticationType authenticationType = - AuthenticationType.get(dataSourceMetadataConfig.get(AUTH_TYPE)); - if (AuthenticationType.BASICAUTH.equals(authenticationType)) { - DatasourceValidationUtils.validateLengthAndRequiredFields( - dataSourceMetadataConfig, Set.of(URI, USERNAME, PASSWORD)); - } else if (AuthenticationType.AWSSIGV4AUTH.equals(authenticationType)) { - DatasourceValidationUtils.validateLengthAndRequiredFields( - dataSourceMetadataConfig, Set.of(URI, ACCESS_KEY, SECRET_KEY, REGION)); - } - } else { - DatasourceValidationUtils.validateLengthAndRequiredFields( - dataSourceMetadataConfig, Set.of(URI)); - } - DatasourceValidationUtils.validateHost( - dataSourceMetadataConfig.get(URI), - settings.getSettingValue(Settings.Key.DATASOURCES_URI_HOSTS_DENY_LIST)); - } - StorageEngine getStorageEngine(Map requiredConfig) { PrometheusClient prometheusClient; prometheusClient = @@ -88,7 +60,8 @@ StorageEngine getStorageEngine(Map requiredConfig) { try { validateDataSourceConfigProperties(requiredConfig); return new PrometheusClientImpl( - getHttpClient(requiredConfig), new URI(requiredConfig.get(URI))); + PrometheusClientUtils.getHttpClient(requiredConfig, settings), + new URI(requiredConfig.get(URI))); } catch (URISyntaxException | UnknownHostException e) { throw new IllegalArgumentException( String.format("Invalid URI in prometheus properties: %s", e.getMessage())); @@ -97,33 +70,25 @@ StorageEngine getStorageEngine(Map requiredConfig) { return new PrometheusStorageEngine(prometheusClient); } - private OkHttpClient getHttpClient(Map config) { - OkHttpClient.Builder okHttpClient = new OkHttpClient.Builder(); - okHttpClient.callTimeout(1, TimeUnit.MINUTES); - okHttpClient.connectTimeout(30, TimeUnit.SECONDS); - okHttpClient.followRedirects(false); - okHttpClient.addInterceptor( - new URIValidatorInterceptor( - settings.getSettingValue(Settings.Key.DATASOURCES_URI_HOSTS_DENY_LIST))); - if (config.get(AUTH_TYPE) != null) { - AuthenticationType authenticationType = AuthenticationType.get(config.get(AUTH_TYPE)); + // Need to refactor to a separate Validator class. + private void validateDataSourceConfigProperties(Map dataSourceMetadataConfig) + throws URISyntaxException, UnknownHostException { + if (dataSourceMetadataConfig.get(AUTH_TYPE) != null) { + AuthenticationType authenticationType = + AuthenticationType.get(dataSourceMetadataConfig.get(AUTH_TYPE)); if (AuthenticationType.BASICAUTH.equals(authenticationType)) { - okHttpClient.addInterceptor( - new BasicAuthenticationInterceptor(config.get(USERNAME), config.get(PASSWORD))); + DatasourceValidationUtils.validateLengthAndRequiredFields( + dataSourceMetadataConfig, Set.of(URI, USERNAME, PASSWORD)); } else if (AuthenticationType.AWSSIGV4AUTH.equals(authenticationType)) { - okHttpClient.addInterceptor( - new AwsSigningInterceptor( - new AWSStaticCredentialsProvider( - new BasicAWSCredentials(config.get(ACCESS_KEY), config.get(SECRET_KEY))), - config.get(REGION), - "aps")); - } else { - throw new IllegalArgumentException( - String.format( - "AUTH Type : %s is not supported with Prometheus Connector", - config.get(AUTH_TYPE))); + DatasourceValidationUtils.validateLengthAndRequiredFields( + dataSourceMetadataConfig, Set.of(URI, ACCESS_KEY, SECRET_KEY, REGION)); } + } else { + DatasourceValidationUtils.validateLengthAndRequiredFields( + dataSourceMetadataConfig, Set.of(URI)); } - return okHttpClient.build(); + DatasourceValidationUtils.validateHost( + dataSourceMetadataConfig.get(URI), + settings.getSettingValue(Settings.Key.DATASOURCES_URI_HOSTS_DENY_LIST)); } } diff --git a/prometheus/src/test/java/org/opensearch/sql/prometheus/client/PrometheusClientImplTest.java b/prometheus/src/test/java/org/opensearch/sql/prometheus/client/PrometheusClientImplTest.java deleted file mode 100644 index d6ca25b206c..00000000000 --- a/prometheus/src/test/java/org/opensearch/sql/prometheus/client/PrometheusClientImplTest.java +++ /dev/null @@ -1,223 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.sql.prometheus.client; - -import static org.junit.jupiter.api.Assertions.assertEquals; -import static org.junit.jupiter.api.Assertions.assertNotNull; -import static org.junit.jupiter.api.Assertions.assertThrows; -import static org.junit.jupiter.api.Assertions.assertTrue; -import static org.opensearch.sql.prometheus.constants.TestConstants.ENDTIME; -import static org.opensearch.sql.prometheus.constants.TestConstants.METRIC_NAME; -import static org.opensearch.sql.prometheus.constants.TestConstants.QUERY; -import static org.opensearch.sql.prometheus.constants.TestConstants.STARTTIME; -import static org.opensearch.sql.prometheus.constants.TestConstants.STEP; -import static org.opensearch.sql.prometheus.utils.TestUtils.getJson; - -import java.io.IOException; -import java.util.ArrayList; -import java.util.Collections; -import java.util.HashMap; -import java.util.List; -import java.util.Map; -import lombok.SneakyThrows; -import okhttp3.HttpUrl; -import okhttp3.OkHttpClient; -import okhttp3.mockwebserver.MockResponse; -import okhttp3.mockwebserver.MockWebServer; -import okhttp3.mockwebserver.RecordedRequest; -import org.json.JSONArray; -import org.json.JSONObject; -import org.junit.jupiter.api.AfterEach; -import org.junit.jupiter.api.BeforeEach; -import org.junit.jupiter.api.Test; -import org.junit.jupiter.api.extension.ExtendWith; -import org.mockito.junit.jupiter.MockitoExtension; -import org.opensearch.sql.prometheus.exceptions.PrometheusClientException; -import org.opensearch.sql.prometheus.request.system.model.MetricMetadata; - -@ExtendWith(MockitoExtension.class) -public class PrometheusClientImplTest { - - private MockWebServer mockWebServer; - private PrometheusClient prometheusClient; - - @BeforeEach - void setUp() throws IOException { - this.mockWebServer = new MockWebServer(); - this.mockWebServer.start(); - this.prometheusClient = - new PrometheusClientImpl(new OkHttpClient(), mockWebServer.url("").uri().normalize()); - } - - @Test - @SneakyThrows - void testQueryRange() { - MockResponse mockResponse = - new MockResponse() - .addHeader("Content-Type", "application/json; charset=utf-8") - .setBody(getJson("query_range_response.json")); - mockWebServer.enqueue(mockResponse); - JSONObject jsonObject = prometheusClient.queryRange(QUERY, STARTTIME, ENDTIME, STEP); - assertTrue(new JSONObject(getJson("query_range_result.json")).similar(jsonObject)); - RecordedRequest recordedRequest = mockWebServer.takeRequest(); - verifyQueryRangeCall(recordedRequest); - } - - @Test - @SneakyThrows - void testQueryRangeWith2xxStatusAndError() { - MockResponse mockResponse = - new MockResponse() - .addHeader("Content-Type", "application/json; charset=utf-8") - .setBody(getJson("error_response.json")); - mockWebServer.enqueue(mockResponse); - PrometheusClientException prometheusClientException = - assertThrows( - PrometheusClientException.class, - () -> prometheusClient.queryRange(QUERY, STARTTIME, ENDTIME, STEP)); - assertEquals("Error", prometheusClientException.getMessage()); - RecordedRequest recordedRequest = mockWebServer.takeRequest(); - verifyQueryRangeCall(recordedRequest); - } - - @Test - @SneakyThrows - void testQueryRangeWithNonJsonResponse() { - MockResponse mockResponse = - new MockResponse() - .addHeader("Content-Type", "application/json; charset=utf-8") - .setBody(getJson("non_json_response.json")); - mockWebServer.enqueue(mockResponse); - PrometheusClientException prometheusClientException = - assertThrows( - PrometheusClientException.class, - () -> prometheusClient.queryRange(QUERY, STARTTIME, ENDTIME, STEP)); - assertEquals( - "Prometheus returned unexpected body, " + "please verify your prometheus server setup.", - prometheusClientException.getMessage()); - RecordedRequest recordedRequest = mockWebServer.takeRequest(); - verifyQueryRangeCall(recordedRequest); - } - - @Test - @SneakyThrows - void testQueryRangeWithNon2xxError() { - MockResponse mockResponse = - new MockResponse() - .addHeader("Content-Type", "application/json; charset=utf-8") - .setResponseCode(400); - mockWebServer.enqueue(mockResponse); - PrometheusClientException prometheusClientException = - assertThrows( - PrometheusClientException.class, - () -> prometheusClient.queryRange(QUERY, STARTTIME, ENDTIME, STEP)); - assertEquals( - "Request to Prometheus is Unsuccessful with code : 400", - prometheusClientException.getMessage()); - RecordedRequest recordedRequest = mockWebServer.takeRequest(); - verifyQueryRangeCall(recordedRequest); - } - - @Test - @SneakyThrows - void testGetLabel() { - MockResponse mockResponse = - new MockResponse() - .addHeader("Content-Type", "application/json; charset=utf-8") - .setBody(getJson("get_labels_response.json")); - mockWebServer.enqueue(mockResponse); - List response = prometheusClient.getLabels(METRIC_NAME); - assertEquals( - new ArrayList() { - { - add("call"); - add("code"); - } - }, - response); - RecordedRequest recordedRequest = mockWebServer.takeRequest(); - verifyGetLabelsCall(recordedRequest); - } - - @Test - @SneakyThrows - void testGetAllMetrics() { - MockResponse mockResponse = - new MockResponse() - .addHeader("Content-Type", "application/json; charset=utf-8") - .setBody(getJson("all_metrics_response.json")); - mockWebServer.enqueue(mockResponse); - Map> response = prometheusClient.getAllMetrics(); - Map> expected = new HashMap<>(); - expected.put( - "go_gc_duration_seconds", - Collections.singletonList( - new MetricMetadata( - "summary", "A summary of the pause duration of garbage collection cycles.", ""))); - expected.put( - "go_goroutines", - Collections.singletonList( - new MetricMetadata("gauge", "Number of goroutines that currently exist.", ""))); - assertEquals(expected, response); - RecordedRequest recordedRequest = mockWebServer.takeRequest(); - verifyGetAllMetricsCall(recordedRequest); - } - - @Test - @SneakyThrows - void testQueryExemplars() { - MockResponse mockResponse = - new MockResponse() - .addHeader("Content-Type", "application/json; charset=utf-8") - .setBody(getJson("query_exemplars_response.json")); - mockWebServer.enqueue(mockResponse); - JSONArray jsonArray = prometheusClient.queryExemplars(QUERY, STARTTIME, ENDTIME); - assertTrue(new JSONArray(getJson("query_exemplars_result.json")).similar(jsonArray)); - RecordedRequest recordedRequest = mockWebServer.takeRequest(); - verifyQueryExemplarsCall(recordedRequest); - } - - @AfterEach - void tearDown() throws IOException { - mockWebServer.shutdown(); - } - - private void verifyQueryRangeCall(RecordedRequest recordedRequest) { - HttpUrl httpUrl = recordedRequest.getRequestUrl(); - assertEquals("GET", recordedRequest.getMethod()); - assertNotNull(httpUrl); - assertEquals("/api/v1/query_range", httpUrl.encodedPath()); - assertEquals(QUERY, httpUrl.queryParameter("query")); - assertEquals(STARTTIME.toString(), httpUrl.queryParameter("start")); - assertEquals(ENDTIME.toString(), httpUrl.queryParameter("end")); - assertEquals(STEP, httpUrl.queryParameter("step")); - } - - private void verifyGetLabelsCall(RecordedRequest recordedRequest) { - HttpUrl httpUrl = recordedRequest.getRequestUrl(); - assertEquals("GET", recordedRequest.getMethod()); - assertNotNull(httpUrl); - assertEquals("/api/v1/labels", httpUrl.encodedPath()); - assertEquals(METRIC_NAME, httpUrl.queryParameter("match[]")); - } - - private void verifyGetAllMetricsCall(RecordedRequest recordedRequest) { - HttpUrl httpUrl = recordedRequest.getRequestUrl(); - assertEquals("GET", recordedRequest.getMethod()); - assertNotNull(httpUrl); - assertEquals("/api/v1/metadata", httpUrl.encodedPath()); - } - - private void verifyQueryExemplarsCall(RecordedRequest recordedRequest) { - HttpUrl httpUrl = recordedRequest.getRequestUrl(); - assertEquals("GET", recordedRequest.getMethod()); - assertNotNull(httpUrl); - assertEquals("/api/v1/query_exemplars", httpUrl.encodedPath()); - assertEquals(QUERY, httpUrl.queryParameter("query")); - assertEquals(STARTTIME.toString(), httpUrl.queryParameter("start")); - assertEquals(ENDTIME.toString(), httpUrl.queryParameter("end")); - } -} diff --git a/prometheus/src/test/java/org/opensearch/sql/prometheus/functions/scan/QueryExemplarsFunctionTableScanBuilderTest.java b/prometheus/src/test/java/org/opensearch/sql/prometheus/functions/scan/QueryExemplarsFunctionTableScanBuilderTest.java index bb7806f824c..f037906a4a9 100644 --- a/prometheus/src/test/java/org/opensearch/sql/prometheus/functions/scan/QueryExemplarsFunctionTableScanBuilderTest.java +++ b/prometheus/src/test/java/org/opensearch/sql/prometheus/functions/scan/QueryExemplarsFunctionTableScanBuilderTest.java @@ -7,9 +7,9 @@ package org.opensearch.sql.prometheus.functions.scan; -import static org.opensearch.sql.prometheus.constants.TestConstants.ENDTIME; -import static org.opensearch.sql.prometheus.constants.TestConstants.QUERY; -import static org.opensearch.sql.prometheus.constants.TestConstants.STARTTIME; +import static org.opensearch.sql.prometheus.constant.TestConstants.ENDTIME; +import static org.opensearch.sql.prometheus.constant.TestConstants.QUERY; +import static org.opensearch.sql.prometheus.constant.TestConstants.STARTTIME; import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.Test; diff --git a/prometheus/src/test/java/org/opensearch/sql/prometheus/functions/scan/QueryExemplarsFunctionTableScanOperatorTest.java b/prometheus/src/test/java/org/opensearch/sql/prometheus/functions/scan/QueryExemplarsFunctionTableScanOperatorTest.java index 5b8cf34fc2f..856af85673a 100644 --- a/prometheus/src/test/java/org/opensearch/sql/prometheus/functions/scan/QueryExemplarsFunctionTableScanOperatorTest.java +++ b/prometheus/src/test/java/org/opensearch/sql/prometheus/functions/scan/QueryExemplarsFunctionTableScanOperatorTest.java @@ -11,9 +11,9 @@ import static org.junit.jupiter.api.Assertions.assertThrows; import static org.mockito.ArgumentMatchers.any; import static org.mockito.Mockito.when; -import static org.opensearch.sql.prometheus.constants.TestConstants.ENDTIME; -import static org.opensearch.sql.prometheus.constants.TestConstants.QUERY; -import static org.opensearch.sql.prometheus.constants.TestConstants.STARTTIME; +import static org.opensearch.sql.prometheus.constant.TestConstants.ENDTIME; +import static org.opensearch.sql.prometheus.constant.TestConstants.QUERY; +import static org.opensearch.sql.prometheus.constant.TestConstants.STARTTIME; import static org.opensearch.sql.prometheus.utils.TestUtils.getJson; import java.io.IOException; diff --git a/prometheus/src/test/java/org/opensearch/sql/prometheus/functions/scan/QueryRangeFunctionTableScanBuilderTest.java b/prometheus/src/test/java/org/opensearch/sql/prometheus/functions/scan/QueryRangeFunctionTableScanBuilderTest.java index dca79d69053..f3d9e366679 100644 --- a/prometheus/src/test/java/org/opensearch/sql/prometheus/functions/scan/QueryRangeFunctionTableScanBuilderTest.java +++ b/prometheus/src/test/java/org/opensearch/sql/prometheus/functions/scan/QueryRangeFunctionTableScanBuilderTest.java @@ -7,10 +7,10 @@ package org.opensearch.sql.prometheus.functions.scan; -import static org.opensearch.sql.prometheus.constants.TestConstants.ENDTIME; -import static org.opensearch.sql.prometheus.constants.TestConstants.QUERY; -import static org.opensearch.sql.prometheus.constants.TestConstants.STARTTIME; -import static org.opensearch.sql.prometheus.constants.TestConstants.STEP; +import static org.opensearch.sql.prometheus.constant.TestConstants.ENDTIME; +import static org.opensearch.sql.prometheus.constant.TestConstants.QUERY; +import static org.opensearch.sql.prometheus.constant.TestConstants.STARTTIME; +import static org.opensearch.sql.prometheus.constant.TestConstants.STEP; import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.Test; diff --git a/prometheus/src/test/java/org/opensearch/sql/prometheus/functions/scan/QueryRangeFunctionTableScanOperatorTest.java b/prometheus/src/test/java/org/opensearch/sql/prometheus/functions/scan/QueryRangeFunctionTableScanOperatorTest.java index e59a2bf7c4e..573f5167b85 100644 --- a/prometheus/src/test/java/org/opensearch/sql/prometheus/functions/scan/QueryRangeFunctionTableScanOperatorTest.java +++ b/prometheus/src/test/java/org/opensearch/sql/prometheus/functions/scan/QueryRangeFunctionTableScanOperatorTest.java @@ -11,10 +11,10 @@ import static org.junit.jupiter.api.Assertions.assertThrows; import static org.mockito.ArgumentMatchers.any; import static org.mockito.Mockito.when; -import static org.opensearch.sql.prometheus.constants.TestConstants.ENDTIME; -import static org.opensearch.sql.prometheus.constants.TestConstants.QUERY; -import static org.opensearch.sql.prometheus.constants.TestConstants.STARTTIME; -import static org.opensearch.sql.prometheus.constants.TestConstants.STEP; +import static org.opensearch.sql.prometheus.constant.TestConstants.ENDTIME; +import static org.opensearch.sql.prometheus.constant.TestConstants.QUERY; +import static org.opensearch.sql.prometheus.constant.TestConstants.STARTTIME; +import static org.opensearch.sql.prometheus.constant.TestConstants.STEP; import static org.opensearch.sql.prometheus.data.constants.PrometheusFieldConstants.LABELS; import static org.opensearch.sql.prometheus.data.constants.PrometheusFieldConstants.TIMESTAMP; import static org.opensearch.sql.prometheus.data.constants.PrometheusFieldConstants.VALUE; diff --git a/prometheus/src/test/java/org/opensearch/sql/prometheus/request/PrometheusDescribeMetricRequestTest.java b/prometheus/src/test/java/org/opensearch/sql/prometheus/request/PrometheusDescribeMetricRequestTest.java index a5cbcb7b7e0..79c1368e5de 100644 --- a/prometheus/src/test/java/org/opensearch/sql/prometheus/request/PrometheusDescribeMetricRequestTest.java +++ b/prometheus/src/test/java/org/opensearch/sql/prometheus/request/PrometheusDescribeMetricRequestTest.java @@ -11,7 +11,7 @@ import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; import static org.opensearch.sql.data.model.ExprValueUtils.stringValue; -import static org.opensearch.sql.prometheus.constants.TestConstants.METRIC_NAME; +import static org.opensearch.sql.prometheus.constant.TestConstants.METRIC_NAME; import static org.opensearch.sql.prometheus.data.constants.PrometheusFieldConstants.TIMESTAMP; import static org.opensearch.sql.prometheus.data.constants.PrometheusFieldConstants.VALUE; diff --git a/prometheus/src/test/java/org/opensearch/sql/prometheus/request/PrometheusListMetricsRequestTest.java b/prometheus/src/test/java/org/opensearch/sql/prometheus/request/PrometheusListMetricsRequestTest.java index 09f63463b54..89eb1957259 100644 --- a/prometheus/src/test/java/org/opensearch/sql/prometheus/request/PrometheusListMetricsRequestTest.java +++ b/prometheus/src/test/java/org/opensearch/sql/prometheus/request/PrometheusListMetricsRequestTest.java @@ -29,8 +29,8 @@ import org.opensearch.sql.data.model.ExprTupleValue; import org.opensearch.sql.data.model.ExprValue; import org.opensearch.sql.prometheus.client.PrometheusClient; +import org.opensearch.sql.prometheus.model.MetricMetadata; import org.opensearch.sql.prometheus.request.system.PrometheusListMetricsRequest; -import org.opensearch.sql.prometheus.request.system.model.MetricMetadata; @ExtendWith(MockitoExtension.class) public class PrometheusListMetricsRequestTest { diff --git a/prometheus/src/test/java/org/opensearch/sql/prometheus/storage/PrometheusMetricScanTest.java b/prometheus/src/test/java/org/opensearch/sql/prometheus/storage/PrometheusMetricScanTest.java index 00ddc973bca..40f8ba020a0 100644 --- a/prometheus/src/test/java/org/opensearch/sql/prometheus/storage/PrometheusMetricScanTest.java +++ b/prometheus/src/test/java/org/opensearch/sql/prometheus/storage/PrometheusMetricScanTest.java @@ -12,10 +12,10 @@ import static org.opensearch.sql.data.type.ExprCoreType.INTEGER; import static org.opensearch.sql.data.type.ExprCoreType.LONG; import static org.opensearch.sql.data.type.ExprCoreType.STRING; -import static org.opensearch.sql.prometheus.constants.TestConstants.ENDTIME; -import static org.opensearch.sql.prometheus.constants.TestConstants.QUERY; -import static org.opensearch.sql.prometheus.constants.TestConstants.STARTTIME; -import static org.opensearch.sql.prometheus.constants.TestConstants.STEP; +import static org.opensearch.sql.prometheus.constant.TestConstants.ENDTIME; +import static org.opensearch.sql.prometheus.constant.TestConstants.QUERY; +import static org.opensearch.sql.prometheus.constant.TestConstants.STARTTIME; +import static org.opensearch.sql.prometheus.constant.TestConstants.STEP; import static org.opensearch.sql.prometheus.data.constants.PrometheusFieldConstants.TIMESTAMP; import static org.opensearch.sql.prometheus.data.constants.PrometheusFieldConstants.VALUE; import static org.opensearch.sql.prometheus.utils.TestUtils.getJson; diff --git a/prometheus/src/test/java/org/opensearch/sql/prometheus/storage/PrometheusMetricTableTest.java b/prometheus/src/test/java/org/opensearch/sql/prometheus/storage/PrometheusMetricTableTest.java index 8bdab9244be..c6b9b63ec5e 100644 --- a/prometheus/src/test/java/org/opensearch/sql/prometheus/storage/PrometheusMetricTableTest.java +++ b/prometheus/src/test/java/org/opensearch/sql/prometheus/storage/PrometheusMetricTableTest.java @@ -54,7 +54,7 @@ import org.opensearch.sql.planner.physical.PhysicalPlan; import org.opensearch.sql.planner.physical.ProjectOperator; import org.opensearch.sql.prometheus.client.PrometheusClient; -import org.opensearch.sql.prometheus.constants.TestConstants; +import org.opensearch.sql.prometheus.constant.TestConstants; import org.opensearch.sql.prometheus.functions.scan.QueryRangeFunctionTableScanBuilder; import org.opensearch.sql.prometheus.request.PrometheusQueryRequest; import org.opensearch.sql.storage.read.TableScanBuilder; diff --git a/prometheus/src/test/java/org/opensearch/sql/prometheus/storage/QueryExemplarsTableTest.java b/prometheus/src/test/java/org/opensearch/sql/prometheus/storage/QueryExemplarsTableTest.java index 7f49de981a1..b7685280958 100644 --- a/prometheus/src/test/java/org/opensearch/sql/prometheus/storage/QueryExemplarsTableTest.java +++ b/prometheus/src/test/java/org/opensearch/sql/prometheus/storage/QueryExemplarsTableTest.java @@ -7,9 +7,9 @@ import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertTrue; -import static org.opensearch.sql.prometheus.constants.TestConstants.ENDTIME; -import static org.opensearch.sql.prometheus.constants.TestConstants.QUERY; -import static org.opensearch.sql.prometheus.constants.TestConstants.STARTTIME; +import static org.opensearch.sql.prometheus.constant.TestConstants.ENDTIME; +import static org.opensearch.sql.prometheus.constant.TestConstants.QUERY; +import static org.opensearch.sql.prometheus.constant.TestConstants.STARTTIME; import java.util.Map; import lombok.SneakyThrows; diff --git a/settings.gradle b/settings.gradle index ba38e3aa427..861a467d2b6 100644 --- a/settings.gradle +++ b/settings.gradle @@ -20,6 +20,8 @@ include 'benchmarks' include 'datasources' include 'async-query-core' include 'async-query' +include 'direct-query-core' +include 'direct-query' // exclude integ-test/doctest in case of offline build since they need downloads if (!gradle.startParameter.offline) {