Prometheus Connector Initial Code

prometheus/build.gradle

@@ -0,0 +1,82 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+plugins {
+    id 'java-library'
+    id "io.freefair.lombok"
+    id 'jacoco'
+}
+
+dependencies {
+    api project(':core')
+    api group: 'org.opensearch', name: 'opensearch', version: "${opensearch_version}"
+    implementation "io.github.resilience4j:resilience4j-retry:1.5.0"
+    implementation group: 'com.fasterxml.jackson.core', name: 'jackson-core', version: "${jackson_version}"
+    implementation group: 'com.fasterxml.jackson.core', name: 'jackson-databind', version: "${jackson_version}"
+    implementation group: 'com.fasterxml.jackson.dataformat', name: 'jackson-dataformat-cbor', version: "${jackson_version}"
+    compileOnly group: 'org.opensearch.client', name: 'opensearch-rest-high-level-client', version: "${opensearch_version}"
+    api group: 'com.squareup.okhttp3', name: 'okhttp', version: '4.9.3'
+    implementation group: 'org.json', name: 'json', version: '20180813'
+
+    testImplementation('org.junit.jupiter:junit-jupiter:5.6.2')
+    testImplementation group: 'org.hamcrest', name: 'hamcrest-library', version: '2.1'
+    testImplementation group: 'org.mockito', name: 'mockito-core', version: '3.12.4'
+    testImplementation group: 'org.mockito', name: 'mockito-junit-jupiter', version: '3.12.4'
+    testImplementation group: 'org.opensearch.client', name: 'opensearch-rest-high-level-client', version: "${opensearch_version}"
+    testImplementation group: 'org.opensearch.test', name: 'framework', version: "${opensearch_version}"
+    testImplementation group: 'com.squareup.okhttp3', name: 'mockwebserver', version: '4.9.3'
+}
+
+test {
+    useJUnitPlatform()
+    testLogging {
+        events "passed", "skipped", "failed"
+        exceptionFormat "full"
+    }
+}
+
+configurations.all {
+    resolutionStrategy.force "org.jetbrains.kotlin:kotlin-stdlib:1.6.0"
+    resolutionStrategy.force "org.jetbrains.kotlin:kotlin-stdlib-common:1.6.0"
+}
+
+jacocoTestReport {
+    reports {
+        html.enabled true
+        xml.enabled true
+    }
+    afterEvaluate {
+        classDirectories.setFrom(files(classDirectories.files.collect {
+            fileTree(dir: it)
+        }))
+    }
+}
+test.finalizedBy(project.tasks.jacocoTestReport)
+
+jacocoTestCoverageVerification {
+    violationRules {
+        rule {
+            element = 'CLASS'
+            excludes = [
+                    'org.opensearch.sql.prometheus.data.constants.*'
+            ]
+            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

prometheus/lombok.config

@@ -0,0 +1,3 @@
+# This file is generated by the 'io.freefair.lombok' Gradle plugin
+config.stopBubbling = true
+lombok.addLombokGeneratedAnnotation = true

prometheus/src/main/java/org/opensearch/sql/prometheus/client/PrometheusClient.java

@@ -0,0 +1,17 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.prometheus.client;
+
+import java.io.IOException;
+import java.util.List;
+import org.json.JSONObject;
+
+public interface PrometheusClient {
+
+  JSONObject queryRange(String query, Long start, Long end, String step) throws IOException;
+
+  List<String> getLabels(String metricName) throws IOException;
+}

prometheus/src/main/java/org/opensearch/sql/prometheus/client/PrometheusClientImpl.java

@@ -0,0 +1,95 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.prometheus.client;
+
+import java.io.IOException;
+import java.net.URI;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Objects;
+import okhttp3.HttpUrl;
+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.JSONObject;
+
+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 {
+    HttpUrl httpUrl = new HttpUrl.Builder()
+        .scheme(uri.getScheme())
+        .host(uri.getHost())
+        .port(uri.getPort())
+        .addPathSegment("api")
+        .addPathSegment("v1")
+        .addPathSegment("query_range")
+        .addQueryParameter("query", query)
+        .addQueryParameter("start", Long.toString(start))
+        .addQueryParameter("end", Long.toString(end))
+        .addQueryParameter("step", step)
+        .build();
+    logger.debug("queryUrl: " + httpUrl);
+    Request request = new Request.Builder()
+        .url(httpUrl)
+        .build();
+    Response response = this.okHttpClient.newCall(request).execute();
+    JSONObject jsonObject = readResponse(response);
+    return jsonObject.getJSONObject("data");
+  }
+
+  @Override
+  public List<String> getLabels(String metricName) throws IOException {
+    String queryUrl = String.format("%sapi/v1/labels?match[]=%s", uri.toString(), metricName);
+    logger.debug("queryUrl: " + queryUrl);
+    Request request = new Request.Builder()
+        .url(queryUrl)
+        .build();
+    Response response = this.okHttpClient.newCall(request).execute();
+    JSONObject jsonObject = readResponse(response);
+    return toListOfStrings(jsonObject.getJSONArray("data"));
+  }
+
+  private List<String> toListOfStrings(JSONArray array) {
+    List<String> result = new ArrayList<>();
+    for (int i = 0; i < array.length(); i++) {
+      result.add(array.optString(i));
+    }
+    return result;
+  }
+
+
+  private JSONObject readResponse(Response response) throws IOException {
+    if (response.isSuccessful()) {
+      JSONObject jsonObject = new JSONObject(Objects.requireNonNull(response.body()).string());
+      if ("success".equals(jsonObject.getString("status"))) {
+        return jsonObject;
+      } else {
+        throw new RuntimeException(jsonObject.getString("error"));
+      }
+    } else {
+      throw new RuntimeException(
+          String.format("Request to Prometheus is Unsuccessful with : %s", response.message()));
+    }
+  }
+
+
+}

prometheus/src/main/java/org/opensearch/sql/prometheus/data/constants/PrometheusFieldConstants.java

@@ -0,0 +1,12 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.prometheus.data.constants;
+
+public class PrometheusFieldConstants {
+  public static final String TIMESTAMP = "@timestamp";
+  public static final String VALUE = "@value";
+  public static final String METRIC = "metric";
+}

prometheus/src/main/java/org/opensearch/sql/prometheus/request/PrometheusDescribeMetricRequest.java

@@ -0,0 +1,77 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+
+package org.opensearch.sql.prometheus.request;
+
+import static org.opensearch.sql.prometheus.data.constants.PrometheusFieldConstants.METRIC;
+import static org.opensearch.sql.prometheus.data.constants.PrometheusFieldConstants.TIMESTAMP;
+import static org.opensearch.sql.prometheus.data.constants.PrometheusFieldConstants.VALUE;
+
+import java.io.IOException;
+import java.security.AccessController;
+import java.security.PrivilegedAction;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import lombok.ToString;
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
+import org.opensearch.sql.data.type.ExprCoreType;
+import org.opensearch.sql.data.type.ExprType;
+import org.opensearch.sql.prometheus.client.PrometheusClient;
+
+/**
+ * Describe Metric metadata request.
+ * This is triggered in case of both query range table function and relation.
+ * In case of table function metric name is null.
+ */
+@ToString(onlyExplicitlyIncluded = true)
+public class PrometheusDescribeMetricRequest {
+
+  private final PrometheusClient prometheusClient;
+
+  @ToString.Include
+  private final Optional<String> metricName;
+
+  private static final Logger LOG = LogManager.getLogger();
+
+
+  public PrometheusDescribeMetricRequest(PrometheusClient prometheusClient,
+                                         String metricName) {
+    this.prometheusClient = prometheusClient;
+    this.metricName = Optional.ofNullable(metricName);
+  }
+
+
+  /**
+   * Get the mapping of field and type.
+   *
+   * @return mapping of field and type.
+   */
+  public Map<String, ExprType> getFieldTypes() {
+    Map<String, ExprType> fieldTypes = new HashMap<>();
+    AccessController.doPrivileged((PrivilegedAction<List<Void>>) () -> {
+      if (metricName.isPresent()) {
+        try {
+          prometheusClient.getLabels(metricName.get())
+              .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.get(), e.getMessage()));
+        }
+      }
+      return null;
+    });
+    fieldTypes.put(VALUE, ExprCoreType.DOUBLE);
+    fieldTypes.put(TIMESTAMP, ExprCoreType.TIMESTAMP);
+    fieldTypes.put(METRIC, ExprCoreType.STRING);
+    return fieldTypes;
+  }
+
+}

prometheus/src/main/java/org/opensearch/sql/prometheus/request/PrometheusQueryRequest.java

@@ -0,0 +1,54 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+
+package org.opensearch.sql.prometheus.request;
+
+import lombok.EqualsAndHashCode;
+import lombok.Getter;
+import lombok.Setter;
+import lombok.ToString;
+import org.opensearch.common.unit.TimeValue;
+
+/**
+ * Prometheus metric query request.
+ */
+@EqualsAndHashCode
+@Getter
+@ToString
+public class PrometheusQueryRequest {
+
+  public static final TimeValue DEFAULT_QUERY_TIMEOUT = TimeValue.timeValueMinutes(1L);
+
+  /**
+   * PromQL.
+   */
+  private final StringBuilder promQl;
+
+  /**
+   * startTime of the query.
+   */
+  @Setter
+  private Long startTime;
+
+  /**
+   * endTime of the query.
+   */
+  @Setter
+  private Long endTime;
+
+  /**
+   * step is the resolution required between startTime and endTime.
+   */
+  @Setter
+  private String step;
+
+  /**
+   * Constructor of PrometheusQueryRequest.
+   */
+  public PrometheusQueryRequest() {
+    this.promQl = new StringBuilder();
+  }
+}