Fix bugs in the code and support a different ruler uri

Signed-off-by: Ashish Agrawal <ashisagr@amazon.com>

Author: lezzago

direct-query-core/src/main/java/org/opensearch/sql/prometheus/client/PrometheusClient.java

@@ -58,14 +58,15 @@ Map<String, List<MetricMetadata>> getAllMetrics(Map<String, String> queryParams)
   JSONObject getAlerts() throws IOException;
 
   /**
-   * Get all recording and alerting rules. Returns the raw response body since the format varies by
-   * backend: Prometheus returns JSON, Cortex/Thanos returns YAML, and AMP returns JSON.
+   * Get all recording and alerting rules, normalized to a consistent JSON format. Handles
+   * Prometheus JSON, Cortex/Thanos YAML, and AMP JSON responses, returning them all as a
+   * {"groups":[...]} structure.
    *
    * @param queryParams Map of query parameters to include in the request
-   * @return String containing the raw response body
+   * @return JSONObject with {"groups":[...]} structure
    * @throws IOException If there is an issue with the request
    */
-  String getRules(Map<String, String> queryParams) throws IOException;
+  JSONObject getRules(Map<String, String> queryParams) throws IOException;
 
   /**
    * Get all alerts from Alertmanager.
@@ -110,15 +111,16 @@ Map<String, List<MetricMetadata>> getAllMetrics(Map<String, String> queryParams)
   String createAlertmanagerSilences(String silenceJson) throws IOException;
 
   /**
-   * Get rules for a specific namespace from the Cortex/Thanos Ruler API. The response is returned
-   * as a raw string since the Ruler API returns YAML (Cortex/Thanos) or JSON (AMP).
+   * Get rules for a specific namespace, normalized to a consistent JSON format. Handles
+   * Cortex/Thanos YAML and AMP JSON responses, returning them all as a {"groups":[...]} structure.
    *
    * @param namespace The rules namespace
    * @param queryParams Map of query parameters to include in the request
-   * @return String containing the raw response body (YAML or JSON)
+   * @return JSONObject with {"groups":[...]} structure
    * @throws IOException If there is an issue with the request
    */
