feat(prod-queries): Implement backend for prod queries (#4398)

snuba/admin/audit_log/querylog.py → snuba/admin/audit_log/query.py

@@ -3,40 +3,38 @@
 from datetime import datetime
 from enum import Enum
 from functools import partial
-from typing import Callable, MutableMapping, Union
-
-from snuba.clickhouse.native import ClickhouseResult
+from typing import Callable, MutableMapping, TypeVar, Union
 
 DATETIME_FORMAT = "%B %d, %Y %H:%M:%S %p"
 
 from snuba.admin.audit_log.action import AuditLogAction
 from snuba.admin.audit_log.base import AuditLog
 
+Return = TypeVar("Return")
+
 
 class QueryExecutionStatus(Enum):
     SUCCEEDED = "succeeded"
     FAILED = "failed"
 
 
-__querylog_audit_log_notification_client = AuditLog()
+__query_audit_log_notification_client = AuditLog()
 
 
-def audit_log(
-    fn: Callable[[str, str], ClickhouseResult]
-) -> Callable[[str, str], ClickhouseResult]:
+def audit_log(fn: Callable[[str, str], Return]) -> Callable[[str, str], Return]:
     """
     Decorator function for querylog query runner.
 
     Logs the user, query, start/end timestamps, and whether or not
     the query was successful.
     """
 
-    def audit_log_wrapper(query: str, user: str) -> ClickhouseResult:
+    def audit_log_wrapper(query: str, user: str) -> Return:
         data: MutableMapping[str, Union[str, QueryExecutionStatus]] = {
             "query": query,
         }
         audit_log_notify = partial(
-            __querylog_audit_log_notification_client.record,
+            __query_audit_log_notification_client.record,
             user=user,
             action=AuditLogAction.RAN_QUERY,
         )

snuba/admin/cardinality_analyzer/cardinality_analyzer.py

@@ -1,4 +1,4 @@
-from snuba.admin.audit_log.querylog import audit_log
+from snuba.admin.audit_log.query import audit_log
 from snuba.admin.clickhouse.common import (
     get_ro_query_node_connection,
     validate_ro_query,

snuba/admin/clickhouse/querylog.py

@@ -1,5 +1,5 @@
 from snuba import state
-from snuba.admin.audit_log.querylog import audit_log
+from snuba.admin.audit_log.query import audit_log
 from snuba.admin.clickhouse.common import (
     get_ro_query_node_connection,
     validate_ro_query,

snuba/admin/production_queries/prod_queries.py

@@ -0,0 +1,51 @@
+from typing import Any, Dict
+
+from flask import Response
+
+from snuba import settings
+from snuba.admin.audit_log.query import audit_log
+from snuba.clickhouse.query_dsl.accessors import get_object_ids_in_query_ast
+from snuba.datasets.dataset import Dataset
+from snuba.datasets.factory import get_dataset
+from snuba.query.exceptions import InvalidQueryException
+from snuba.query.query_settings import HTTPQuerySettings
+from snuba.query.snql.parser import parse_snql_query
+from snuba.request.schema import RequestSchema
+from snuba.utils.metrics.timer import Timer
+from snuba.web.views import dataset_query
+
+
+def run_snql_query(body: Dict[str, Any], user: str) -> Response:
+    """
+    Validates, audit logs, and executes given query.
+    """
+
+    @audit_log
+    def run_query_with_audit(query: str, user: str) -> Response:
+        dataset = get_dataset(body.pop("dataset"))
+        body["dry_run"] = True
+        response = dataset_query(dataset, body, Timer("admin"))
+        if response.status_code != 200:
+            return response
+
+        body["dry_run"] = False
+        _validate_projects_in_query(body, dataset)
+        return dataset_query(dataset, body, Timer("admin"))
+
+    return run_query_with_audit(body["query"], user)
+
+
+def _validate_projects_in_query(body: Dict[str, Any], dataset: Dataset) -> None:
+    request_parts = RequestSchema.build(HTTPQuerySettings).validate(body)
+    query = parse_snql_query(request_parts.query["query"], dataset)[0]
+    project_ids = get_object_ids_in_query_ast(query, "project_id")
+    if project_ids is None:
+        raise InvalidQueryException("Missing project ID")
+
+    disallowed_project_ids = project_ids.difference(
+        set(settings.ADMIN_ALLOWED_PROD_PROJECTS)
+    )
+    if len(disallowed_project_ids) > 0:
+        raise InvalidQueryException(
+            f"Cannot access the following project ids: {disallowed_project_ids}"
+        )

snuba/admin/static/api_client.tsx

@@ -53,6 +53,7 @@ interface Client {
   getClickhouseNodes: () => Promise<[ClickhouseNodeData]>;
   getSnubaDatasetNames: () => Promise<SnubaDatasetName[]>;
   convertSnQLQuery: (query: SnQLRequest) => Promise<SnQLResult>;
+  executeSnQLQuery: (query: SnQLRequest) => Promise<any>;
   getPredefinedQueryOptions: () => Promise<[PredefinedQuery]>;
   executeSystemQuery: (req: QueryRequest) => Promise<QueryResult>;
   executeTracingQuery: (req: TracingRequest) => Promise<TracingResult>;
@@ -205,6 +206,24 @@ function Client() {
       });
     },
 
+    executeSnQLQuery: (query: SnQLRequest) => {
+      const url = baseUrl + "production_snql_query";
+      return fetch(url, {
+        headers: { "Content-Type": "application/json" },
+        method: "POST",
+        body: JSON.stringify(query),
+      }).then((res) => {
+        if (res.ok) {
+          return Promise.resolve(res.json());
+        } else {
+          return res.json().then((err) => {
+            let errMsg = err?.error.message || "Could not execute SnQL";
+            throw new Error(errMsg);
+          });
+        }
+      });
+    },
+
     getPredefinedQueryOptions: () => {
       const url = baseUrl + "clickhouse_queries";
       return fetch(url).then((resp) => resp.json());

snuba/admin/static/production_queries/index.tsx

@@ -1,13 +1,134 @@
-import React, { useState, useEffect } from "react";
+import React, { useEffect, useState } from "react";
 import Client from "../api_client";
-import QueryDisplay from "./query_display";
+import { Table } from "../table";
+import { QueryResult, QueryResultColumnMeta, SnQLRequest } from "./types";
+import { executeActionsStyle, executeButtonStyle, selectStyle } from "./styles";
 
 function ProductionQueries(props: { api: Client }) {
+  const [datasets, setDatasets] = useState<string[]>([]);
+  const [snql_query, setQuery] = useState<Partial<SnQLRequest>>({});
+  const [queryResultHistory, setQueryResultHistory] = useState<QueryResult[]>(
+    []
+  );
+  const [isExecuting, setIsExecuting] = useState<boolean>(false);
+
+  useEffect(() => {
+    props.api.getSnubaDatasetNames().then((res) => {
+      setDatasets(res);
+    });
+  }, []);
+
+  function selectDataset(dataset: string) {
+    setQuery((prevQuery) => {
+      return {
+        ...prevQuery,
+        dataset,
+      };
+    });
+  }
+
+  function updateQuerySql(query: string) {
+    setQuery((prevQuery) => {
+      return {
+        ...prevQuery,
+        query,
+      };
+    });
+  }
+
+  function executeQuery() {
+    if (isExecuting) {
+      window.alert("A query is already running");
+    }
+    setIsExecuting(true);
+    props.api
+      .executeSnQLQuery(snql_query as SnQLRequest)
+      .then((result) => {
+        const result_columns = result.meta.map(
+          (col: QueryResultColumnMeta) => col.name
+        );
+        const query_result: QueryResult = {
+          input_query: snql_query.query,
+          columns: result_columns,
+          rows: result.data.map((obj: object) =>
+            result_columns.map(
+              (col_name: string) => obj[col_name as keyof typeof obj]
+            )
+          ),
+        };
+        setQueryResultHistory((prevHistory) => [query_result, ...prevHistory]);
+      })
+      .catch((err) => {
+        console.log("ERROR", err);
+        window.alert("An error occurred: " + err.message);
+      })
+      .finally(() => {
+        setIsExecuting(false);
+      });
+  }
+
   return (
     <div>
-      {QueryDisplay({
-        api: props.api,
-      })}
+      <form>
+        <h2>Run a SnQL Query</h2>
+        <p>Currently, we only support queries with project_id = 1</p>
+        <div>
+          <textarea
+            spellCheck={false}
+            value={snql_query.query || ""}
+            onChange={(evt) => updateQuerySql(evt.target.value)}
+            style={{ width: "100%", height: 100 }}
+            placeholder={"Write your query here"}
+          />
+        </div>
+        <div style={executeActionsStyle}>
+          <div>
+            <select
+              value={snql_query.dataset || ""}
+              onChange={(evt) => selectDataset(evt.target.value)}
+              style={selectStyle}
+            >
+              <option disabled value="">
+                Select a dataset
+              </option>
+              {datasets.map((dataset) => (
+                <option key={dataset} value={dataset}>
+                  {dataset}
+                </option>
+              ))}
+            </select>
+          </div>
+          <div>
+            <button
+              onClick={(_) => executeQuery()}
+              style={executeButtonStyle}
+              disabled={
+                isExecuting ||
+                snql_query.dataset == undefined ||
+                snql_query.query == undefined
+              }
+            >
+              Execute Query
+            </button>
+          </div>
+        </div>
+      </form>
+      <div>
+        <h2>Query results</h2>
+        {queryResultHistory.map((queryResult, idx) => {
+          if (idx === 0) {
+            console.log(queryResult);
+            return (
+              <div>
+                <Table
+                  headerData={queryResult.columns}
+                  rowData={queryResult.rows}
+                />
+              </div>
+            );
+          }
+        })}
+      </div>
     </div>
   );
 }

snuba/admin/static/production_queries/styles.tsx

@@ -0,0 +1,25 @@
+const executeActionsStyle = {
+  display: "flex",
+  justifyContent: "space-between",
+  marginTop: 8,
+};
+
+const executeButtonStyle = {
+  height: 30,
+  border: 0,
+  padding: "4px 20px",
+};
+
+const selectStyle = {
+  marginRight: 8,
+  height: 30,
+};
+
+let collapsibleStyle = { listStyleType: "none", fontFamily: "Monaco" };
+
+export {
+  executeActionsStyle,
+  executeButtonStyle,
+  selectStyle,
+  collapsibleStyle,
+};

snuba/admin/static/production_queries/types.tsx

@@ -0,0 +1,17 @@
+type SnQLRequest = {
+  dataset: string;
+  query: string;
+};
+
+type QueryResult = {
+  input_query?: string;
+  columns: [string];
+  rows: [[string]];
+};
+
+type QueryResultColumnMeta = {
+  name: string;
+  type: string;
+};
+
+export { SnQLRequest, QueryResult, QueryResultColumnMeta };

snuba/admin/tool_policies.py

@@ -27,6 +27,7 @@ class AdminTools(Enum):
     AUDIT_LOG = "audit-log"
     KAFKA = "kafka"
     CAPACITY_MANAGEMENT = "capacity-management"
+    PRODUCTION_QUERIES = "production-queries"
     CARDINALITY_ANALYZER = "cardinality-analyzer"
     SNUBA_EXPLAIN = "snuba_explain"