hack(starfish): cardinality analyzer (#4359)

* add minimal query UI

* WIP: added backend but server is not working

* add roles for tool

snuba/admin/auth_roles.py

@@ -79,6 +79,7 @@ class InteractToolAction(ToolAction):
 TOOL_RESOURCES = {
     "snql-to-sql": ToolResource("snql-to-sql"),
     "tracing": ToolResource("tracing"),
+    "cardinality-analyzer": ToolResource("cardinality-analyzer"),
     "all": ToolResource("all"),
 }
 
@@ -176,6 +177,10 @@ def generate_tool_test_role(tool: str) -> Role:
             )
         },
     ),
+    "CardinalityAnalyzer": Role(
+        name="cardinality-analyzer",
+        actions={InteractToolAction([TOOL_RESOURCES["cardinality-analyzer"]])},
+    ),
 }
 
 DEFAULT_ROLES = [

snuba/admin/cardinality_analyzer/cardinality_analyzer.py

@@ -0,0 +1,48 @@
+from snuba.admin.audit_log.querylog import audit_log
+from snuba.admin.clickhouse.common import (
+    get_ro_query_node_connection,
+    validate_ro_query,
+)
+from snuba.clickhouse.native import ClickhouseResult
+from snuba.clusters.cluster import ClickhouseClientSettings
+from snuba.datasets.schemas.tables import TableSchema
+from snuba.datasets.storages.factory import get_storage
+from snuba.datasets.storages.storage_key import StorageKey
+
+# HACK (VOLO): Everything in this file is a hack
+
+
+@audit_log
+def run_metrics_query(query: str, user: str) -> ClickhouseResult:
+    """
+    Validates, audit logs, and executes given query against Querylog
+    table in ClickHouse. `user` param is necessary for audit_log
+    decorator.
+    """
+    schema = get_storage(StorageKey("generic_metrics_distributions")).get_schema()
+    assert isinstance(schema, TableSchema)
+    validate_ro_query(
+        sql_query=query,
+        allowed_tables={
+            schema.get_table_name(),
+            "generic_metric_distributions_aggregated_dist",
+        },
+    )
+    return __run_query(query)
+
+
+def __run_query(query: str) -> ClickhouseResult:
+    """
+    Runs given Query against metrics distributions in ClickHouse. This function assumes valid
+    query and does not validate/sanitize query or response data.
+    """
+    connection = get_ro_query_node_connection(
+        StorageKey("generic_metrics_distributions").value,
+        ClickhouseClientSettings.TRACING,
+    )
+
+    query_result = connection.execute(
+        query=query,
+        with_column_types=True,
+    )
+    return query_result

snuba/admin/static/api_client.tsx

@@ -25,6 +25,7 @@ import { SnQLRequest, SnQLResult, SnubaDatasetName } from "./snql_to_sql/types";
 
 import { KafkaTopicData } from "./kafka/types";
 import { QuerylogRequest, QuerylogResult } from "./querylog/types";
+import { CardinalityQueryRequest, CardinalityQueryResult } from "./cardinality_analyzer/types";
 
 import { AllocationPolicy } from "./capacity_management/types";
 
@@ -56,6 +57,7 @@ interface Client {
   getPredefinedQuerylogOptions: () => Promise<[PredefinedQuery]>;
   getQuerylogSchema: () => Promise<QuerylogResult>;
   executeQuerylogQuery: (req: QuerylogRequest) => Promise<QuerylogResult>;
+  executeCardinalityQuery: (req: CardinalityQueryRequest) => Promise<CardinalityQueryResult>;
   getAllMigrationGroups: () => Promise<MigrationGroupResult[]>;
   runMigration: (req: RunMigrationRequest) => Promise<RunMigrationResult>;
   getAllowedTools: () => Promise<AllowedTools>;
@@ -267,6 +269,20 @@ function Client() {
         }
       });
     },
