fix: add 24h read/write metrics to d1 info (#3434)

* fix: add 24h read/write metrics to d1 info

* add test

* add test

* filter by db.uuid

* fix: we can only fetch metrics for new dbs

* fix: use toLocaleString to get human-readable numbers for now

* fix: don't prettify json output

* file_size -> database_size

* remove date-fns, fix up some logic

.changeset/six-horses-punch.md

@@ -0,0 +1,5 @@
+---
+"wrangler": patch
+---
+
+fix: add the number of read queries and write queries in the last 24 hours to the `d1 info` command

packages/wrangler/src/__tests__/fetch-graphql-result.test.ts

@@ -0,0 +1,77 @@
+import { rest } from "msw";
+import { fetchGraphqlResult } from "../cfetch";
+import { mockAccountId, mockApiToken } from "./helpers/mock-account-id";
+import { mockOAuthFlow } from "./helpers/mock-oauth-flow";
+import { msw } from "./helpers/msw";
+
+describe("fetchGraphqlResult", () => {
+	mockAccountId({ accountId: null });
+	mockApiToken();
+	const { mockOAuthServerCallback } = mockOAuthFlow();
+
+	it("should make a request against the graphql endpoint by default", async () => {
+		mockOAuthServerCallback();
+		msw.use(
+			rest.post("*/graphql", async (req, res, ctx) => {
+				return res(
+					ctx.status(200),
+					ctx.json({
+						data: {
+							viewer: {
+								__typename: "viewer",
+							},
+						},
+						errors: null,
+					})
+				);
+			})
+		);
+		expect(
+			await fetchGraphqlResult({
+				body: JSON.stringify({
+					query: `{
+                    viewer {
+                      __typename
+                    }
+                  }`,
+				}),
+			})
+		).toEqual({ data: { viewer: { __typename: "viewer" } }, errors: null });
+	});
+
+	it("should accept a request with no init, but return no data", async () => {
+		mockOAuthServerCallback();
+		const now = new Date().toISOString();
+		msw.use(
+			rest.post("*/graphql", async (req, res, ctx) => {
+				return res(
+					ctx.status(200),
+					ctx.json({
+						data: null,
+						errors: [
+							{
+								message: "failed to recognize JSON request: 'EOF'",
+								path: null,
+								extensions: {
+									timestamp: now,
+								},
+							},
+						],
+					})
+				);
+			})
+		);
+		expect(await fetchGraphqlResult()).toEqual({
+			data: null,
+			errors: [
+				{
+					message: "failed to recognize JSON request: 'EOF'",
+					path: null,
+					extensions: {
+						timestamp: now,
+					},
+				},
+			],
+		});
+	});
+});

packages/wrangler/src/cfetch/index.ts

@@ -43,6 +43,26 @@ export async function fetchResult<ResponseType>(
 	}
 }
 
