Add TSDB query rewriter

9 files changed, 397 insertions, 0 deletions

diff --git a/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/restapi/horizon/TsdbQueryRewriter.java b/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/restapi/horizon/TsdbQueryRewriter.java
new file mode 100644
index 00000000000..40bc9145ce2
--- /dev/null
+++ b/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/restapi/horizon/TsdbQueryRewriter.java
@@ -0,0 +1,102 @@
+// Copyright Verizon Media. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+package com.yahoo.vespa.hosted.controller.restapi.horizon;
+
+import com.fasterxml.jackson.databind.JsonNode;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.fasterxml.jackson.databind.node.ArrayNode;
+import com.fasterxml.jackson.databind.node.ObjectNode;
+import com.yahoo.config.provision.SystemName;
+import com.yahoo.config.provision.TenantName;
+import com.yahoo.vespa.hosted.controller.api.role.Role;
+import com.yahoo.vespa.hosted.controller.api.role.RoleDefinition;
+import com.yahoo.vespa.hosted.controller.api.role.TenantRole;
+
+import java.io.IOException;
+import java.util.EnumSet;
+import java.util.Optional;
+import java.util.Set;
+import java.util.stream.Collectors;
+
+/**
+ * @author valerijf
+ */
+public class TsdbQueryRewriter {
+
+    private static final ObjectMapper mapper = new ObjectMapper();
+    private static final EnumSet<RoleDefinition> operatorRoleDefinitions =
+            EnumSet.of(RoleDefinition.hostedOperator, RoleDefinition.hostedSupporter);
+
+    public static byte[] rewrite(byte[] data, Set<Role> roles, SystemName systemName) throws IOException {
+        boolean operator = roles.stream().map(Role::definition).anyMatch(operatorRoleDefinitions::contains);
+
+        // Anyone with any tenant relation can view metrics for apps within those tenants
+        Set<TenantName> authorizedTenants = roles.stream()
+                .filter(TenantRole.class::isInstance)
+                .map(role -> ((TenantRole) role).tenant())
+                .collect(Collectors.toUnmodifiableSet());
+        if (!operator && authorizedTenants.isEmpty())
+            throw new UnauthorizedException();
+
+        JsonNode root = mapper.readTree(data);
+        getField(root, "executionGraph", ArrayNode.class)
+                .ifPresent(graph -> rewriteExecutionGraph(graph, authorizedTenants, operator, systemName));
+        getField(root, "filters", ArrayNode.class)
+                .ifPresent(filters -> rewriteFilters(filters, authorizedTenants, operator, systemName));
+
+        return mapper.writeValueAsBytes(root);
+    }
+
+    private static void rewriteExecutionGraph(ArrayNode executionGraph, Set<TenantName> tenantNames, boolean operator, SystemName systemName) {
+        for (int i = 0; i < executionGraph.size(); i++) {
+            JsonNode execution = executionGraph.get(i);
+
+            // Will be handled by rewriteFilters()
+            if (execution.has("filterId")) continue;
+
+            rewriteFilter((ObjectNode) execution, tenantNames, operator, systemName);
+        }
+    }
+
+    private static void rewriteFilters(ArrayNode filters, Set<TenantName> tenantNames, boolean operator, SystemName systemName) {
+        for (int i = 0; i < filters.size(); i++)
+            rewriteFilter((ObjectNode) filters.get(i), tenantNames, operator, systemName);
+    }
+
+    private static void rewriteFilter(ObjectNode parent, Set<TenantName> tenantNames, boolean operator, SystemName systemName) {
+        ObjectNode prev = ((ObjectNode) parent.get("filter"));
+        ArrayNode filters;
+        // If we dont already have a filter object, or the object that we have is not an AND filter
+        if (prev == null || !"Chain".equals(prev.get("type").asText()) || !"AND".equals(prev.get("op").asText())) {
+            // Create new filter object
+            filters = parent.putObject("filter")
+                    .put("type", "Chain")
+                    .put("op", "AND")
+                    .putArray("filters");
+
+            // Add the previous filter to the AND expression
+            if (prev != null) filters.add(prev);
+        } else filters = (ArrayNode) prev.get("filters");
+
+        // Make sure we only show metrics in the relevant system
+        ObjectNode systemFilter = filters.addObject();
+        systemFilter.put("type", "TagValueLiteralOr");
+        systemFilter.put("filter", systemName.name());
+        systemFilter.put("tagKey", "system");
+
+        // Make sure non-operators cannot see metrics outside of their tenants
+        if (!operator) {
+            ObjectNode appFilter = filters.addObject();
+            appFilter.put("type", "TagValueRegex");
+            appFilter.put("filter",
+                    tenantNames.stream().map(TenantName::value).sorted().collect(Collectors.joining("|", "(", ")\\..*")));
+            appFilter.put("tagKey", "applicationId");
+        }
+    }
+
+    private static <T extends JsonNode> Optional<T> getField(JsonNode object, String fieldName, Class<T> clazz) {
+        return Optional.ofNullable(object.get(fieldName)).filter(clazz::isInstance).map(clazz::cast);
+    }
+
+    static class UnauthorizedException extends RuntimeException { }
+
+}
diff --git a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/restapi/horizon/TsdbQueryRewriterTest.java b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/restapi/horizon/TsdbQueryRewriterTest.java
new file mode 100644
index 00000000000..937b0e95440
--- /dev/null
+++ b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/restapi/horizon/TsdbQueryRewriterTest.java
@@ -0,0 +1,55 @@
+package com.yahoo.vespa.hosted.controller.restapi.horizon;
+
+import com.yahoo.config.provision.SystemName;
+import com.yahoo.config.provision.TenantName;
+import com.yahoo.slime.JsonFormat;
+import com.yahoo.slime.SlimeUtils;
+import com.yahoo.vespa.hosted.controller.api.role.Role;
+import org.junit.Test;
+
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.nio.file.Files;
+import java.nio.file.Paths;
+import java.util.Set;
+
+import static org.junit.Assert.assertEquals;
+
+/**
+ * @author valerijf
+ */
+public class TsdbQueryRewriterTest {
+
+    @Test
+    public void rewrites_query() throws IOException {
+        assertRewrite("filters-complex.json", "filters-complex.expected.json", Role.reader(TenantName.from("tenant2")));
+
+        assertRewrite("filter-in-execution-graph.json",
+                "filter-in-execution-graph.expected.json",
+                Role.reader(TenantName.from("tenant2")), Role.athenzTenantAdmin(TenantName.from("tenant3")));
+
+        assertRewrite("filter-in-execution-graph.json",
+                "filter-in-execution-graph.expected.operator.json",
+                Role.reader(TenantName.from("tenant2")), Role.athenzTenantAdmin(TenantName.from("tenant3")), Role.hostedOperator());
+
+        assertRewrite("no-filters.json",
+                "no-filters.expected.json",
+                Role.reader(TenantName.from("tenant2")), Role.athenzTenantAdmin(TenantName.from("tenant3")));
+    }
+
+    @Test(expected = TsdbQueryRewriter.UnauthorizedException.class)
+    public void throws_if_no_roles() throws IOException {
+        assertRewrite("filters-complex.json", "filters-complex.expected.json");
+    }
+
+    private static void assertRewrite(String initialFilename, String expectedFilename, Role... roles) throws IOException {
+        byte[] data = Files.readAllBytes(Paths.get("src/test/resources/horizon", initialFilename));
+        data = TsdbQueryRewriter.rewrite(data, Set.of(roles), SystemName.Public);
+
+        ByteArrayOutputStream baos = new ByteArrayOutputStream();
+        new JsonFormat(false).encode(baos, SlimeUtils.jsonToSlime(data));
+        String expectedJson = Files.readString(Paths.get("src/test/resources/horizon", expectedFilename));
+
+        assertEquals(expectedJson, baos.toString());
+    }
+}
\ No newline at end of file
diff --git a/controller-server/src/test/resources/horizon/filter-in-execution-graph.expected.json b/controller-server/src/test/resources/horizon/filter-in-execution-graph.expected.json
new file mode 100644
index 00000000000..7c279442f1d
--- /dev/null
+++ b/controller-server/src/test/resources/horizon/filter-in-execution-graph.expected.json
@@ -0,0 +1,37 @@
+{
+ "start": 1619301600000,
+ "end": 1623161217471,
+ "executionGraph": [
+  {
+   "id": "q1_m1",
+   "type": "TimeSeriesDataSource",
+   "metric": {
+    "type": "MetricLiteral",
+    "metric": "Vespa.vespa.distributor.vds.distributor.docsstored.average"
+   },
+   "sourceId": null,
+   "fetchLast": false,
+   "filter": {
+    "type": "Chain",
+    "op": "AND",
+    "filters": [
+     {
+      "type": "TagValueLiteralOr",
+      "filter": "tenant1.application1.instance1",
+      "tagKey": "applicationId"
+     },
+     {
+      "type": "TagValueLiteralOr",
+      "filter": "Public",
+      "tagKey": "system"
+     },
+     {
+      "type": "TagValueRegex",
+      "filter": "(tenant2|tenant3)\\..*",
+      "tagKey": "applicationId"
+     }
+    ]
+   }
+  }
+ ]
+}
diff --git a/controller-server/src/test/resources/horizon/filter-in-execution-graph.expected.operator.json b/controller-server/src/test/resources/horizon/filter-in-execution-graph.expected.operator.json
new file mode 100644
index 00000000000..5eb5b403b2f
--- /dev/null
+++ b/controller-server/src/test/resources/horizon/filter-in-execution-graph.expected.operator.json
@@ -0,0 +1,32 @@
+{
+ "start": 1619301600000,
+ "end": 1623161217471,
+ "executionGraph": [
+  {
+   "id": "q1_m1",
+   "type": "TimeSeriesDataSource",
+   "metric": {
+    "type": "MetricLiteral",
+    "metric": "Vespa.vespa.distributor.vds.distributor.docsstored.average"
+   },
+   "sourceId": null,
+   "fetchLast": false,
+   "filter": {
+    "type": "Chain",
+    "op": "AND",
+    "filters": [
+     {
+      "type": "TagValueLiteralOr",
+      "filter": "tenant1.application1.instance1",
+      "tagKey": "applicationId"
+     },
+     {
+      "type": "TagValueLiteralOr",
+      "filter": "Public",
+      "tagKey": "system"
+     }
+    ]
+   }
+  }
+ ]
+}
diff --git a/controller-server/src/test/resources/horizon/filter-in-execution-graph.json b/controller-server/src/test/resources/horizon/filter-in-execution-graph.json
new file mode 100644
index 00000000000..6a2512c3642
--- /dev/null
+++ b/controller-server/src/test/resources/horizon/filter-in-execution-graph.json
@@ -0,0 +1,21 @@
+{
+  "start": 1619301600000,
+  "end": 1623161217471,
+  "executionGraph": [
+    {
+      "id": "q1_m1",
+      "type": "TimeSeriesDataSource",
+      "metric": {
+        "type": "MetricLiteral",
+        "metric": "Vespa.vespa.distributor.vds.distributor.docsstored.average"
+      },
+      "sourceId": null,
+      "fetchLast": false,
+      "filter": {
+        "type": "TagValueLiteralOr",
+        "filter": "tenant1.application1.instance1",
+        "tagKey": "applicationId"
+      }
+    }
+  ]
+}
\ No newline at end of file
diff --git a/controller-server/src/test/resources/horizon/filters-complex.expected.json b/controller-server/src/test/resources/horizon/filters-complex.expected.json
new file mode 100644
index 00000000000..333e79150e4
--- /dev/null
+++ b/controller-server/src/test/resources/horizon/filters-complex.expected.json
@@ -0,0 +1,56 @@
+{
+ "start": 1623080040000,
+ "end": 1623166440000,
+ "executionGraph": [
+  {
+   "id": "q1_m1",
+   "type": "TimeSeriesDataSource",
+   "metric": {
+    "type": "MetricLiteral",
+    "metric": "Vespa.vespa.qrserver.documents_covered.count"
+   },
+   "sourceId": null,
+   "fetchLast": false,
+   "filterId": "filter-ni8"
+  }
+ ],
+ "filters": [
+  {
+   "filter": {
+    "type": "Chain",
+    "op": "AND",
+    "filters": [
+     {
+      "type": "NOT",
+      "filter": {
+       "type": "TagValueLiteralOr",
+       "filter": "tenant1.app1.instance1",
+       "tagKey": "applicationId"
+      }
+     },
+     {
+      "type": "TagValueLiteralOr",
+      "filter": "Public",
+      "tagKey": "system"
+     },
+     {
+      "type": "TagValueRegex",
+      "filter": "(tenant2)\\..*",
+      "tagKey": "applicationId"
+     }
+    ]
+   },
+   "id": "filter-ni8"
+  }
+ ],
+ "serdesConfigs": [
+  {
+   "id": "JsonV3QuerySerdes",
+   "filter": [
+    "summarizer"
+   ]
+  }
+ ],
+ "logLevel": "ERROR",
+ "cacheMode": null
+}
diff --git a/controller-server/src/test/resources/horizon/filters-complex.json b/controller-server/src/test/resources/horizon/filters-complex.json
new file mode 100644
index 00000000000..3acc7fe5044
--- /dev/null
+++ b/controller-server/src/test/resources/horizon/filters-complex.json
@@ -0,0 +1,46 @@
+{
+  "start": 1623080040000,
+  "end": 1623166440000,
+  "executionGraph": [
+    {
+      "id": "q1_m1",
+      "type": "TimeSeriesDataSource",
+      "metric": {
+        "type": "MetricLiteral",
+        "metric": "Vespa.vespa.qrserver.documents_covered.count"
+      },
+      "sourceId": null,
+      "fetchLast": false,
+      "filterId": "filter-ni8"
+    }
+  ],
+  "filters": [
+    {
+      "filter": {
+        "type": "Chain",
+        "op": "AND",
+        "filters": [
+          {
+            "type": "NOT",
+            "filter": {
+              "type": "TagValueLiteralOr",
+              "filter": "tenant1.app1.instance1",
+              "tagKey": "applicationId"
+            }
+          }
+        ]
+      },
+      "id": "filter-ni8"
+    }
+  ],
+  "serdesConfigs": [
+    {
+      "id": "JsonV3QuerySerdes",
+      "filter": [
+        "summarizer"
+      ]
+    }
+  ],
+  "logLevel": "ERROR",
+  "cacheMode": null
+}
diff --git a/controller-server/src/test/resources/horizon/no-filters.expected.json b/controller-server/src/test/resources/horizon/no-filters.expected.json
new file mode 100644
index 00000000000..d2d6407e6b4
--- /dev/null
+++ b/controller-server/src/test/resources/horizon/no-filters.expected.json
@@ -0,0 +1,32 @@
+{
+ "start": 1619301600000,
+ "end": 1623161217471,
+ "executionGraph": [
+  {
+   "id": "q1_m1",
+   "type": "TimeSeriesDataSource",
+   "metric": {
+    "type": "MetricLiteral",
+    "metric": "Vespa.vespa.distributor.vds.distributor.docsstored.average"
+   },
+   "sourceId": null,
+   "fetchLast": false,
+   "filter": {
+    "type": "Chain",
+    "op": "AND",
+    "filters": [
+     {
+      "type": "TagValueLiteralOr",
+      "filter": "Public",
+      "tagKey": "system"
+     },
+     {
+      "type": "TagValueRegex",
+      "filter": "(tenant2|tenant3)\\..*",
+      "tagKey": "applicationId"
+     }
+    ]
+   }
+  }
+ ]
+}
diff --git a/controller-server/src/test/resources/horizon/no-filters.json b/controller-server/src/test/resources/horizon/no-filters.json
new file mode 100644
index 00000000000..3ff80feba02
--- /dev/null
+++ b/controller-server/src/test/resources/horizon/no-filters.json
@@ -0,0 +1,16 @@
+{
+  "start": 1619301600000,
+  "end": 1623161217471,
+  "executionGraph": [
+    {
+      "id": "q1_m1",
+      "type": "TimeSeriesDataSource",
+      "metric": {
+        "type": "MetricLiteral",
+        "metric": "Vespa.vespa.distributor.vds.distributor.docsstored.average"
+      },
+      "sourceId": null,
+      "fetchLast": false
+    }
+  ]
+}
\ No newline at end of file