Merge branch 'master' into vlad/doc_admin

snuba/admin/clickhouse/system_queries.py

@@ -1,8 +1,9 @@
+import re
 from dataclasses import dataclass
-from typing import Any, Dict, Optional, Sequence, Tuple, Type, cast
+from typing import Dict, Optional, Sequence, Type
 
 from snuba import settings
-from snuba.clickhouse.native import ClickhousePool
+from snuba.clickhouse.native import ClickhousePool, ClickhouseResult
 from snuba.clusters.cluster import ClickhouseClientSettings, ClickhouseCluster
 from snuba.datasets.storages import StorageKey
 from snuba.datasets.storages.factory import get_storage
@@ -21,6 +22,14 @@ class InvalidStorageError(SerializableException):
     pass
 
 
+class InvalidResultError(SerializableException):
+    pass
+
+
+class InvalidCustomQuery(SerializableException):
+    pass
+
+
 class _QueryRegistry:
     """Keep a mapping of SystemQueries to their names"""
 
@@ -90,21 +99,16 @@ class ActivePartitions(SystemQuery):
 
 
 def _is_valid_node(host: str, port: int, cluster: ClickhouseCluster) -> bool:
-    connection_id = cluster.get_connection_id()
-    return host == connection_id.hostname and port == connection_id.tcp_port
-
-
-def run_system_query_on_host_by_name(
-    clickhouse_host: str,
-    clickhouse_port: int,
-    storage_name: str,
-    system_query_name: str,
-) -> Tuple[Sequence[Any], Sequence[Tuple[str, str]]]:
-    query = SystemQuery.from_name(system_query_name)
+    nodes = cluster.get_local_nodes()
+    return any(node.host_name == host and node.port == port for node in nodes)
 
-    if not query:
-        raise NonExistentSystemQuery(extra_data={"query_name": system_query_name})
 
+def _run_sql_query_on_host(
+    clickhouse_host: str, clickhouse_port: int, storage_name: str, sql: str
+) -> ClickhouseResult:
+    """
+    Run the SQL query. It should be validated before getting to this point
+    """
     storage_key = None
     try:
         storage_key = StorageKey(storage_name)
@@ -130,6 +134,83 @@ def run_system_query_on_host_by_name(
         # force read-only
         client_settings=ClickhouseClientSettings.QUERY.value.settings,
     )
-    query_result = connection.execute(query=query.sql, with_column_types=True)
+    query_result = connection.execute(query=sql, with_column_types=True)
     connection.close()
-    return cast(Tuple[Sequence[Any], Sequence[Tuple[str, str]]], query_result,)
+
+    return query_result
+
+
+def run_system_query_on_host_by_name(
+    clickhouse_host: str,
+    clickhouse_port: int,
+    storage_name: str,
+    system_query_name: str,
+) -> ClickhouseResult:
+    query = SystemQuery.from_name(system_query_name)
+
+    if not query:
+        raise NonExistentSystemQuery(extra_data={"query_name": system_query_name})
+
+    return _run_sql_query_on_host(
+        clickhouse_host, clickhouse_port, storage_name, query.sql
+    )
+
+
+SYSTEM_QUERY_RE = re.compile(
+    r"""
+        ^ # Start
+        (SELECT|select)
+        \s
+        (?P<select_statement>[\w\s,\(\)]+|\*)
+        \s
+        (FROM|from)
+        \s
+        system.[a-z_]+
+        (?P<extra>\s[\w\s,=+\(\)']+)?
+        ;? # Optional semicolon
+        $ # End
+    """,
+    re.VERBOSE,
+)
+
+
+def run_system_query_on_host_with_sql(
+    clickhouse_host: str, clickhouse_port: int, storage_name: str, system_query_sql: str
+) -> ClickhouseResult:
+    validate_system_query(system_query_sql)
+    return _run_sql_query_on_host(
+        clickhouse_host, clickhouse_port, storage_name, system_query_sql
+    )
+
+
+def validate_system_query(sql_query: str) -> None:
+    """
+    Simple validation to ensure query only attempts to access system tables and not
+    any others. Will be replaced by AST parser eventually.
+
+    Raises InvalidCustomQuery if query is invalid or not allowed.
+    """
+    sql_query = " ".join(sql_query.split())
+
+    disallowed_keywords = ["select", "insert", "join"]
+
+    match = SYSTEM_QUERY_RE.match(sql_query)
+
+    if match is None:
+        raise InvalidCustomQuery("Query is invalid")
+
+    select_statement = match.group("select_statement")
+
+    # Extremely quick and dirty way of ensuring there is not a nested select, insert or a join
+    for kw in disallowed_keywords:
+        if kw in select_statement.lower():
+            raise InvalidCustomQuery(f"{kw} is not allowed here")
+
+    extra = match.group("extra")
+
+    # Unfortunately "extra" is pretty permissive right now, just ensure
+    # there is no attempt to do a select, insert or join in there
+    if extra is not None:
+        for kw in disallowed_keywords:
+            if kw in extra.lower():
+                raise InvalidCustomQuery(f"{kw} is not allowed here")

