diff --git a/datasources/src/main/java/org/opensearch/sql/datasources/exceptions/DataSourceClientException.java b/datasources/src/main/java/org/opensearch/sql/datasources/exceptions/DataSourceClientException.java new file mode 100644 index 0000000000..a7a0218761 --- /dev/null +++ b/datasources/src/main/java/org/opensearch/sql/datasources/exceptions/DataSourceClientException.java @@ -0,0 +1,19 @@ +/* + * + * * Copyright OpenSearch Contributors + * * SPDX-License-Identifier: Apache-2.0 + * + */ + +package org.opensearch.sql.datasources.exceptions; + +/** + * This is the base class for all DataSource Client Exceptions. All datasource connector modules + * will extend this exception for their respective connection clients. + */ +public class DataSourceClientException extends RuntimeException { + + public DataSourceClientException(String message) { + super(message); + } +} diff --git a/plugin/src/main/java/org/opensearch/sql/plugin/rest/RestPPLQueryAction.java b/plugin/src/main/java/org/opensearch/sql/plugin/rest/RestPPLQueryAction.java index 0f4d2f7d0c..922b58d6f4 100644 --- a/plugin/src/main/java/org/opensearch/sql/plugin/rest/RestPPLQueryAction.java +++ b/plugin/src/main/java/org/opensearch/sql/plugin/rest/RestPPLQueryAction.java @@ -18,6 +18,7 @@ import java.util.function.Supplier; import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; +import org.opensearch.OpenSearchSecurityException; import org.opensearch.action.ActionListener; import org.opensearch.client.node.NodeClient; import org.opensearch.index.IndexNotFoundException; @@ -28,6 +29,7 @@ import org.opensearch.rest.RestStatus; import org.opensearch.sql.common.antlr.SyntaxCheckException; import org.opensearch.sql.common.setting.Settings; +import org.opensearch.sql.datasources.exceptions.DataSourceClientException; import org.opensearch.sql.exception.ExpressionEvaluationException; import org.opensearch.sql.exception.QueryEngineException; import org.opensearch.sql.exception.SemanticCheckException; @@ -67,7 +69,9 @@ private static boolean isClientError(Exception e) { || e instanceof SemanticCheckException || e instanceof ExpressionEvaluationException || e instanceof QueryEngineException - || e instanceof SyntaxCheckException; + || e instanceof SyntaxCheckException + || e instanceof DataSourceClientException + || e instanceof IllegalAccessException; } @Override @@ -129,8 +133,9 @@ public void onFailure(Exception e) { channel, INTERNAL_SERVER_ERROR, "Failed to explain the query due to error: " + e.getMessage()); - } else if (e instanceof IllegalAccessException) { - reportError(channel, e, BAD_REQUEST); + } else if (e instanceof OpenSearchSecurityException) { + OpenSearchSecurityException exception = (OpenSearchSecurityException) e; + reportError(channel, exception, exception.status()); } else { LOG.error("Error happened during query handling", e); if (isClientError(e)) { 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 index 512c014564..6249e96ad8 100644 --- a/prometheus/src/main/java/org/opensearch/sql/prometheus/client/PrometheusClientImpl.java +++ b/prometheus/src/main/java/org/opensearch/sql/prometheus/client/PrometheusClientImpl.java @@ -22,7 +22,9 @@ 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 { @@ -98,18 +100,26 @@ private List toListOfLabels(JSONArray array) { private JSONObject readResponse(Response response) throws IOException { if (response.isSuccessful()) { - JSONObject jsonObject = new JSONObject(Objects.requireNonNull(response.body()).string()); + 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 RuntimeException(jsonObject.getString("error")); + throw new PrometheusClientException(jsonObject.getString("error")); } } else { - throw new RuntimeException( - String.format("Request to Prometheus is Unsuccessful with : %s", Objects.requireNonNull( - response.body(), "Response body can't be null").string())); + throw new PrometheusClientException( + String.format( + "Request to Prometheus is Unsuccessful with : %s", + Objects.requireNonNull(response.body(), "Response body can't be null").string())); } } -} \ No newline at end of file +} 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 new file mode 100644 index 0000000000..4f429e00e4 --- /dev/null +++ b/prometheus/src/main/java/org/opensearch/sql/prometheus/exceptions/PrometheusClientException.java @@ -0,0 +1,17 @@ +/* + * + * * 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 2e0d46b3e8..4b94b511db 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 @@ -28,6 +28,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.storage.PrometheusMetricDefaultSchema; /** @@ -73,18 +74,25 @@ public PrometheusDescribeMetricRequest(PrometheusClient prometheusClient, */ public Map getFieldTypes() { Map fieldTypes = new HashMap<>(); - AccessController.doPrivileged((PrivilegedAction>) () -> { - try { - prometheusClient.getLabels(metricName) - .forEach(label -> fieldTypes.put(label, ExprCoreType.STRING)); - } catch (IOException e) { - LOG.error("Error while fetching labels for {} from prometheus: {}", - metricName, e.getMessage()); - throw new RuntimeException(String.format("Error while fetching labels " - + "for %s from prometheus: %s", metricName, e.getMessage())); - } - return null; - }); + AccessController.doPrivileged( + (PrivilegedAction>) + () -> { + try { + prometheusClient + .getLabels(metricName) + .forEach(label -> fieldTypes.put(label, ExprCoreType.STRING)); + } catch (IOException e) { + LOG.error( + "Error while fetching labels for {} from prometheus: {}", + metricName, + e.getMessage()); + throw new PrometheusClientException( + String.format( + "Error while fetching labels " + "for %s from prometheus: %s", + metricName, e.getMessage())); + } + return null; + }); fieldTypes.putAll(PrometheusMetricDefaultSchema.DEFAULT_MAPPING.getMapping()); return fieldTypes; } 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 b3ecd25af3..2a30d6634b 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 @@ -101,6 +101,7 @@ private OkHttpClient getHttpClient(Map config) { OkHttpClient.Builder okHttpClient = new OkHttpClient.Builder(); okHttpClient.callTimeout(1, TimeUnit.MINUTES); okHttpClient.connectTimeout(30, TimeUnit.SECONDS); + okHttpClient.followRedirects(false); if (config.get(AUTH_TYPE) != null) { AuthenticationType authenticationType = AuthenticationType.get(config.get(AUTH_TYPE)); if (AuthenticationType.BASICAUTH.equals(authenticationType)) { @@ -158,8 +159,9 @@ private void validateURI(Map config) throws URISyntaxException { Matcher matcher = allowHostsPattern.matcher(host); if (!matcher.matches()) { throw new IllegalArgumentException( - String.format("Disallowed hostname in the uri: %s. Validate with %s config", - config.get(URI), Settings.Key.DATASOURCES_URI_ALLOWHOSTS.getKeyValue())); + String.format( + "Disallowed hostname in the uri. Validate with %s config", + Settings.Key.DATASOURCES_URI_ALLOWHOSTS.getKeyValue())); } } } 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 index 76abb05751..0f5ad2ae5c 100644 --- a/prometheus/src/test/java/org/opensearch/sql/prometheus/client/PrometheusClientImplTest.java +++ b/prometheus/src/test/java/org/opensearch/sql/prometheus/client/PrometheusClientImplTest.java @@ -34,6 +34,7 @@ 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) @@ -72,10 +73,30 @@ void testQueryRangeWith2xxStatusAndError() { .addHeader("Content-Type", "application/json; charset=utf-8") .setBody(getJson("error_response.json")); mockWebServer.enqueue(mockResponse); - RuntimeException runtimeException - = assertThrows(RuntimeException.class, - () -> prometheusClient.queryRange(QUERY, STARTTIME, ENDTIME, STEP)); - assertEquals("Error", runtimeException.getMessage()); + 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); } @@ -87,11 +108,14 @@ void testQueryRangeWithNon2xxError() { .addHeader("Content-Type", "application/json; charset=utf-8") .setResponseCode(400); mockWebServer.enqueue(mockResponse); - RuntimeException runtimeException - = assertThrows(RuntimeException.class, - () -> prometheusClient.queryRange(QUERY, STARTTIME, ENDTIME, STEP)); + PrometheusClientException prometheusClientException = + assertThrows( + PrometheusClientException.class, + () -> prometheusClient.queryRange(QUERY, STARTTIME, ENDTIME, STEP)); assertTrue( - runtimeException.getMessage().contains("Request to Prometheus is Unsuccessful with :")); + prometheusClientException + .getMessage() + .contains("Request to Prometheus is Unsuccessful with :")); RecordedRequest recordedRequest = mockWebServer.takeRequest(); verifyQueryRangeCall(recordedRequest); } diff --git a/prometheus/src/test/java/org/opensearch/sql/prometheus/storage/PrometheusStorageFactoryTest.java b/prometheus/src/test/java/org/opensearch/sql/prometheus/storage/PrometheusStorageFactoryTest.java index d6a934a015..0925910731 100644 --- a/prometheus/src/test/java/org/opensearch/sql/prometheus/storage/PrometheusStorageFactoryTest.java +++ b/prometheus/src/test/java/org/opensearch/sql/prometheus/storage/PrometheusStorageFactoryTest.java @@ -259,8 +259,11 @@ void createDataSourceWithHostnameNotMatchingWithAllowHostsConfig() { RuntimeException exception = Assertions.assertThrows(RuntimeException.class, () -> prometheusStorageFactory.createDataSource(metadata)); Assertions.assertTrue( - exception.getMessage().contains("Disallowed hostname in the uri: http://localhost.com:9090. " - + "Validate with plugins.query.datasources.uri.allowhosts config")); + exception + .getMessage() + .contains( + "Disallowed hostname in the uri. " + + "Validate with plugins.query.datasources.uri.allowhosts config")); } @Test diff --git a/prometheus/src/test/resources/non_json_response.json b/prometheus/src/test/resources/non_json_response.json new file mode 100644 index 0000000000..8a2586e7b7 --- /dev/null +++ b/prometheus/src/test/resources/non_json_response.json @@ -0,0 +1 @@ +fadsfadsfasdfasdfadsfasdfadsfadsf \ No newline at end of file