Add support for filtering in wrangler tail. (#462)

Add filtering to wrangler tail, so you can now `wrangler tail <name> --status ok`, for example. Supported options: - `--status cancelled --status error` --> you can filter on `ok`, `error`, and `cancelled` to only tail logs that have that status - `--header X-CUSTOM-HEADER:somevalue` --> you can filter on headers, including ones that have specific values (`"somevalue"`) or just that contain any header (e.g. `--header X-CUSTOM-HEADER` with no colon) - `--method POST --method PUT` --> filter on the HTTP method used to trigger the worker - `--search catch-this` --> only shows messages that contain the phrase `"catch-this"`. Does not (yet!) support regular expressions - `--ip self --ip 192.0.2.232` --> only show logs from requests that originate from the given IP addresses. `"self"` will be replaced with the IP address of the computer that sent the tail request.
cloudflare · Feb 14, 2022 · a173c80 · a173c80
1 parent 705f988
commit a173c80
Show file tree

Hide file tree

Showing 3 changed files with 228 additions and 21 deletions.
diff --git a/.changeset/gentle-crews-lay.md b/.changeset/gentle-crews-lay.md
@@ -0,0 +1,11 @@
+---
+"wrangler": patch
+---
+
+Add filtering to wrangler tail, so you can now `wrangler tail <name> --status ok`, for example. Supported options:
+
+- `--status cancelled --status error` --> you can filter on `ok`, `error`, and `cancelled` to only tail logs that have that status
+- `--header X-CUSTOM-HEADER:somevalue` --> you can filter on headers, including ones that have specific values (`"somevalue"`) or just that contain any header (e.g. `--header X-CUSTOM-HEADER` with no colon)
+- `--method POST --method PUT` --> filter on the HTTP method used to trigger the worker
+- `--search catch-this` --> only shows messages that contain the phrase `"catch-this"`. Does not (yet!) support regular expressions
+- `--ip self --ip 192.0.2.232` --> only show logs from requests that originate from the given IP addresses. `"self"` will be replaced with the IP address of the computer that sent the tail request.
diff --git a/packages/wrangler/src/index.tsx b/packages/wrangler/src/index.tsx
@@ -30,7 +30,13 @@ import { pages } from "./pages";
 import publish from "./publish";
 import { createR2Bucket, deleteR2Bucket, listR2Buckets } from "./r2";
 import { getAssetPaths } from "./sites";
-import { createTail } from "./tail";
+import {
+  createTail,
+  jsonPrintLogs,
+  prettyPrintLogs,
+  translateCliFiltersToApiFilters,
+} from "./tail";
+import type { TailCLIFilters } from "./tail";
 import {
   login,
   logout,
@@ -45,6 +51,7 @@ import { whoami } from "./whoami";
 import type { Entry } from "./bundle";
 import type { Config } from "./config";
 import type Yargs from "yargs";
+import type { RawData } from "ws";
 
 const resetColor = "\x1b[0m";
 const fgGreenColor = "\x1b[32m";
@@ -1043,6 +1050,7 @@ export async function main(argv: string[]): Promise<void> {
           .option("status", {
             choices: ["ok", "error", "canceled"],
             describe: "Filter by invocation status",
+            array: true,
           })
           .option("header", {
             type: "string",
@@ -1051,6 +1059,7 @@ export async function main(argv: string[]): Promise<void> {
           .option("method", {
             type: "string",
             describe: "Filter by HTTP method",
+            array: true,
           })
           .option("sampling-rate", {
             type: "number",
@@ -1060,12 +1069,18 @@ export async function main(argv: string[]): Promise<void> {
             type: "string",
             describe: "Filter by a text match in console.log messages",
           })
+          .option("ip", {
+            type: "string",
+            describe:
+              'Filter by the IP address the request originates from. Use "self" to filter for your own IP',
+            array: true,
+          })
+          // TODO: is this deprecated now with services / environments / etc?
           .option("env", {
             type: "string",
             describe: "Perform on a specific environment",
           })
       );
-      // TODO: filter by client ip, which can be 'self' or an ip address
     },
     async (args) => {
       if (args.local) {
@@ -1076,12 +1091,11 @@ export async function main(argv: string[]): Promise<void> {
 
       const config = await readConfig(args.config as ConfigPath);
 
-      if (!(args.name || config.name)) {
+      const shortScriptName = args.name || config.name;
+      if (!shortScriptName) {
         throw new Error("Missing script name");
       }
-      const scriptName = `${args.name || config.name}${
-        args.env ? `-${args.env}` : ""
-      }`;
+      const scriptName = `${shortScriptName}${args.env ? `-${args.env}` : ""}`;
 
       // -- snip, extract --
       const loggedIn = await loginOrRefreshIfRequired();
@@ -1101,14 +1115,17 @@ export async function main(argv: string[]): Promise<void> {
 
       const accountId = config.account_id;
 
-      const filters = {
-        status: args.status as "ok" | "error" | "canceled",
+      const cliFilters: TailCLIFilters = {
+        status: args.status as Array<"ok" | "error" | "canceled">,
         header: args.header,
         method: args.method,
-        "sampling-rate": args["sampling-rate"],
+        samplingRate: args["sampling-rate"],
         search: args.search,
+        clientIp: args.ip,
       };
 
+      const filters = translateCliFiltersToApiFilters(cliFilters);
+
       const { tail, expiration, /* sendHeartbeat, */ deleteTail } =
         await createTail(accountId, scriptName, filters);
 
@@ -1121,9 +1138,10 @@ export async function main(argv: string[]): Promise<void> {
         await deleteTail();
       });
 
-      tail.on("message", (data) => {
-        console.log(JSON.stringify(JSON.parse(data.toString()), null, "  "));
-      });
+      const printLog: (data: RawData) => void =
+        args.format === "pretty" ? prettyPrintLogs : jsonPrintLogs;
+
+      tail.on("message", printLog);
 
       while (tail.readyState !== tail.OPEN) {
         switch (tail.readyState) {
@@ -1139,6 +1157,11 @@ export async function main(argv: string[]): Promise<void> {
       }
 
       console.log(`Connected to ${scriptName}, waiting for logs...`);
+
+      tail.on("close", async () => {
+        tail.terminate();
+        await deleteTail();
+      });
     }
   );
 

diff --git a/packages/wrangler/src/tail.tsx b/packages/wrangler/src/tail.tsx
@@ -8,6 +8,61 @@ export type TailApiResponse = {
   expires_at: Date;
 };
 
+export type TailCLIFilters = {
+  status?: Array<"ok" | "error" | "canceled">;
+  header?: string;
+  method?: string[];
+  search?: string;
+  samplingRate?: number;
+  clientIp?: string[];
+};
+
+// due to the trace worker being built around wrangler 1 and
+// some other stuff, the filters we send to the API are slightly
+// different than the ones we read from the CLI
+type SamplingRateFilter = {
+  sampling_rate: number;
+};
+
+type OutcomeFilter = {
+  outcome: string[];
+};
+
+type MethodFilter = {
+  method: string[];
+};
+
+type HeaderFilter = {
+  header: {
+    key: string;
+    query?: string;
+  };
+};
+
+type ClientIpFilter = {
+  client_ip: string[];
+};
+
+type QueryFilter = {
+  query: string;
+};
+
+type ApiFilter =
+  | SamplingRateFilter
+  | OutcomeFilter
+  | MethodFilter
+  | HeaderFilter
+  | ClientIpFilter
+  | QueryFilter;
+
+type ApiFilterMessage = {
+  filters: ApiFilter[];
+  debug: boolean;
+};
+
+// TODO: make this a real type if we wanna pretty-print
+type TailMessage = string;
+
 function makeCreateTailUrl(accountId: string, workerName: string): string {
   return `/accounts/${accountId}/workers/scripts/${workerName}/tails`;
 }
@@ -35,7 +90,7 @@ async function createTailButDontConnect(
 export async function createTail(
   accountId: string,
   workerName: string,
-  _filters: Filters
+  filters: ApiFilter[]
 ): Promise<{
   tail: WebSocket;
   expiration: Date;
@@ -60,14 +115,132 @@ export async function createTail(
     },
   });
 
-  // TODO: send filters as well
+  // check if there's any filters to send
+  if (filters.length === 0) {
+    const message: ApiFilterMessage = {
+      filters,
+      // if debug is set to true, then all logs will be sent through.
+      // logs that _would_ have been blocked will result with a message
+      // telling you what filter would have rejected it
+      debug: false,
+    };
+
+    tail.on("open", function () {
+      tail.send(
+        JSON.stringify(message),
+        { binary: false, compress: false, mask: false, fin: true },
+        (err) => {
+          if (err) {
+            throw err;
+          }
+        }
+      );
+    });
+  }
+
   return { tail, expiration, deleteTail };
 }
 
-export type Filters = {
-  status?: "ok" | "error" | "canceled";
-  header?: string;
-  method?: string;
-  "sampling-rate"?: number;
-  search?: string;
-};
+// TODO: should this validation step happen before connecting to the tail?
+export function translateCliFiltersToApiFilters(
+  cliFilters: TailCLIFilters
+): ApiFilter[] {
+  const apiFilters: ApiFilter[] = [];
+
+  // TODO: do these all need to be their own functions or should
+  // they just be inlined
+  if (cliFilters.samplingRate) {
+    apiFilters.push(parseSamplingRate(cliFilters.samplingRate));
+  }
+
+  if (cliFilters.status) {
+    apiFilters.push(parseOutcome(cliFilters.status));
+  }
+
+  if (cliFilters.method) {
+    apiFilters.push(parseMethod(cliFilters.method));
+  }
+
+  if (cliFilters.header) {
+    apiFilters.push(parseHeader(cliFilters.header));
+  }
+
+  if (cliFilters.clientIp) {
+    apiFilters.push(parseClientIp(cliFilters.clientIp));
+  }
+
+  if (cliFilters.search) {
+    apiFilters.push(parseQuery(cliFilters.search));
+  }
+
+  return apiFilters;
+}
+
+function parseSamplingRate(samplingRate: number): SamplingRateFilter {
+  if (samplingRate <= 0 || samplingRate >= 1) {
+    throw new Error("A sampling rate must be between 0 and 1");
+  }
+
+  return { sampling_rate: samplingRate };
+}
+
+function parseOutcome(
+  statuses: Array<"ok" | "error" | "canceled">
+): OutcomeFilter {
+  const outcomes = new Set<string>();
+  for (const status in statuses) {
+    switch (status) {
+      case "ok":
+        outcomes.add("ok");
+        break;
+      case "canceled":
+        outcomes.add("canceled");
+        break;
+      // there's more than one way to error
+      case "error":
+        outcomes.add("exception");
+        outcomes.add("exceededCpu");
+        outcomes.add("unknown");
+        break;
+      default:
+        break;
+    }
+  }
+
+  return {
+    outcome: Array.from(outcomes),
+  };
+}
+
+// we actually don't need to do anything here
+function parseMethod(method: string[]): MethodFilter {
+  return { method };
+}
+
+function parseHeader(header: string): HeaderFilter {
+  // headers of the form "HEADER-KEY: VALUE" get split.
+  // the query is optional
+  const [headerKey, headerQuery] = header.split(":", 2);
+  return {
+    header: {
+      key: headerKey.trim(),
+      query: headerQuery?.trim(),
+    },
+  };
+}
+
+function parseClientIp(client_ip: string[]): ClientIpFilter {
+  return { client_ip };
+}
+
+function parseQuery(query: string): QueryFilter {
+  return { query };
+}
+
+export function prettyPrintLogs(_data: WebSocket.RawData): void {
+  throw new Error("TODO!");
+}
+
+export function jsonPrintLogs(data: WebSocket.RawData): void {
+  console.log(JSON.stringify(JSON.parse(data.toString()), null, 2));
+}