+    executeCardinalityQuery: (query: CardinalityQueryRequest) => {
+      const url = baseUrl + "cardinality_query";
+      return fetch(url, {
+        headers: { "Content-Type": "application/json" },
+        method: "POST",
+        body: JSON.stringify(query),
+      }).then((resp) => {
+        if (resp.ok) {
+          return resp.json();
+        } else {
+          return resp.json().then(Promise.reject.bind(Promise));
+        }
+      });
+    },
     getAllMigrationGroups: () => {
       const url = baseUrl + "migrations/groups";
       return fetch(url, {

snuba/admin/static/cardinality_analyzer/index.tsx

@@ -0,0 +1,53 @@
+import React, { useState, useEffect } from "react";
+import Client from "../api_client";
+import { Table } from "../table";
+import QueryDisplay from "./query_display";
+import { CardinalityQueryResult, PredefinedQuery } from "./types";
+
+function CardinalityQueries(props: { api: Client }) {
+  const [predefinedQueryOptions, setPredefinedQueryOptions] = useState<
+    PredefinedQuery[]
+  >([]);
+
+
+  function tablePopulator(queryResult: CardinalityQueryResult) {
+    return (
+      <div style={scroll}>
+        <Table
+          headerData={queryResult.column_names}
+          rowData={queryResult.rows}
+        />
+      </div>
+    );
+  }
+
+  function formatSQL(sql: string) {
+    const formatted = sql
+      .split("\n")
+      .map((line) => line.substring(4, line.length))
+      .join("\n");
+    return formatted.trim();
+  }
+
+  return (
+    <div>
+      {QueryDisplay({
+        api: props.api,
+        resultDataPopulator: tablePopulator,
+        predefinedQueryOptions: [],
+      })}
+    </div>
+  );
+}
+
+const selectStyle = {
+  marginBottom: 8,
+  height: 30,
+};
+
+const scroll = {
+  overflowX: "scroll" as const,
+  width: "100%",
+};
+
+export default CardinalityQueries;

snuba/admin/static/cardinality_analyzer/query_display.tsx

@@ -0,0 +1,172 @@
+import React, { useState } from "react";
+import Client from "../api_client";
+import { Collapse } from "../collapse";
+import QueryEditor from "../query_editor";
+
+import { CardinalityQueryRequest, CardinalityQueryResult, PredefinedQuery } from "./types";
+
+type QueryState = Partial<CardinalityQueryRequest>;
+
+function QueryDisplay(props: {
+  api: Client;
+  resultDataPopulator: (queryResult: CardinalityQueryResult) => JSX.Element;
+  predefinedQueryOptions: Array<PredefinedQuery>;
+}) {
+  const [query, setQuery] = useState<QueryState>({});
+  const [queryResultHistory, setCardinalityQueryResultHistory] = useState<
+    CardinalityQueryResult[]
+  >([]);
+
+  function updateQuerySql(sql: string) {
+    setQuery((prevQuery) => {
+      return {
+        ...prevQuery,
+        sql,
+      };
+    });
+  }
+
+  function executeQuery() {
+    props.api
+      .executeCardinalityQuery(query as CardinalityQueryRequest)
+      .then((result) => {
+        result.input_query = query.sql || "<Input Query>";
+        setCardinalityQueryResultHistory((prevHistory) => [result, ...prevHistory]);
+      })
+      .catch((err) => {
+        console.log("ERROR", err);
+        window.alert("An error occurred: " + err.error.message);
+      });
+  }
+
+  function convertResultsToCSV(queryResult: CardinalityQueryResult) {
+    let output = queryResult.column_names.join(",");
+    for (const row of queryResult.rows) {
+      const escaped = row.map((v) =>
+        typeof v == "string" && v.includes(",") ? '"' + v + '"' : v
+      );
+      output = output + "\n" + escaped.join(",");
+    }
+    return output;
+  }
+
+  function copyText(queryResult: CardinalityQueryResult, format: string) {
+    const formatter = format == "csv" ? convertResultsToCSV : JSON.stringify;
+    window.navigator.clipboard.writeText(formatter(queryResult));
+  }
+
+  return (
+    <div>
+      <h2>Construct a Metrics Query</h2>
+      <QueryEditor
+        onQueryUpdate={(sql) => {
+          updateQuerySql(sql);
+        }}
+        predefinedQueryOptions={props.predefinedQueryOptions}
+      />
+      <div style={executeActionsStyle}>
+        <div>
+          <button
+            onClick={(evt) => {
+              evt.preventDefault();
+              executeQuery();
+            }}
+            style={executeButtonStyle}
+            disabled={!query.sql}
+          >
+            Execute Query
+          </button>
+        </div>
+      </div>
+      <div>
+        <h2>Query results</h2>
+        {queryResultHistory.map((queryResult, idx) => {
+          if (idx === 0) {
+            return (
+              <div key={idx}>
+                <p>{queryResult.input_query}</p>
+                <p>
+                  <button
+                    style={executeButtonStyle}
+                    onClick={() => copyText(queryResult, "json")}
+                  >
+                    Copy to clipboard (JSON)
+                  </button>
+                </p>
+                <p>
+                  <button
+                    style={executeButtonStyle}
+                    onClick={() => copyText(queryResult, "csv")}
+                  >
+                    Copy to clipboard (CSV)
+                  </button>
+                </p>
+                {props.resultDataPopulator(queryResult)}
+              </div>
+            );
+          }
+
+          return (
+            <Collapse key={idx} text={queryResult.input_query}>
+              <button
+                style={executeButtonStyle}
+                onClick={() => copyText(queryResult, "json")}
+              >
+                Copy to clipboard (JSON)
+              </button>
+              <button
+                style={executeButtonStyle}
+                onClick={() => copyText(queryResult, "csv")}
+              >
+                Copy to clipboard (CSV)
+              </button>
+              {props.resultDataPopulator(queryResult)}
+            </Collapse>
+          );
+        })}
+      </div>
+    </div>
+  );
+}
+
+const executeActionsStyle = {
+  display: "flex",
+  justifyContent: "space-between",
+  marginTop: 8,
+};
+
+const executeButtonStyle = {
+  height: 30,
+  border: 0,
+  padding: "4px 20px",
+  marginRight: 10,
+};
+
+const selectStyle = {
+  marginRight: 8,
+  height: 30,
+};
+
+function TextArea(props: {
+  value: string;
+  onChange: (nextValue: string) => void;
+}) {
+  const { value, onChange } = props;
+  return (
+    <textarea
+      spellCheck={false}
+      value={value}
+      onChange={(evt) => onChange(evt.target.value)}
+      style={{ width: "100%", height: 140 }}
+      placeholder={"Write your query here"}
+    />
+  );
+}
+
+const queryDescription = {
+  minHeight: 10,
+  width: "auto",
+  fontSize: 16,
+  padding: "10px 5px",
+};
+export default QueryDisplay;

snuba/admin/static/cardinality_analyzer/types.tsx

@@ -0,0 +1,22 @@
+type QueryResultColumnMetadata = [string];
+type QueryResultRow = [string];
+
+type CardinalityQueryRequest= {
+  sql: string;
+};
+
+type CardinalityQueryResult = {
+  input_query: string;
+  timestamp: number;
+  column_names: QueryResultColumnMetadata;
+  rows: [QueryResultRow];
+  error?: string;
+};
+
+type PredefinedQuery = {
+  name: string;
+  sql: string;
+  description: string;
+};
+
+export { CardinalityQueryRequest, CardinalityQueryResult, PredefinedQuery };

snuba/admin/static/data.tsx

@@ -8,6 +8,7 @@ import Kafka from "./kafka";
 import QuerylogQueries from "./querylog";
 import CapacityManagement from "./capacity_management";
 import DeadLetterQueue from "./dead_letter_queue";
+import CardinalityAnalyzer from "./cardinality_analyzer";
 import ProductionQueries from "./production_queries";
 
 function Placeholder(props: any) {
@@ -62,6 +63,11 @@ const NAV_ITEMS = [
     display: "♻️ Dead Letter Queue",
     component: DeadLetterQueue,
   },
+  {
+    id: "cardinality-analyzer",
+    display: "🔢Cardinality Analyzer!!!",
+    component: CardinalityAnalyzer,
+  },
   {
     id: "production-queries",
     display: "🔦 Production Queries",

snuba/admin/tool_policies.py

@@ -27,6 +27,7 @@ class AdminTools(Enum):
     AUDIT_LOG = "audit-log"
     KAFKA = "kafka"
     CAPACITY_MANAGEMENT = "capacity-management"
+    CARDINALITY_ANALYZER = "cardinality-analyzer"
 
 
 DEVELOPER_TOOLS: set[AdminTools] = {AdminTools.SNQL_TO_SQL, AdminTools.QUERY_TRACING}