snuba/admin/static/api_client.tsx

@@ -5,14 +5,22 @@ import {
   ConfigChange,
 } from "./runtime_config/types";
 
-import { ClickhouseNodeData } from "./clickhouse_queries/types";
+import {
+  ClickhouseCannedQuery,
+  ClickhouseNodeData,
+  QueryRequest,
+  QueryResult,
+} from "./clickhouse_queries/types";
 
 interface Client {
   getConfigs: () => Promise<Config[]>;
   createNewConfig: (key: ConfigKey, value: ConfigValue) => Promise<Config>;
   deleteConfig: (key: ConfigKey) => Promise<void>;
+  editConfig: (key: ConfigKey, value: ConfigValue) => Promise<Config>;
   getAuditlog: () => Promise<ConfigChange[]>;
   getClickhouseNodes: () => Promise<[ClickhouseNodeData]>;
+  getClickhouseCannedQueries: () => Promise<[ClickhouseCannedQuery]>;
+  executeQuery: (req: QueryRequest) => Promise<QueryResult>;
 }
 
 function Client() {
@@ -55,6 +63,21 @@ function Client() {
         }
       });
     },
+    editConfig: (key: ConfigKey, value: ConfigValue) => {
+      const url = baseUrl + "configs/" + encodeURIComponent(key);
+      return fetch(url, {
+        headers: { "Content-Type": "application/json" },
+        method: "PUT",
+        body: JSON.stringify({ value }),
+      }).then((res) => {
+        if (res.ok) {
+          return Promise.resolve(res.json());
+        } else {
+          throw new Error("Could not edit config");
+        }
+      });
+    },
+
     getAuditlog: () => {
       const url = baseUrl + "config_auditlog";
       return fetch(url).then((resp) => resp.json());
@@ -71,6 +94,18 @@ function Client() {
           })
       );
     },
+    getClickhouseCannedQueries: () => {
+      const url = baseUrl + "clickhouse_queries";
+      return fetch(url).then((resp) => resp.json());
+    },
+    executeQuery: (query: QueryRequest) => {
+      const url = baseUrl + "run_clickhouse_system_query";
+      return fetch(url, {
+        headers: { "Content-Type": "application/json" },
+        method: "POST",
+        body: JSON.stringify(query),
+      }).then((resp) => resp.json());
+    },
   };
 }
 

snuba/admin/static/clickhouse_queries/index.tsx

@@ -1,34 +1,41 @@
 import React, { useEffect, useState } from "react";
 import Client from "../api_client";
+import { Table } from "../table";
 
-import { ClickhouseNodeData } from "./types";
+import {
+  ClickhouseNodeData,
+  ClickhouseCannedQuery,
+  QueryRequest,
+  QueryResult,
+} from "./types";
 
