[D1] teach wrangler how to fetch insights about D1's queries (#4909)

* [D1] teach wrangler how to fetch insights about D1's queries * add options for sorting and number of results * add a means of filtering the datetime geq/leq * add more options, fix tests * add wrangler banner, add warning text * add changeset * rename last to timePeriod * use string format over numbers as >1mo's data could be possible * Update red-icons-flow.md * make it possible to order by count * Update red-icons-flow.md * PR feedback * Update .changeset/red-icons-flow.md Co-authored-by: Pete Bacon Darwin <[email protected]> * address PR feedback --------- Co-authored-by: Pete Bacon Darwin <[email protected]>
cloudflare · Feb 6, 2024 · 34b6ea1 · 34b6ea1
1 parent e61dba5
commit 34b6ea1
Show file tree

Hide file tree

Showing 6 changed files with 239 additions and 1 deletion.
diff --git a/.changeset/red-icons-flow.md b/.changeset/red-icons-flow.md
@@ -0,0 +1,27 @@
+---
+"wrangler": patch
+---
+
+feat: add an experimental `insights` command to `wrangler d1`
+
+This PR adds a `wrangler d1 insights <DB_NAME>` command, to let D1 users figure out which of their queries to D1 need to be optimised.
+
+This command defaults to fetching the top 5 queries that took the longest to run in total over the last 24 hours.
+
+You can also fetch the top 5 queries that consumed the most rows read over the last week, for example:
+
+```bash
+npx wrangler d1 insights northwind --sortBy reads --timePeriod 7d
+```
+
+Or the top 5 queries that consumed the most rows written over the last month, for example:
+
+```bash
+npx wrangler d1 insights northwind --sortBy writes --timePeriod 31d
+```
+
+Or the top 5 most frequently run queries in the last 24 hours, for example:
+
+```bash
+npx wrangler d1 insights northwind --sortBy count
+```
diff --git a/packages/wrangler/src/__tests__/d1/d1.test.ts b/packages/wrangler/src/__tests__/d1/d1.test.ts
@@ -19,6 +19,7 @@ describe("d1", () => {
 		Commands:
 		  wrangler d1 list                List D1 databases
 		  wrangler d1 info <name>         Get information about a D1 database, including the current database size and state.
+		  wrangler d1 insights <name>     Experimental command. Get information about the queries run on a D1 database.
 		  wrangler d1 create <name>       Create D1 database
 		  wrangler d1 delete <name>       Delete D1 database
 		  wrangler d1 backup              Interact with D1 Backups
@@ -59,6 +60,7 @@ describe("d1", () => {
 		Commands:
 		  wrangler d1 list                List D1 databases
 		  wrangler d1 info <name>         Get information about a D1 database, including the current database size and state.
+		  wrangler d1 insights <name>     Experimental command. Get information about the queries run on a D1 database.
 		  wrangler d1 create <name>       Create D1 database
 		  wrangler d1 delete <name>       Delete D1 database
 		  wrangler d1 backup              Interact with D1 Backups

diff --git a/packages/wrangler/src/d1/index.ts b/packages/wrangler/src/d1/index.ts
@@ -3,6 +3,7 @@ import * as Create from "./create";
 import * as Delete from "./delete";
 import * as Execute from "./execute";
 import * as Info from "./info";
+import * as Insights from "./insights";
 import * as List from "./list";
 import * as Migrations from "./migrations";
 import * as TimeTravel from "./timeTravel";
@@ -19,6 +20,12 @@ export function d1(yargs: CommonYargsArgv) {
 				Info.Options,
 				Info.Handler
 			)
+			.command(
+				"insights <name>",
+				"Experimental command. Get information about the queries run on a D1 database.",
+				Insights.Options,
+				Insights.Handler
+			)
 			.command(
 				"create <name>",
 				"Create D1 database",

diff --git a/packages/wrangler/src/d1/info.tsx b/packages/wrangler/src/d1/info.tsx
@@ -49,7 +49,7 @@ export const Handler = withConfig<HandlerOptions>(
 			output["database_size"] = output["file_size"];
 			delete output["file_size"];
 		}
-		if (result.version === "beta") {
+		if (result.version !== "alpha") {
 			const today = new Date();
 			const yesterday = new Date(new Date(today).setDate(today.getDate() - 1));
 

diff --git a/packages/wrangler/src/d1/insights.ts b/packages/wrangler/src/d1/insights.ts
@@ -0,0 +1,170 @@
+import { printWranglerBanner } from "..";
+import { fetchGraphqlResult } from "../cfetch";
+import { withConfig } from "../config";
+import { logger } from "../logger";
+import { requireAuth } from "../user";
+import {
+	d1BetaWarning,
+	getDatabaseByNameOrBinding,
+	getDatabaseInfoFromId,
+} from "./utils";
+import type {
+	CommonYargsArgv,
+	StrictYargsOptionsToInterface,
+} from "../yargs-types";
+import type { D1QueriesGraphQLResponse, Database } from "./types";
+
+export function Options(d1ListYargs: CommonYargsArgv) {
+	return d1ListYargs
+		.positional("name", {
+			describe: "The name of the DB",
+			type: "string",
+			demandOption: true,
+		})
+		.option("timePeriod", {
+			choices: ["1d", "7d", "31d"] as const,
+			describe: "Fetch data from now to the provided time period",
+			default: "1d" as const,
+		})
+		.option("sort-type", {
+			choices: ["sum", "avg"] as const,
+			describe: "Choose the operation you want to sort insights by",
+			default: "sum" as const,
+		})
+		.option("sort-by", {
+			choices: ["time", "reads", "writes", "count"] as const,
+			describe: "Choose the field you want to sort insights by",
+			default: "time" as const,
+		})
+		.option("sort-direction", {
+			choices: ["ASC", "DESC"] as const,
+			describe: "Choose a sort direction",
+			default: "DESC" as const,
+		})
+		.option("count", {
+			describe: "fetch insights about the first X queries",
+			type: "number",
+			default: 5,
+		})
+		.option("json", {
+			describe: "return output as clean JSON",
+			type: "boolean",
+			default: false,
+		})
+		.epilogue(d1BetaWarning);
+}
+
+const cliOptionToGraphQLOption = {
+	time: "queryDurationMs",
+	reads: "rowsRead",
+	writes: "rowsWritten",
+	count: "count",
+};
+
+type HandlerOptions = StrictYargsOptionsToInterface<typeof Options>;
+export const Handler = withConfig<HandlerOptions>(
+	async ({
+		name,
+		config,
+		json,
+		count,
+		timePeriod,
+		sortType,
+		sortBy,
+		sortDirection,
+	}): Promise<void> => {
+		const accountId = await requireAuth(config);
+		const db: Database = await getDatabaseByNameOrBinding(
+			config,
+			accountId,
+			name
+		);
+
+		const result = await getDatabaseInfoFromId(accountId, db.uuid);
+
+		const output: Record<string, string | number>[] = [];
+
+		if (result.version !== "alpha") {
+			const convertedTimePeriod = Number(timePeriod.replace("d", ""));
+			const endDate = new Date();
+			const startDate = new Date(
+				new Date(endDate).setDate(endDate.getDate() - convertedTimePeriod)
+			);
+			const parsedSortBy = cliOptionToGraphQLOption[sortBy];
+			const orderByClause =
+				parsedSortBy === "count"
+					? `${parsedSortBy}_${sortDirection}`
+					: `${sortType}_${parsedSortBy}_${sortDirection}`;
+			const graphqlQueriesResult =
+				await fetchGraphqlResult<D1QueriesGraphQLResponse>({
+					method: "POST",
+					body: JSON.stringify({
+						query: `query getD1QueriesOverviewQuery($accountTag: string, $filter: ZoneWorkersRequestsFilter_InputObject) {
+								viewer {
+									accounts(filter: {accountTag: $accountTag}) {
+										d1QueriesAdaptiveGroups(limit: ${count}, filter: $filter, orderBy: [${orderByClause}]) {
+											sum {
+												queryDurationMs
+												rowsRead
+												rowsWritten
+											}
+											avg {
+												queryDurationMs
+												rowsRead
+												rowsWritten
+											}
+											count
+											dimensions {
+													query
+											}
+										}
+									}
+								}
+							}`,
+						operationName: "getD1QueriesOverviewQuery",
+						variables: {
+							accountTag: accountId,
+							filter: {
+								AND: [
+									{
+										datetimeHour_geq: startDate.toISOString(),
+										datetimeHour_leq: endDate.toISOString(),
+										databaseId: db.uuid,
+									},
+								],
+							},
+						},
+					}),
+					headers: {
+						"Content-Type": "application/json",
+					},
+				});
+
+			graphqlQueriesResult?.data?.viewer?.accounts[0]?.d1QueriesAdaptiveGroups?.forEach(
+				(row) => {
+					if (!row.dimensions.query) return;
+					output.push({
+						query: row.dimensions.query,
+						avgRowsRead: row?.avg?.rowsRead ?? 0,
+						totalRowsRead: row?.sum?.rowsRead ?? 0,
+						avgRowsWritten: row?.avg?.rowsWritten ?? 0,
+						totalRowsWritten: row?.sum?.rowsWritten ?? 0,
+						avgDurationMs: row?.avg?.queryDurationMs ?? 0,
+						totalDurationMs: row?.sum?.queryDurationMs ?? 0,
+						numberOfTimesRun: row?.count ?? 0,
+					});
+				}
+			);
+		}
+
+		if (json) {
+			logger.log(JSON.stringify(output, null, 2));
+		} else {
+			await printWranglerBanner();
+			logger.log(
+				"-------------------\n🚧 `wrangler d1 insights` is an experimental command.\n🚧 Flags for this command, their descriptions, and output may change between wrangler versions.\n-------------------\n"
+			);
+			logger.log(JSON.stringify(output, null, 2));
+		}
+	}
+);
diff --git a/packages/wrangler/src/d1/types.ts b/packages/wrangler/src/d1/types.ts
@@ -72,3 +72,35 @@ export interface D1MetricsGraphQLResponse {
 		};
 	};
 }
+
+export interface D1Queries {
+	avg?: {
+		queryDurationMs?: number;
+		rowsRead?: number;
+		rowsWritten?: number;
+	};
+	sum?: {
+		queryDurationMs?: number;
+		rowsRead?: number;
+		rowsWritten?: number;
+	};
+	count?: number;
+	dimensions: {
+		query?: string;
+		databaseId?: string;
+		date?: string;
+		datetime?: string;
+		datetimeMinute?: string;
+		datetimeFiveMinutes?: string;
+		datetimeFifteenMinutes?: string;
+		datetimeHour?: string;
+	};
+}
+
+export interface D1QueriesGraphQLResponse {
+	data: {
+		viewer: {
+			accounts: { d1QueriesAdaptiveGroups?: D1Queries[] }[];
+		};
+	};
+}