+/**
+ * Make a fetch request to the GraphQL API, and return the JSON response.
+ */
+export async function fetchGraphqlResult<ResponseType>(
+	init: RequestInit = {},
+	abortSignal?: AbortSignal
+): Promise<ResponseType> {
+	const json = await fetchInternal<ResponseType>(
+		"/graphql",
+		{ ...init, method: "POST" }, //Cloudflare API v4 doesn't allow GETs to /graphql
+		undefined,
+		abortSignal
+	);
+	if (json) {
+		return json;
+	} else {
+		throw new Error("A request to the Cloudflare API (/graphql) failed.");
+	}
+}
+
 /**
  * Make a fetch request for a list of values,
  * extracting the `result` from the JSON response,

packages/wrangler/src/d1/info.tsx

@@ -1,7 +1,7 @@
 import Table from "ink-table";
 import prettyBytes from "pretty-bytes";
 import React from "react";
-import { fetchResult } from "../cfetch";
+import { fetchGraphqlResult, fetchResult } from "../cfetch";
 import { withConfig } from "../config";
 import { logger } from "../logger";
 import { requireAuth } from "../user";
@@ -11,7 +11,7 @@ import type {
 	CommonYargsArgv,
 	StrictYargsOptionsToInterface,
 } from "../yargs-types";
-import type { Database } from "./types";
+import type { D1MetricsGraphQLResponse, Database } from "./types";
 
 export function Options(d1ListYargs: CommonYargsArgv) {
 	return d1ListYargs
@@ -46,16 +46,86 @@ export const Handler = withConfig<HandlerOptions>(
 				},
 			}
 		);
+
+		const output: Record<string, string | number> = { ...result };
+		if (output["file_size"]) {
+			output["database_size"] = output["file_size"];
+			delete output["file_size"];
+		}
+		if (result.version === "beta") {
+			const today = new Date();
+			const yesterday = new Date(new Date(today).setDate(today.getDate() - 1));
+
+			const graphqlResult = await fetchGraphqlResult<D1MetricsGraphQLResponse>({
+				method: "POST",
+				body: JSON.stringify({
+					query: `query getD1MetricsOverviewQuery($accountTag: string, $filter: ZoneWorkersRequestsFilter_InputObject) {
+								viewer {
+									accounts(filter: {accountTag: $accountTag}) {
+										d1AnalyticsAdaptiveGroups(limit: 10000, filter: $filter) {
+											sum {
+												readQueries
+												writeQueries
+										}
+										dimensions {
+												datetimeHour
+										}
+									}
+								}
+							}
+						}`,
+					operationName: "getD1MetricsOverviewQuery",
+					variables: {
+						accountTag: accountId,
+						filter: {
+							AND: [
+								{
+									datetimeHour_geq: yesterday.toISOString(),
+									datetimeHour_leq: today.toISOString(),
+									databaseId: db.uuid,
+								},
+							],
+						},
+					},
+				}),
+				headers: {
+					"Content-Type": "application/json",
+				},
+			});
+
+			const metrics = { readQueries: 0, writeQueries: 0 };
+			if (graphqlResult) {
+				graphqlResult.data?.viewer?.accounts[0]?.d1AnalyticsAdaptiveGroups?.forEach(
+					(row) => {
+						metrics.readQueries += row?.sum?.readQueries ?? 0;
+						metrics.writeQueries += row?.sum?.writeQueries ?? 0;
+					}
+				);
+				output.read_queries_24h = metrics.readQueries;
+				output.write_queries_24h = metrics.writeQueries;
+			}
+		}
+
 		if (json) {
-			logger.log(JSON.stringify(result, null, 2));
+			logger.log(JSON.stringify(output, null, 2));
 		} else {
 			// Snip off the "uuid" property from the response and use those as the header
 
-			const entries = Object.entries(result).filter(([k, _v]) => k !== "uuid");
-			const data = entries.map(([k, v]) => ({
-				[db.binding || ""]: k,
-				[db.uuid]: k === "file_size" ? prettyBytes(Number(v)) : v,
-			}));
+			const entries = Object.entries(output).filter(([k, _v]) => k !== "uuid");
+			const data = entries.map(([k, v]) => {
+				let value;
+				if (k === "database_size") {
+					value = prettyBytes(Number(v));
+				} else if (k === "read_queries_24h" || k === "write_queries_24h") {
+					value = v.toLocaleString();
+				} else {
+					value = v;
+				}
+				return {
+					[db.binding || ""]: k,
+					[db.uuid]: value,
+				};
+			});
 
 			logger.log(renderToString(<Table data={data} />));
 		}

packages/wrangler/src/d1/types.ts

@@ -30,3 +30,34 @@ export type Migration = {
 	name: string;
 	applied_at: string;
 };
+
+export interface D1Metrics {
+	sum?: {
+		readQueries?: number;
+		writeQueries?: number;
+		queryBatchResponseBytes?: number;
+	};
+	quantiles?: {
+		queryBatchTimeMsP90?: number;
+	};
+	avg?: {
+		queryBatchTimeMs?: number;
+	};
+	dimensions: {
+		databaseId?: string;
+		date?: string;
+		datetime?: string;
+		datetimeMinute?: string;
+		datetimeFiveMinutes?: string;
+		datetimeFifteenMinutes?: string;
+		datetimeHour?: string;
+	};
+}
+
+export interface D1MetricsGraphQLResponse {
+	data: {
+		viewer: {
+			accounts: { d1AnalyticsAdaptiveGroups?: D1Metrics[] }[];
+		};
+	};
+}