-type QueryState = {
-  storage: string | null;
-  host: string | null;
-  port: number | null;
-};
+type QueryState = Partial<QueryRequest>;
 
 function ClickhouseQueries(props: { api: Client }) {
   const [nodeData, setNodeData] = useState<ClickhouseNodeData[]>([]);
-
-  const [query, setQuery] = useState<QueryState>({
-    storage: null,
-    host: null,
-    port: null,
-  });
+  const [cannedQueries, setCannedQueries] = useState<ClickhouseCannedQuery[]>(
+    []
+  );
+  const [query, setQuery] = useState<QueryState>({});
+  const [queryResultHistory, setQueryResultHistory] = useState<QueryResult[]>(
+    []
+  );
 
   useEffect(() => {
     props.api.getClickhouseNodes().then((res) => {
       setNodeData(res);
     });
+    props.api.getClickhouseCannedQueries().then((res) => {
+      setCannedQueries(res);
+    });
   }, []);
 
   function selectStorage(storage: string) {
-    setQuery({
-      storage,
-      host: null,
-      port: null,
+    setQuery((prevQuery) => {
+      return {
+        ...prevQuery,
+        storage: storage,
+      };
     });
   }
 
@@ -44,9 +51,26 @@ function ClickhouseQueries(props: { api: Client }) {
     });
   }
 
+  function selectCannedQuery(queryName: string) {
+    setQuery((prevQuery) => {
+      return {
+        ...prevQuery,
+        query_name: queryName,
+      };
+    });
+  }
+
+  function executeQuery() {
+    props.api.executeQuery(query as QueryRequest).then((result) => {
+      result.input_query = `${query.query_name}(${query.storage},${query.host}:${query.port})`;
+      setQueryResultHistory((prevHistory) => [result, ...prevHistory]);
+    });
+  }
+
   return (
     <div>
       <form>
+        <h2>Construct a query</h2>
         <select
           value={query.storage || ""}
           onChange={(evt) => selectStorage(evt.target.value)}
@@ -82,7 +106,38 @@ function ClickhouseQueries(props: { api: Client }) {
               ))}
           </select>
         )}
+        {query.storage && query.host && query.port && (
+          <select
+            value={query.query_name || ""}
+            onChange={(evt) => selectCannedQuery(evt.target.value)}
+          >
+            <option disabled value="">
+              Select a query
+            </option>
+            {cannedQueries.map((cannedQuery) => (
+              <option key={`${cannedQuery.name}`} value={`${cannedQuery.name}`}>
+                {cannedQuery.name}: {cannedQuery.description}
+              </option>
+            ))}
+          </select>
+        )}
+        {query.storage && query.host && query.port && query.query_name && (
+          <button onClick={(_) => executeQuery()}>Execute query</button>
+        )}
       </form>
+      <div>
+        <h2>Query results</h2>
+        <Table
+          headerData={["Query", "Response"]}
+          rowData={queryResultHistory.map((queryResult) => [
+            <span>{queryResult.input_query}</span>,
+            <Table
+              headerData={queryResult.column_names}
+              rowData={queryResult.rows}
+            />,
+          ])}
+        />
+      </div>
     </div>
   );
 }

snuba/admin/static/clickhouse_queries/types.tsx

@@ -9,4 +9,27 @@ type ClickhouseNodeData = {
   local_nodes: ClickhouseNode[];
 };
 
-export { ClickhouseNodeData };
+type ClickhouseCannedQuery = {
+  description: string | null;
+  name: string;
+  sql: string;
+};
+
+type QueryRequest = {
+  storage: string;
+  host: string;
+  port: number;
+  query_name: string;
+};
+
+type QueryResultColumnMetadata = [string];
+type QueryResultRow = [string];
+
+type QueryResult = {
+  input_query?: string;
+  timestamp: number;
+  column_names: QueryResultColumnMetadata;
+  rows: [QueryResultRow];
+};
+
+export { ClickhouseNodeData, ClickhouseCannedQuery, QueryRequest, QueryResult };

snuba/admin/static/runtime_config/index.tsx

@@ -76,7 +76,34 @@ function RuntimeConfig(props: { api: Client }) {
                 return { ...prev, value: newValue };
               });
             },
-            () => {}, // TODO: Editing existing row
+            () => {
+              if (
+                window.confirm(
+                  `Are you sure you want to update ${key} to ${currentRowData.value}?`
+                )
+              ) {
+                api
+                  .editConfig(key, currentRowData.value)
+                  .then((res) => {
+                    setData((prev) => {
+                      if (prev) {
+                        const row = prev.find(
+                          (config) => config.key === res.key
+                        );
+                        if (!row) {
+                          throw new Error("An error occurred");
+                        }
+                        row.value = res.value;
+                      }
+                      return prev;
+                    });
+                    resetForm();
+                  })
+                  .catch((err) => {
+                    window.alert(err);
+                  });
+              }
+            },
             () => {
               if (window.confirm(`Are you sure you want to delete ${key}?`)) {
                 api.deleteConfig(key).then(() => {