Print included query info fields as yaml (#526)

* Print included query info fields as yaml * disable line folding in yaml * fix color change * enable query info in shell * implement query info toggling in shell * help suggestions * include summary by default * testing * formatting --------- Co-authored-by: E. Cooper <[email protected]>
fauna · Dec 16, 2024 · b381a73 · b381a73
1 parent ae7b43f
commit b381a73
Show file tree

Hide file tree

Showing 7 changed files with 204 additions and 35 deletions.
diff --git a/src/commands/query.mjs b/src/commands/query.mjs
@@ -13,8 +13,8 @@ import {
 } from "../lib/errors.mjs";
 import {
   formatError,
+  formatQueryInfo,
   formatQueryResponse,
-  formatQuerySummary,
   getSecret,
 } from "../lib/fauna-client.mjs";
 import { isTTY } from "../lib/misc.mjs";
@@ -89,8 +89,8 @@ async function queryCommand(argv) {
       typecheck,
       apiVersion,
       performanceHints,
-      summary,
       color,
+      include,
     } = argv;
 
     // If we're writing to a file, don't colorize the output regardless of the user's preference
@@ -110,12 +110,16 @@ async function queryCommand(argv) {
       color: useColor,
     });
 
-    // If performance hints are enabled, print the summary to stderr.
+    // If any query info should be displayed, print to stderr.
     // This is only supported in v10.