-  String getRulesByNamespace(String namespace, Map<String, String> queryParams) throws IOException;
+  JSONObject getRulesByNamespace(String namespace, Map<String, String> queryParams)
+      throws IOException;
 
   /**
    * Create or update a rule group in a namespace via the Cortex/Thanos Ruler API.

direct-query-core/src/main/java/org/opensearch/sql/prometheus/client/PrometheusClientImpl.java

@@ -7,6 +7,7 @@
 
 import com.fasterxml.jackson.core.type.TypeReference;
 import com.fasterxml.jackson.databind.ObjectMapper;
+import com.fasterxml.jackson.dataformat.yaml.YAMLFactory;
 import java.io.IOException;
 import java.net.URI;
 import java.net.URLEncoder;
@@ -36,30 +37,44 @@
 public class PrometheusClientImpl implements PrometheusClient {
 
   private static final Logger logger = LogManager.getLogger(PrometheusClientImpl.class);
+  private static final ObjectMapper YAML_MAPPER = new ObjectMapper(new YAMLFactory());
 
   private final OkHttpClient prometheusHttpClient;
   private final OkHttpClient alertmanagerHttpClient;
 
   private final URI prometheusUri;
   private final URI alertmanagerUri;
+  private final URI rulerUri;
 
   public PrometheusClientImpl(OkHttpClient prometheusHttpClient, URI prometheusUri) {
     this(
         prometheusHttpClient,
         prometheusUri,
         prometheusHttpClient,
-        URI.create(prometheusUri.toString().replaceAll("/$", "") + "/alertmanager"));
+        URI.create(prometheusUri.toString().replaceAll("/$", "") + "/alertmanager"),
+        prometheusUri);
   }
 
   public PrometheusClientImpl(
       OkHttpClient prometheusHttpClient,
       URI prometheusUri,
       OkHttpClient alertmanagerHttpClient,
       URI alertmanagerUri) {
+    this(prometheusHttpClient, prometheusUri, alertmanagerHttpClient, alertmanagerUri,
+        prometheusUri);
+  }
+
+  public PrometheusClientImpl(
+      OkHttpClient prometheusHttpClient,
+      URI prometheusUri,
+      OkHttpClient alertmanagerHttpClient,
+      URI alertmanagerUri,
+      URI rulerUri) {
     this.prometheusHttpClient = prometheusHttpClient;
     this.prometheusUri = prometheusUri;
     this.alertmanagerHttpClient = alertmanagerHttpClient;
     this.alertmanagerUri = alertmanagerUri;
+    this.rulerUri = rulerUri;
   }
 
   private String paramsToQueryString(Map<String, String> queryParams) {
@@ -229,38 +244,18 @@ public JSONObject getAlerts() throws IOException {
   }
 
   @Override
-  public String getRules(Map<String, String> queryParams) throws IOException {
+  public JSONObject getRules(Map<String, String> queryParams) throws IOException {
     String queryString = this.paramsToQueryString(queryParams);
     String queryUrl =
         String.format(
             "%s/api/v1/rules%s", prometheusUri.toString().replaceAll("/$", ""), queryString);
-    logger.debug("Making Ruler GET request for all rules: {}", queryUrl);
+    logger.debug("Making Ruler GET request for all rules");
     Request request = new Request.Builder().url(queryUrl).build();
-    Response response =
+    try (Response response =
         AccessController.doPrivilegedChecked(
-            () -> this.prometheusHttpClient.newCall(request).execute());
-
-    if (response.isSuccessful()) {
-      String bodyString = Objects.requireNonNull(response.body()).string();
-      // Try to extract data from Prometheus JSON format; return raw body otherwise
-      try {
-        JSONObject jsonObject = new JSONObject(bodyString);
-        if ("success".equals(jsonObject.optString("status"))
-            && jsonObject.has("data")) {
-          return jsonObject.getJSONObject("data").toString();
-        }
-      } catch (JSONException e) {
-        // Not JSON (e.g. Cortex YAML) — fall through
-      }
-      return bodyString;
-    } else {
-      String errorBody = response.body() != null ? response.body().string() : "No response body";
-      logger.error(
-          "Ruler GET request failed with code: {}, error body: {}", response.code(), errorBody);
-      throw new PrometheusClientException(
-          String.format(
-              "Ruler request failed with code: %s. Error details: %s",
-              response.code(), errorBody));
+            () -> this.prometheusHttpClient.newCall(request).execute())) {
+      String body = readRulerResponse(response, "GET all rules");
+      return normalizeRulesResponse(body, null);
     }
   }
 
@@ -341,31 +336,22 @@ public String createAlertmanagerSilences(String silenceJson) throws IOException
   }
 
   @Override
-  public String getRulesByNamespace(String namespace, Map<String, String> queryParams)
+  public JSONObject getRulesByNamespace(String namespace, Map<String, String> queryParams)
       throws IOException {
     String queryString = this.paramsToQueryString(queryParams);
     String queryUrl =
         String.format(
             "%s/api/v1/rules/%s%s",
-            prometheusUri.toString().replaceAll("/$", ""),
+            rulerUri.toString().replaceAll("/$", ""),
             URLEncoder.encode(namespace, StandardCharsets.UTF_8),
             queryString);
-    logger.debug("Making Ruler GET request for namespace: {}", queryUrl);
+    logger.debug("Making Ruler GET request for namespace");
     Request request = new Request.Builder().url(queryUrl).build();
-    Response response =
+    try (Response response =
         AccessController.doPrivilegedChecked(
-            () -> this.prometheusHttpClient.newCall(request).execute());
-
-    if (response.isSuccessful()) {
-      return Objects.requireNonNull(response.body()).string();
-    } else {
-      String errorBody = response.body() != null ? response.body().string() : "No response body";
-      logger.error(
-          "Ruler GET request failed with code: {}, error body: {}", response.code(), errorBody);
-      throw new PrometheusClientException(
-          String.format(
-              "Ruler request failed with code: %s. Error details: %s",
-              response.code(), errorBody));
+            () -> this.prometheusHttpClient.newCall(request).execute())) {
+      String body = readRulerResponse(response, "GET namespace " + namespace);
+      return normalizeRulesResponse(body, namespace);
     }
   }
 
@@ -374,30 +360,20 @@ public String createOrUpdateRuleGroup(String namespace, String yamlBody) throws
     String queryUrl =
         String.format(
             "%s/api/v1/rules/%s",
-            prometheusUri.toString().replaceAll("/$", ""),
+            rulerUri.toString().replaceAll("/$", ""),
             URLEncoder.encode(namespace, StandardCharsets.UTF_8));
-    logger.debug("Making Ruler POST request to create/update rule group: {}", queryUrl);
+    logger.debug("Making Ruler POST request to create/update rule group");
     Request request =
         new Request.Builder()
             .url(queryUrl)
             .header("Content-Type", "application/yaml")
             .post(RequestBody.create(yamlBody.getBytes(StandardCharsets.UTF_8)))
             .build();
-    Response response =
+    try (Response response =
         AccessController.doPrivilegedChecked(
-            () -> this.prometheusHttpClient.newCall(request).execute());
-
-    if (response.isSuccessful()) {
-      String body = response.body() != null ? response.body().string() : "";
+            () -> this.prometheusHttpClient.newCall(request).execute())) {
+      String body = readRulerResponse(response, "POST create/update rule group");
       return body.isEmpty() ? "{\"status\":\"success\"}" : body;
-    } else {
-      String errorBody = response.body() != null ? response.body().string() : "No response body";
-      logger.error(
-          "Ruler POST request failed with code: {}, error body: {}", response.code(), errorBody);
-      throw new PrometheusClientException(
-          String.format(
-              "Ruler request failed with code: %s. Error details: %s",
-              response.code(), errorBody));
     }
   }
 
@@ -406,24 +382,15 @@ public String deleteRuleNamespace(String namespace) throws IOException {
     String queryUrl =
         String.format(
             "%s/api/v1/rules/%s",
-            prometheusUri.toString().replaceAll("/$", ""),
+            rulerUri.toString().replaceAll("/$", ""),
             URLEncoder.encode(namespace, StandardCharsets.UTF_8));
-    logger.debug("Making Ruler DELETE request for namespace: {}", queryUrl);
+    logger.debug("Making Ruler DELETE request for namespace");
     Request request = new Request.Builder().url(queryUrl).delete().build();
-    Response response =
+    try (Response response =
         AccessController.doPrivilegedChecked(
-            () -> this.prometheusHttpClient.newCall(request).execute());
-
-    if (response.isSuccessful()) {
-      return "{\"status\":\"success\"}";
-    } else {
-      String errorBody = response.body() != null ? response.body().string() : "No response body";
-      logger.error(
-          "Ruler DELETE request failed with code: {}, error body: {}", response.code(), errorBody);
-      throw new PrometheusClientException(
-          String.format(
-              "Ruler request failed with code: %s. Error details: %s",
-              response.code(), errorBody));
+            () -> this.prometheusHttpClient.newCall(request).execute())) {
+      String body = readRulerResponse(response, "DELETE namespace " + namespace);
+      return body.isEmpty() ? "{\"status\":\"success\"}" : body;
     }
   }
 
@@ -432,28 +399,116 @@ public String deleteRuleGroup(String namespace, String groupName) throws IOExcep
     String queryUrl =
         String.format(
             "%s/api/v1/rules/%s/%s",
-            prometheusUri.toString().replaceAll("/$", ""),
+            rulerUri.toString().replaceAll("/$", ""),
             URLEncoder.encode(namespace, StandardCharsets.UTF_8),
             URLEncoder.encode(groupName, StandardCharsets.UTF_8));
-    logger.debug("Making Ruler DELETE request for group: {}", queryUrl);
+    logger.debug("Making Ruler DELETE request for group");
     Request request = new Request.Builder().url(queryUrl).delete().build();
-    Response response =
+    try (Response response =
         AccessController.doPrivilegedChecked(
-            () -> this.prometheusHttpClient.newCall(request).execute());
+            () -> this.prometheusHttpClient.newCall(request).execute())) {
+      String body = readRulerResponse(response, "DELETE group " + groupName);
+      return body.isEmpty() ? "{\"status\":\"success\"}" : body;
+    }
+  }
 
+  /**
+   * Reads a Ruler API response, returning the body string on success or throwing on failure.
+   * Consolidates the error-handling pattern shared by all Ruler methods.
+   *
+   * @param response The HTTP response
+   * @param operationDescription Description for log messages (e.g., "GET all rules")
+   * @return The response body as a string (empty string if body is null)
+   * @throws IOException If there is an issue reading the response
+   */
+  private String readRulerResponse(Response response, String operationDescription)
+      throws IOException {
     if (response.isSuccessful()) {
-      return "{\"status\":\"success\"}";
+      return Objects.requireNonNull(response.body(), "Ruler response body is null").string();
     } else {
       String errorBody = response.body() != null ? response.body().string() : "No response body";
       logger.error(
-          "Ruler DELETE request failed with code: {}, error body: {}", response.code(), errorBody);
+          "Ruler {} request failed with code: {}, error body: {}",
+          operationDescription,
+          response.code(),
+          errorBody);
       throw new PrometheusClientException(
           String.format(
               "Ruler request failed with code: %s. Error details: %s",
               response.code(), errorBody));
     }
   }
 
