Skip to content

Commit 30f558e

Browse files
laplabpenalosa
laplab
and
penalosa
authored
feat: add R2 SQL command (#10558)
* feat: r2 sql commands * fix: format results, fix bugs * fix: renames * fix: add tests * fix: use common framework for the token env variable * fix: add changeset * fix: snapshots * fix: formatting * fix: review * fix: replace UserError with APIError where appropriate * feat: add spinner while waiting for query result * fix: tests * fix: remove enable/disable commands * fix: update query url * fix: refactor Co-authored-by: Somhairle MacLeòid <[email protected]> * fix: docs link and refactor * fix: style * fix: tests * fix: lint --------- Co-authored-by: Somhairle MacLeòid <[email protected]>
1 parent 0837a8d commit 30f558e

File tree

9 files changed

+441
-7
lines changed

9 files changed

+441
-7
lines changed

‎.changeset/nice-guests-smoke.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+
Add commands to send queries and manage R2 SQL product.

‎packages/wrangler/src/__tests__/index.test.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -294,6 +294,7 @@ describe("wrangler", () => {
294294
COMMANDS
295295
wrangler r2 object Manage R2 objects
296296
wrangler r2 bucket Manage R2 buckets
297+
wrangler r2 sql Send queries and manage R2 SQL [open-beta]
297298
298299
GLOBAL FLAGS
299300
-c, --config Path to Wrangler configuration file [string]

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -99,6 +99,7 @@ describe("r2", () => {
9999
COMMANDS
100100
wrangler r2 object Manage R2 objects
101101
wrangler r2 bucket Manage R2 buckets
102+
wrangler r2 sql Send queries and manage R2 SQL [open-beta]
102103
103104
GLOBAL FLAGS
104105
-c, --config Path to Wrangler configuration file [string]
@@ -129,6 +130,7 @@ describe("r2", () => {
129130
COMMANDS
130131
wrangler r2 object Manage R2 objects
131132
wrangler r2 bucket Manage R2 buckets
133+
wrangler r2 sql Send queries and manage R2 SQL [open-beta]
132134
133135
GLOBAL FLAGS
134136
-c, --config Path to Wrangler configuration file [string]

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

Lines changed: 244 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,244 @@
1+
import { http, HttpResponse } from "msw";
2+
import { endEventLoop } from "../helpers/end-event-loop";
3+
import { mockAccountId, mockApiToken } from "../helpers/mock-account-id";
4+
import { mockConsoleMethods } from "../helpers/mock-console";
5+
import { msw } from "../helpers/msw";
6+
import { runInTempDir } from "../helpers/run-in-tmp";
7+
import { runWrangler } from "../helpers/run-wrangler";
8+
9+
describe("r2 sql", () => {
10+
const std = mockConsoleMethods();
11+
mockAccountId();
12+
mockApiToken();
13+
runInTempDir();
14+
15+
describe("help", () => {
16+
it("should show help when no subcommand is passed", async () => {
17+
await runWrangler("r2 sql");
18+
await endEventLoop();
19+
expect(std.out).toContain("wrangler r2 sql");
20+
expect(std.out).toContain("Send queries and manage R2 SQL");
21+
expect(std.out).toContain("wrangler r2 sql query <warehouse> <query>");
22+
});
23+
24+
it("should show help for query command", async () => {
25+
await runWrangler("r2 sql query --help");
26+
await endEventLoop();
27+
expect(std.out).toContain("Execute SQL query against R2 Data Catalog");
28+
expect(std.out).toContain("warehouse");
29+
expect(std.out).toContain("R2 Data Catalog warehouse name");
30+
expect(std.out).toContain("query");
31+
expect(std.out).toContain("The SQL query to execute");
32+
});
33+
});
34+
35+
describe("query", () => {
36+
const mockWarehouse = "account123_mybucket";
37+
const mockQuery = "SELECT * FROM data";
38+
const mockToken = "test-token-123";
39+
40+
beforeEach(() => {
41+
vi.stubEnv("CLOUDFLARE_API_TOKEN", mockToken);
42+
});
43+
44+
it("should require warehouse and query arguments", async () => {
45+
await expect(runWrangler("r2 sql query")).rejects.toThrow(
46+
"Not enough non-option arguments: got 0, need at least 2"
47+
);
48+
49+
await expect(runWrangler("r2 sql query testWarehouse")).rejects.toThrow(
50+
"Not enough non-option arguments: got 1, need at least 2"
51+
);
52+
});
53+
54+
it("should require CLOUDFLARE_API_TOKEN environment variable", async () => {
55+
vi.stubEnv("CLOUDFLARE_API_TOKEN", undefined);
56+
57+
await expect(
58+
runWrangler(`r2 sql query ${mockWarehouse} "${mockQuery}"`)
59+
).rejects.toThrow("Missing CLOUDFLARE_API_TOKEN environment variable");
60+
});
61+
62+
it("should validate warehouse name format", async () => {
63+
await expect(
64+
runWrangler(`r2 sql query invalidwarehouse "${mockQuery}"`)
65+
).rejects.toThrow("Warehouse name has invalid format");
66+
});
67+
68+
it("should execute a successful query and display results", async () => {
69+
const mockResponse = {
70+
success: true,
71+
errors: [],
72+
messages: [],
73+
result: {
74+
column_order: ["id", "name", "age"],
75+
rows: [
76+
{ id: 1, name: "Alice", age: 30 },
77+
{ id: 2, name: "Bob", age: 25 },
78+
{ id: 3, name: "Charlie", age: 35 },
79+
],
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+
},
88+
},
89+
};
90+
91+
msw.use(
92+
http.post(
93+
"https://api.sql.cloudflarestorage.com/api/v1/accounts/:accountId/r2-sql/query/:bucketName",
94+
async ({ request, params }) => {
95+
const { accountId, bucketName } = params;
96+
expect(accountId).toEqual("account123");
97+
expect(bucketName).toEqual("mybucket");
98+
99+
const body = (await request.json()) as {
100+
warehouse: string;
101+
query: string;
102+
};
103+
expect(body.warehouse).toEqual(mockWarehouse);
104+
expect(body.query).toEqual(mockQuery);
105+
106+
const authHeader = request.headers.get("Authorization");
107+
expect(authHeader).toEqual(`Bearer ${mockToken}`);
108+
109+
return HttpResponse.json(mockResponse);
110+
},
111+
{ once: true }
112+
)
113+
);
114+
115+
await runWrangler(`r2 sql query ${mockWarehouse} "${mockQuery}"`);
116+
117+
// Check that results are displayed in a table format.
118+
expect(std.out).toContain("id");
119+
expect(std.out).toContain("name");
120+
expect(std.out).toContain("age");
121+
expect(std.out).toContain("Alice");
122+
expect(std.out).toContain("Bob");
123+
expect(std.out).toContain("Charlie");
124+
expect(std.out).toContain("across 2 files from R2");
125+
// Not checking MB/s speed as it depends on timing.
126+
});
127+
128+
it("should handle queries with no results", async () => {
129+
const mockResponse = {
130+
success: true,
131+
errors: [],
132+
messages: [],
133+
result: {
134+
column_order: [],
135+
rows: [],
136+
},
137+
};
138+
139+
msw.use(
140+
http.post(
141+
"https://api.sql.cloudflarestorage.com/api/v1/accounts/:accountId/r2-sql/query/:bucketName",
142+
async () => {
143+
return HttpResponse.json(mockResponse);
144+
},
145+
{ once: true }
146+
)
147+
);
148+
149+
await runWrangler(`r2 sql query ${mockWarehouse} "${mockQuery}"`);
150+
expect(std.out).toContain("Query executed successfully with no results");
151+
});
152+
153+
it("should handle query failures", async () => {
154+
const mockResponse = {
155+
success: false,
156+
errors: [
157+
{ code: 1001, message: "Syntax error in SQL query" },
158+
{ code: 1002, message: "Table not found" },
159+
],
160+
messages: [],
161+
result: null,
162+
};
163+
164+
msw.use(
165+
http.post(
166+
"https://api.sql.cloudflarestorage.com/api/v1/accounts/:accountId/r2-sql/query/:bucketName",
167+
async () => {
168+
return HttpResponse.json(mockResponse, { status: 500 });
169+
},
170+
{ once: true }
171+
)
172+
);
173+
174+
await runWrangler(`r2 sql query ${mockWarehouse} "${mockQuery}"`);
175+
expect(std.err).toContain(
176+
"Query failed because of the following errors:"
177+
);
178+
expect(std.err).toContain("1001: Syntax error in SQL query");
179+
expect(std.err).toContain("1002: Table not found");
180+
});
181+
182+
it("should handle API connection errors", async () => {
183+
msw.use(
184+
http.post(
185+
"https://api.sql.cloudflarestorage.com/api/v1/accounts/:accountId/r2-sql/query/:bucketName",
186+
async () => {
187+
return HttpResponse.error();
188+
},
189+
{ once: true }
190+
)
191+
);
192+
193+
await expect(
194+
runWrangler(`r2 sql query ${mockWarehouse} "${mockQuery}"`)
195+
).rejects.toThrow("Failed to connect to R2 SQL API");
196+
});
197+
198+
it("should handle invalid JSON responses", async () => {
199+
msw.use(
200+
http.post(
201+
"https://api.sql.cloudflarestorage.com/api/v1/accounts/:accountId/r2-sql/query/:bucketName",
202+
async () => {
203+
return HttpResponse.text("Invalid JSON", { status: 200 });
204+
},
205+
{ once: true }
206+
)
207+
);
208+
209+
await expect(
210+
runWrangler(`r2 sql query ${mockWarehouse} "${mockQuery}"`)
211+
).rejects.toThrow("Received a malformed response from the API");
212+
});
213+
214+
it("should handle null values in query results", async () => {
215+
const mockResponse = {
216+
success: true,
217+
errors: [],
218+
messages: [],
219+
result: {
220+
column_order: ["id", "name", "email"],
221+
rows: [
222+
{ id: 1, name: "Alice", email: null },
223+
{ id: 2, name: null, email: "[email protected]" },
224+
],
225+
},
226+
};
227+
228+
msw.use(
229+
http.post(
230+
"https://api.sql.cloudflarestorage.com/api/v1/accounts/:accountId/r2-sql/query/:bucketName",
231+
async () => {
232+
return HttpResponse.json(mockResponse);
233+
},
234+
{ once: true }
235+
)
236+
);
237+
238+
await runWrangler(`r2 sql query ${mockWarehouse} "${mockQuery}"`);
239+
// The output should handle null values gracefully (displayed as empty strings).
240+
expect(std.out).toContain("Alice");
241+
expect(std.out).toContain("[email protected]");
242+
});
243+
});
244+
});

‎packages/wrangler/src/cfetch/internal.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -203,7 +203,7 @@ export async function fetchInternal<ResponseType>(
203203
}
204204
}
205205

206-
function truncate(text: string, maxLength: number): string {
206+
export function truncate(text: string, maxLength: number): string {
207207
const { length } = text;
208208
if (length <= maxLength) {
209209
return text;

‎packages/wrangler/src/core/teams.d.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ export type Teams =
1010
| "Product: KV"
1111
| "Product: R2"
1212
| "Product: R2 Data Catalog"
13+
| "Product: R2 SQL"
1314
| "Product: D1"
1415
| "Product: Queues"
1516
| "Product: AI"

‎packages/wrangler/src/index.ts‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -248,6 +248,7 @@ import {
248248
r2BucketSippyGetCommand,
249249
r2BucketSippyNamespace,
250250
} from "./r2/sippy";
251+
import { r2SqlNamespace, r2SqlQueryCommand } from "./r2/sql";
251252
import {
252253
secretBulkCommand,
253254
secretDeleteCommand,
@@ -957,6 +958,14 @@ export function createCLIParser(argv: string[]) {
957958
command: "wrangler r2 bucket lock set",
958959
definition: r2BucketLockSetCommand,
959960
},
961+
{
962+
command: "wrangler r2 sql",
963+
definition: r2SqlNamespace,
964+
},
965+
{
966+
command: "wrangler r2 sql query",
967+
definition: r2SqlQueryCommand,
968+
},
960969
]);
961970
registry.registerNamespace("r2");
962971

‎packages/wrangler/src/logger.ts‎

Lines changed: 14 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -122,17 +122,25 @@ export class Logger {
122122
log = (...args: unknown[]) => this.doLog("log", args);
123123
warn = (...args: unknown[]) => this.doLog("warn", args);
124124
error = (...args: unknown[]) => this.doLog("error", args);
125-
table<Keys extends string>(data: TableRow<Keys>[]) {
126-
const keys: Keys[] =
125+
table<Keys extends string>(
126+
data: TableRow<Keys>[],
127+
options?: { wordWrap: boolean; head?: Keys[] }
128+
) {
129+
const derivedHead =
127130
data.length === 0 ? [] : (Object.keys(data[0]) as Keys[]);
128-
const t = new CLITable({
129-
head: keys,
131+
const wordWrap = options?.wordWrap ?? false;
132+
const head = options?.head ?? derivedHead;
133+
134+
const tableOptions = {
130135
style: {
131136
head: chalk.level ? ["blue"] : [],
132137
border: chalk.level ? ["gray"] : [],
133138
},
134-
});
135-
t.push(...data.map((row) => keys.map((k) => row[k])));
139+
wordWrap,
140+
head,
141+
};
142+
const t = new CLITable(tableOptions);
143+
t.push(...data.map((row) => head.map((k) => row[k])));
136144
return this.doLog("log", [t.toString()]);
137145
}
138146
console<M extends Exclude<keyof Console, "Console">>(

0 commit comments

Comments
 (0)