-    if ((summary || performanceHints) && apiVersion === "10") {
-      const formattedSummary = formatQuerySummary(results.summary);
-      if (formattedSummary) {
-        logger.stderr(formattedSummary);
+    if (include.length > 0 && apiVersion === "10") {
+      const queryInfo = formatQueryInfo(results, {
+        apiVersion,
+        color: useColor,
+        include,
+      });
+      if (queryInfo) {
+        logger.stderr(queryInfo);
       }
     }
 

diff --git a/src/commands/shell.mjs b/src/commands/shell.mjs
@@ -7,13 +7,14 @@ import * as esprima from "esprima";
 
 import { container } from "../cli.mjs";
 import {
+  QUERY_INFO_CHOICES,
   resolveFormat,
   validateDatabaseOrSecret,
   yargsWithCommonConfigurableQueryOptions,
 } from "../lib/command-helpers.mjs";
 import {
+  formatQueryInfo,
   formatQueryResponse,
-  formatQuerySummary,
   getSecret,
 } from "../lib/fauna-client.mjs";
 import { clearHistoryStorage, initHistoryStorage } from "../lib/file-util.mjs";
@@ -66,6 +67,8 @@ async function shellCommand(argv) {
     shell.on("exit", resolve);
   });
 
+  shell.context.include = argv.include;
+
   [
     {
       cmd: "clear",
@@ -114,13 +117,21 @@ async function shellCommand(argv) {
       },
     },
     {
-      cmd: "toggleSummary",
-      help: "Enable or disable the summary field of the API response. Disabled by default. If enabled, outputs the summary field of the API response.",
+      cmd: "toggleInfo",
+      help: "Enable or disable the query info fields of the API response. Disabled by default. If enabled, outputs the included fields of the API response.",
       action: () => {
-        shell.context.summary = !shell.context.summary;
+        shell.context.include =
+          shell.context.include.length === 0
+            ? // if we are toggling on and no include was provided, turn everything on
+              argv.include.length === 0
+              ? QUERY_INFO_CHOICES
+              : argv.include
+            : [];
+
         logger.stderr(
-          `Summary in shell: ${shell.context.summary ? "on" : "off"}`,
+          `Query info in shell: ${shell.context.include.length === 0 ? "off" : shell.context.include.join(", ")}`,
         );
+
         shell.prompt();
       },
     },
@@ -150,8 +161,8 @@ async function buildCustomEval(argv) {
 
       // These are options used for querying and formatting the response
       const { apiVersion, color } = argv;
+      const include = getArgvOrCtx("include", argv, ctx);
       const performanceHints = getArgvOrCtx("performanceHints", argv, ctx);
-      const summary = getArgvOrCtx("summary", argv, ctx);
 
       // Using --json output takes precedence over --format
       const outputFormat = resolveFormat({ ...argv });
@@ -167,7 +178,7 @@ async function buildCustomEval(argv) {
       let res;
       try {
         const secret = await getSecret();
-        const { url, timeout, typecheck } = argv;
+        const { color, timeout, typecheck, url } = argv;
 
         res = await runQueryFromString(cmd, {
           apiVersion,
@@ -179,10 +190,16 @@ async function buildCustomEval(argv) {
           format: outputFormat,
         });
 
-        if ((summary || performanceHints) && apiVersion === "10") {
-          const formattedSummary = formatQuerySummary(res.summary);
-          if (formattedSummary) {
-            logger.stdout(formattedSummary);
+        // If any query info should be displayed, print to stderr.
+        // This is only supported in v10.
+        if (include.length > 0 && apiVersion === "10") {
+          const queryInfo = formatQueryInfo(res, {
+            apiVersion,
+            color,
+            include,
+          });
+          if (queryInfo) {
+            logger.stdout(queryInfo);
           }
         }
       } catch (err) {

diff --git a/src/lib/command-helpers.mjs b/src/lib/command-helpers.mjs
@@ -79,6 +79,14 @@ const COMMON_QUERY_OPTIONS = {
   },
 };
 
+export const QUERY_INFO_CHOICES = [
+  "txnTs",
+  "schemaVersion",
+  "summary",
+  "queryTags",
+  "stats",
+];
+
 // used for queries customers can configure
 const COMMON_CONFIGURABLE_QUERY_OPTIONS = {
   ...COMMON_QUERY_OPTIONS,
@@ -114,28 +122,47 @@ const COMMON_CONFIGURABLE_QUERY_OPTIONS = {
     default: 5000,
     group: "API:",
   },
-  summary: {
-    type: "boolean",
-    description:
-      "Output the summary field of the API response or nothing when it's empty. Only applies to v10 queries.",
-    default: false,
-    group: "API:",
-  },
   performanceHints: {
     type: "boolean",
     description:
-      "Output the performance hints for the current query or nothing when no hints are available. Only applies to v10 queries.",
+      "Output the performance hints for the current query or nothing when no hints are available. Only applies to v10 queries. Sets '--include summary'",
     default: false,
     group: "API:",
   },
+  include: {
+    type: "array",
+    choices: ["all", "none", ...QUERY_INFO_CHOICES],
+    default: ["summary"],
+    describe: "Include additional query response data in the output.",
+  },
 };
 
 export function yargsWithCommonQueryOptions(yargs) {
   return yargsWithCommonOptions(yargs, COMMON_QUERY_OPTIONS);
 }
 
 export function yargsWithCommonConfigurableQueryOptions(yargs) {
-  return yargsWithCommonOptions(yargs, COMMON_CONFIGURABLE_QUERY_OPTIONS);
+  return yargsWithCommonOptions(
+    yargs,
+    COMMON_CONFIGURABLE_QUERY_OPTIONS,
+  ).middleware((argv) => {
+    if (argv.include.includes("none")) {
+      if (argv.include.length !== 1) {
+        throw new ValidationError(
+          `'--include none' cannot be used with other include options. Provided options: '${argv.include.join(", ")}'`,
+        );
+      }
+      argv.include = [];
+    }
+
+    if (argv.include.includes("all")) {
+      argv.include = [...QUERY_INFO_CHOICES];
+    }
+
+    if (argv.performanceHints && !argv.include.includes("summary")) {
+      argv.include.push("summary");
+    }
+  });
 }
 
 export function yargsWithCommonOptions(yargs, options) {

diff --git a/src/lib/fauna-client.mjs b/src/lib/fauna-client.mjs
@@ -179,3 +179,70 @@ export const formatQuerySummary = (summary) => {
     return summary;
   }
 };
+
+/**
+ * Selects a subset of query info fields from a v10 query response.
+ * @param {import("fauna").QueryInfo} response - The query response
+ * @param {string[]} include - The query info fields to include
+ * @returns {object} An object with the selected query info fields
+ */
+const pickAndCamelCaseQueryInfo = (response, include) => {
+  const queryInfo = {};
+
+  if (include.includes("txnTs") && response.txn_ts)
+    queryInfo.txnTs = response.txn_ts;
+  if (include.includes("schemaVersion") && response.schema_version)
+    queryInfo.schemaVersion = response.schema_version.toString();
+  if (include.includes("summary") && response.summary)
+    queryInfo.summary = response.summary;
+  if (include.includes("queryTags") && response.query_tags)
+    queryInfo.queryTags = response.query_tags;
+  if (include.includes("stats") && response.stats)
+    queryInfo.stats = response.stats;
+
+  return queryInfo;
+};
+
+/**
+ *
+ * @param {object} response - The v4 or v10 query response with query info
+ * @param {object} opts
+ * @param {string} opts.apiVersion - The API version
+ * @param {boolean} opts.color - Whether to colorize the error
+ * @param {string[]} opts.include - The query info fields to include
+ * @returns
+ */
+export const formatQueryInfo = (response, { apiVersion, color, include }) => {
+  if (apiVersion === "4" && include.includes("stats")) {
+    /** @type {import("faunadb").MetricsResponse} */
+    const metricsResponse = response;
+    const colorized = colorize(
+      { metrics: metricsResponse.metrics },
+      { color, format: Format.YAML },
+    );
+
+    return `${colorized}\n`;
+  } else if (apiVersion === "10") {
+    const queryInfoToDisplay = pickAndCamelCaseQueryInfo(response, include);
+
+    if (Object.keys(queryInfoToDisplay).length === 0) return "";
+
+    const SUMMARY_IN_QUERY_INFO_FQL_REGEX = /^(\s\s\s\s\|)|(\d\s\|)/;
+    const colorized = colorize(queryInfoToDisplay, {
+      color,
+      format: Format.YAML,
+    })
+      .split("\n")
+      .map((line) => {
+        if (!line.match(SUMMARY_IN_QUERY_INFO_FQL_REGEX)) {
+          return line;
+        }
+        return colorize(line, { format: Format.FQL });
+      })
+      .join("\n");
+
+    return `${colorized}\n`;
+  }
+
+  return "";
+};
diff --git a/src/lib/formatting/codeToAnsi.mjs b/src/lib/formatting/codeToAnsi.mjs
@@ -3,6 +3,7 @@ import { createHighlighterCoreSync } from "shiki/core";
 import { createJavaScriptRegexEngine } from "shiki/engine/javascript";
 import json from "shiki/langs/json.mjs";
 import log from "shiki/langs/log.mjs";
+import yaml from "shiki/langs/yaml.mjs";
 import githubDarkHighContrast from "shiki/themes/github-dark-high-contrast.mjs";
 
 import { isTTY } from "../misc.mjs";
@@ -13,7 +14,7 @@ const THEME = "github-dark-high-contrast";
 export const createHighlighter = () => {
   const highlighter = createHighlighterCoreSync({
     themes: [githubDarkHighContrast],
-    langs: [fql, log, json],
+    langs: [fql, log, json, yaml],
     engine: createJavaScriptRegexEngine(),
   });
 

diff --git a/src/lib/formatting/colorize.mjs b/src/lib/formatting/colorize.mjs
@@ -1,4 +1,5 @@
 import stripAnsi from "strip-ansi";
+import YAML from "yaml";
 
 import { container } from "../../cli.mjs";
 import { codeToAnsi } from "./codeToAnsi.mjs";
@@ -8,6 +9,7 @@ export const Format = {
   LOG: "log",
   JSON: "json",
   TEXT: "text",
+  YAML: "yaml",
 };
 
 const objToString = (obj) => JSON.stringify(obj, null, 2);
@@ -52,6 +54,18 @@ const logToAnsi = (obj) => {
   return res.trim();
 };
 
+const yamlToAnsi = (obj) => {
+  const codeToAnsi = container.resolve("codeToAnsi");
+  const stringified = YAML.stringify(obj, { lineWidth: 0 });
+  const res = codeToAnsi(stringified, "yaml");
+
+  if (!res) {
+    return "";
+  }
+
+  return res.trim();
+};
+
 /**
  * Formats an object for display with ANSI color codes.
  * @param {any} obj - The object to format
@@ -67,6 +81,8 @@ export const toAnsi = (obj, { format = Format.TEXT } = {}) => {
       return jsonToAnsi(obj);
     case Format.LOG:
       return logToAnsi(obj);
+    case Format.YAML:
+      return yamlToAnsi(obj);
     default:
       return textToAnsi(obj);
   }