Skip to content

Commit a4e2439

Browse files
laplablaplab
authored
fix: update R2 SQL response format and env variable (#10654)
* fix: update response format * fix: use different env variable for R2 SQL token * fix: changeset
1 parent 924fdde commit a4e2439

File tree

6 files changed

+87
-44
lines changed

6 files changed

+87
-44
lines changed

‎.changeset/easy-guests-end.md‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
---
2+
"wrangler": minor
3+
---
4+
5+
Switch to WRANGLER_R2_SQL_AUTH_TOKEN env variable for R2 SQL secret. Update the response format for R2 SQL

‎packages/wrangler/src/__tests__/r2/sql.test.ts‎

Lines changed: 31 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -38,7 +38,7 @@ describe("r2 sql", () => {
3838
const mockToken = "test-token-123";
3939

4040
beforeEach(() => {
41-
vi.stubEnv("CLOUDFLARE_API_TOKEN", mockToken);
41+
vi.stubEnv("WRANGLER_R2_SQL_AUTH_TOKEN", mockToken);
4242
});
4343

4444
it("should require warehouse and query arguments", async () => {
@@ -51,12 +51,15 @@ describe("r2 sql", () => {
5151
);
5252
});
5353

54-
it("should require CLOUDFLARE_API_TOKEN environment variable", async () => {
54+
it("should require WRANGLER_R2_SQL_AUTH_TOKEN environment variable", async () => {
55+
vi.stubEnv("WRANGLER_R2_SQL_AUTH_TOKEN", undefined);
5556
vi.stubEnv("CLOUDFLARE_API_TOKEN", undefined);
5657

5758
await expect(
5859
runWrangler(`r2 sql query ${mockWarehouse} "${mockQuery}"`)
59-
).rejects.toThrow("Missing CLOUDFLARE_API_TOKEN environment variable");
60+
).rejects.toThrow(
61+
"Missing WRANGLER_R2_SQL_AUTH_TOKEN environment variable"
62+
);
6063
});
6164

6265
it("should validate warehouse name format", async () => {
@@ -71,19 +74,20 @@ describe("r2 sql", () => {
7174
errors: [],
7275
messages: [],
7376
result: {
74-
column_order: ["id", "name", "age"],
77+
schema: [
78+
{ name: "id", type: "Int64" },
79+
{ name: "name", type: "Utf8" },
80+
{ name: "age", type: "Int64" },
81+
],
7582
rows: [
7683
{ id: 1, name: "Alice", age: 30 },
7784
{ id: 2, name: "Bob", age: 25 },
7885
{ id: 3, name: "Charlie", age: 35 },
7986
],
80-
stats: {
81-
total_r2_requests: 5,
82-
total_r2_bytes_read: 1024 * 1024,
83-
total_r2_bytes_written: 0,
84-
total_bytes_matched: 512,
85-
total_rows_skipped: 0,
86-
total_files_scanned: 2,
87+
metrics: {
88+
r2_requests_count: 5,
89+
files_scanned: 2,
90+
bytes_scanned: 1024 * 1024,
8791
},
8892
},
8993
};
@@ -131,9 +135,14 @@ describe("r2 sql", () => {
131135
errors: [],
132136
messages: [],
133137
result: {
134-
column_order: [],
138+
schema: [],
135139
rows: [],
136140
},
141+
metrics: {
142+
r2_requests_count: 0,
143+
files_scanned: 0,
144+
bytes_scanned: 0,
145+
},
137146
};
138147

139148
msw.use(
@@ -217,11 +226,20 @@ describe("r2 sql", () => {
217226
errors: [],
218227
messages: [],
219228
result: {
220-
column_order: ["id", "name", "email"],
229+
schema: [
230+
{ name: "id", type: "Int64" },
231+
{ name: "name", type: "Utf8" },
232+
{ name: "email", type: "Utf8" },
233+
],
221234
rows: [
222235
{ id: 1, name: "Alice", email: null },
223236
{ id: 2, name: null, email: "[email protected]" },
224237
],
238+
metrics: {
239+
r2_requests_count: 5,
240+
files_scanned: 2,
241+
bytes_scanned: 1024 * 1024,
242+
},
225243
},
226244
};
227245

‎packages/wrangler/src/environment-variables/factory.ts‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,8 @@ type VariableNames =
1919
| "CLOUDFLARE_API_BASE_URL"
2020
/** Set to "fedramp_high" for FedRAMP High compliance region. This will update the API/AUTH URLs used to make requests to Cloudflare. */
2121
| "CLOUDFLARE_COMPLIANCE_REGION"
22+
/** API token for R2 SQL service. */
23+
| "WRANGLER_R2_SQL_AUTH_TOKEN"
2224

2325
// ## Development & Local Testing
2426

‎packages/wrangler/src/r2/sql.ts‎

Lines changed: 43 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -6,34 +6,35 @@ import { createCommand, createNamespace } from "../core/create-command";
66
import { UserError } from "../errors";
77
import { logger } from "../logger";
88
import { APIError, parseJSON } from "../parse";
9-
import { getCloudflareAPITokenFromEnv } from "../user/auth-variables";
9+
import {
10+
getCloudflareAPITokenFromEnv,
11+
getWranglerR2SqlAuthToken,
12+
} from "../user/auth-variables";
1013

11-
interface SqlQueryResult {
14+
interface SqlQueryResponse {
1215
result?: {
13-
column_order: string[];
16+
request_id?: string;
17+
schema: { name: string; type: string }[];
1418
rows: Record<string, unknown>[];
15-
stats?: {
16-
total_r2_requests: number;
17-
total_r2_bytes_read: number;
18-
total_r2_bytes_written: number;
19-
total_bytes_matched: number;
20-
total_rows_skipped: number;
21-
total_files_scanned: number;
19+
metrics: {
20+
r2_requests_count: number;
21+
files_scanned: number;
22+
bytes_scanned: number;
2223
};
2324
};
2425
success: boolean;
2526
errors: { code: number; message: string }[];
2627
messages: string[];
2728
}
2829

29-
function formatSqlResults(data: SqlQueryResult, duration: number): void {
30+
function formatSqlResults(data: SqlQueryResponse, duration: number): void {
3031
if (!data?.result?.rows || data.result.rows.length === 0) {
3132
logger.log("Query executed successfully with no results");
3233
return;
3334
}
3435

35-
const { column_order, rows, stats } = data.result;
36-
36+
const { schema, rows, metrics } = data.result;
37+
const column_order = schema.map((field) => field.name);
3738
logger.table(
3839
rows.map((row) =>
3940
Object.fromEntries(
@@ -43,15 +44,12 @@ function formatSqlResults(data: SqlQueryResult, duration: number): void {
4344
{ wordWrap: true, head: column_order }
4445
);
4546

46-
// Print stats if available.
47-
if (stats) {
48-
logger.log(
49-
`Read ${prettyBytes(stats.total_r2_bytes_read)} across ${stats.total_files_scanned} files from R2`
50-
);
51-
if (duration > 0) {
52-
const bytesPerSecond = (stats.total_r2_bytes_read / duration) * 1000;
53-
logger.log(`On average, ${prettyBytes(bytesPerSecond)} / s`);
54-
}
47+
logger.log(
48+
`Read ${prettyBytes(metrics.bytes_scanned)} across ${metrics.files_scanned} files from R2`
49+
);
50+
if (duration > 0) {
51+
const bytesPerSecond = (metrics.bytes_scanned / duration) * 1000;
52+
logger.log(`On average, ${prettyBytes(bytesPerSecond)} / s`);
5553
}
5654
}
5755

@@ -83,14 +81,22 @@ export const r2SqlQueryCommand = createCommand({
8381
},
8482
},
8583
async handler({ warehouse, query }) {
86-
const token = getCloudflareAPITokenFromEnv();
84+
let token = getWranglerR2SqlAuthToken();
8785
if (!token) {
88-
throw new UserError(
89-
"Missing CLOUDFLARE_API_TOKEN environment variable. " +
90-
"Please follow instructions in https://developers.cloudflare.com/r2/sql/platform/troubleshooting/ to create a token. " +
91-
"Once done, you can prefix the command with the variable definition like so: `CLOUDFLARE_API_TOKEN=... wrangler r2 sql query ...`. " +
92-
"There also other ways to provide the value of this variable, see https://developers.cloudflare.com/workers/wrangler/system-environment-variables/ for more details."
93-
);
86+
token = getCloudflareAPITokenFromEnv();
87+
if (!token) {
88+
throw new UserError(
89+
"Missing WRANGLER_R2_SQL_AUTH_TOKEN environment variable. " +
90+
"Tried to fallback to CLOUDFLARE_API_TOKEN, didn't find it either. " +
91+
"Please follow instructions in https://developers.cloudflare.com/r2/sql/platform/troubleshooting/ to create a token. " +
92+
"Once done, you can prefix the command with the variable definition like so: `WRANGLER_R2_SQL_AUTH_TOKEN=... wrangler r2 sql query ...`. " +
93+
"There also other ways to provide the value of this variable, see https://developers.cloudflare.com/workers/wrangler/system-environment-variables/ for more details."
94+
);
95+
} else {
96+
logger.warn(
97+
"Missing WRANGLER_R2_SQL_AUTH_TOKEN environment variable, falling back to CLOUDFLARE_API_TOKEN"
98+
);
99+
}
94100
}
95101

96102
const splitIndex = warehouse.indexOf("_");
@@ -132,9 +138,16 @@ export const r2SqlQueryCommand = createCommand({
132138
});
133139
}
134140

141+
if (responseStatus === 403) {
142+
logger.error(
143+
"Please check that token in WRANGLER_R2_SQL_AUTH_TOKEN or CLOUDFLARE_API_TOKEN has the correct permissions. " +
144+
"See https://developers.cloudflare.com/r2/sql/platform/troubleshooting/ for more details."
145+
);
146+
}
147+
135148
let parsed = null;
136149
try {
137-
parsed = parseJSON(text) as SqlQueryResult;
150+
parsed = parseJSON(text) as SqlQueryResponse;
138151
} catch {
139152
throw new APIError({
140153
text: "Received a malformed response from the API",

‎packages/wrangler/src/user/auth-variables.ts‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -95,6 +95,10 @@ export const getRevokeUrlFromEnv = getEnvironmentVariableFactory({
9595
defaultValue: () => `https://${getAuthDomainFromEnv()}/oauth2/revoke`,
9696
});
9797

98+
export const getWranglerR2SqlAuthToken = getEnvironmentVariableFactory({
99+
variableName: "WRANGLER_R2_SQL_AUTH_TOKEN",
100+
});
101+
98102
/**
99103
* Set the `WRANGLER_CF_AUTHORIZATION_TOKEN` to the CF_Authorization token found at https://dash.staging.cloudflare.com/bypass-limits
100104
* if you want to access the staging environment, triggered by `WRANGLER_API_ENVIRONMENT=staging`.

‎packages/wrangler/turbo.json‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -52,7 +52,8 @@
5252
"WRANGLER_SEND_METRICS",
5353
"WRANGLER_WORKER_REGISTRY_PORT",
5454
"DOCKER_HOST",
55-
"WRANGLER_DOCKER_HOST"
55+
"WRANGLER_DOCKER_HOST",
56+
"WRANGLER_R2_SQL_AUTH_TOKEN"
5657
]
5758
},
5859
"test:ci": {

0 commit comments

Comments
 (0)