+  /**
+   * Normalizes a raw rule response body into a consistent {"groups":[...]} JSONObject. Handles
+   * three response formats:
+   *
+   * <ul>
+   *   <li>Prometheus/AMP JSON: {"status":"success","data":{"groups":[...]}} - extracts data
+   *   <li>Cortex/Thanos YAML (all rules): Map of namespace to list of rule groups
+   *   <li>Cortex/Thanos YAML (single namespace): List of rule groups
+   * </ul>
+   *
+   * @param body The raw response body string
+   * @param namespace Optional namespace name used as the "file" field on groups from YAML
+   *     single-namespace responses. When null, indicates the body may be a YAML map of namespaces.
+   * @return JSONObject with {"groups":[...]} structure
+   */
+  @SuppressWarnings("unchecked")
+  private JSONObject normalizeRulesResponse(String body, String namespace) {
+    if (body.isEmpty()) {
+      return new JSONObject().put("groups", new JSONArray());
+    }
+
+    // 1. Try Prometheus/AMP JSON format first
+    try {
+      JSONObject jsonObject = new JSONObject(body);
+      if ("success".equals(jsonObject.optString("status")) && jsonObject.has("data")) {
+        return jsonObject.getJSONObject("data");
+      }
+      if (jsonObject.has("groups")) {
+        return jsonObject;
+      }
+    } catch (JSONException e) {
+      // Not JSON — fall through to YAML parsing
+    }
+
+    // 2. Parse as YAML (Cortex/Thanos format)
+    try {
+      Object parsed = YAML_MAPPER.readValue(body, Object.class);
+      JSONArray groupsArray = new JSONArray();
+      addGroupsFromParsed(parsed, namespace, groupsArray);
+      return new JSONObject().put("groups", groupsArray);
+    } catch (Exception e) {
+      logger.warn(
+          "Failed to parse rules response body, returning empty groups: {}", e.getMessage());
+      return new JSONObject().put("groups", new JSONArray());
+    }
+  }
+
+  @SuppressWarnings("unchecked")
+  private void addGroupsFromParsed(Object parsed, String namespace, JSONArray groupsArray) {
+    if (parsed instanceof Map) {
+      // All-namespaces format: Map<String, List<RuleGroup>>
+      Map<String, Object> namespacesMap = (Map<String, Object>) parsed;
+      for (Map.Entry<String, Object> entry : namespacesMap.entrySet()) {
+        if (entry.getValue() instanceof List) {
+          addGroupsFromParsed(entry.getValue(), entry.getKey(), groupsArray);
+        }
+      }
+    } else if (parsed instanceof List) {
+      // Single-namespace format: List<RuleGroup>
+      List<Map<String, Object>> groups = (List<Map<String, Object>>) parsed;
+      for (Map<String, Object> group : groups) {
+        JSONObject groupObj = new JSONObject(group);
+        if (namespace != null) {
+          groupObj.put("file", namespace);
+        }
+        groupsArray.put(groupObj);
+      }
+    }
+  }
+
   /**
    * Reads and processes an Alertmanager API response.
    *

direct-query-core/src/main/java/org/opensearch/sql/prometheus/query/PrometheusQueryHandler.java

@@ -148,15 +148,15 @@ public GetDirectQueryResourcesResponse<?> getResources(
           }
         case RULES:
           {
-            String rulesBody;
+            JSONObject rules;
             if (request.getResourceName() != null && !request.getResourceName().isEmpty()) {
-              rulesBody =
+              rules =
                   client.getRulesByNamespace(
                       request.getResourceName(), request.getQueryParams());
             } else {
-              rulesBody = client.getRules(request.getQueryParams());
+              rules = client.getRules(request.getQueryParams());
             }
-            return GetDirectQueryResourcesResponse.withStringList(List.of(rulesBody));
+            return GetDirectQueryResourcesResponse.withMap(rules.toMap());
           }
         case ALERTMANAGER_ALERTS:
           {