Skip to content

Commit 44b6577

Browse files
ayrtonayrton
committed
Function overloading for better type checking
1 parent a556f3d commit 44b6577

File tree

2 files changed

+123
-44
lines changed

2 files changed

+123
-44
lines changed

‎__tests__/index.test.ts‎

Lines changed: 72 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -143,7 +143,7 @@ describe('execute', () => {
143143
timing: 1
144144
}
145145

146-
const want: ExecutedQuery = {
146+
const want = {
147147
headers: [':vtg1', 'null'],
148148
types: { ':vtg1': 'INT32', null: 'NULL' },
149149
fields: [
@@ -192,7 +192,7 @@ describe('execute', () => {
192192
timing: 1
193193
}
194194

195-
const want: ExecutedQuery = {
195+
const want = {
196196
headers: ['null'],
197197
types: { null: 'NULL' },
198198
fields: [{ name: 'null', type: 'NULL' }],
@@ -238,7 +238,7 @@ describe('execute', () => {
238238
timing: 1
239239
}
240240

241-
const want: ExecutedQuery = {
241+
const want = {
242242
headers: [':vtg1'],
243243
types: { ':vtg1': 'INT32' },
244244
rows: [[1]],
@@ -273,7 +273,7 @@ describe('execute', () => {
273273
mockPool.intercept({ path: EXECUTE_PATH, method: 'POST' }).reply(200, mockResponse)
274274

275275
const query = 'CREATE TABLE `foo` (bar json);'
276-
const want: ExecutedQuery = {
276+
const want = {
277277
headers: [],
278278
types: {},
279279
fields: [],
@@ -303,7 +303,7 @@ describe('execute', () => {
303303
mockPool.intercept({ path: EXECUTE_PATH, method: 'POST' }).reply(200, mockResponse)
304304

305305
const query = "UPDATE `foo` SET bar='planetscale'"
306-
const want: ExecutedQuery = {
306+
const want = {
307307
headers: [],
308308
types: {},
309309
fields: [],
@@ -334,7 +334,7 @@ describe('execute', () => {
334334
mockPool.intercept({ path: EXECUTE_PATH, method: 'POST' }).reply(200, mockResponse)
335335

336336
const query = "INSERT INTO `foo` (bar) VALUES ('planetscale');"
337-
const want: ExecutedQuery = {
337+
const want = {
338338
headers: [],
339339
types: {},
340340
fields: [],
@@ -418,7 +418,7 @@ describe('execute', () => {
418418
timing: 1
419419
}
420420

421-
const want: ExecutedQuery = {
421+
const want = {
422422
headers: [':vtg1'],
423423
rows: [{ ':vtg1': 1 }],
424424
types: { ':vtg1': 'INT32' },
@@ -452,7 +452,7 @@ describe('execute', () => {
452452
timing: 1
453453
}
454454

455-
const want: ExecutedQuery = {
455+
const want = {
456456
headers: [':vtg1'],
457457
types: { ':vtg1': 'INT32' },
458458
fields: [{ name: ':vtg1', type: 'INT32' }],
@@ -486,7 +486,7 @@ describe('execute', () => {
486486
timing: 1
487487
}
488488

489-
const want: ExecutedQuery = {
489+
const want = {
490490
headers: [':vtg1'],
491491
types: { ':vtg1': 'INT64' },
492492
fields: [{ name: ':vtg1', type: 'INT64' }],
@@ -521,7 +521,7 @@ describe('execute', () => {
521521
timing: 1
522522
}
523523

524-
const want: ExecutedQuery = {
524+
const want = {
525525
headers: [':vtg1'],
526526
types: { ':vtg1': 'INT64' },
527527
fields: [{ name: ':vtg1', type: 'INT64' }],
@@ -558,7 +558,7 @@ describe('execute', () => {
558558
timing: 1
559559
}
560560

561-
const want: ExecutedQuery = {
561+
const want = {
562562
headers: ['document'],
563563
types: { document: 'JSON' },
564564
fields: [{ name: 'document', type: 'JSON' }],
@@ -593,7 +593,67 @@ describe('execute', () => {
593593
timing: 1
594594
}
595595

596-
const want: ExecutedQuery = {
596+
const want = {
597+
headers: [':vtg1', 'null'],
598+
types: { ':vtg1': 'INT32', null: 'NULL' },
599+
fields: [
600+
{ name: ':vtg1', type: 'INT32' },
601+
{ name: 'null', type: 'NULL' }
602+
],
603+
rows: [{ ':vtg1': 1, null: null }],
604+
size: 1,
605+
statement: 'SELECT 1, null from dual;',
606+
rowsAffected: 0,
607+
insertId: '0',
608+
time: 1000
609+
}
610+
611+
interface Got {
612+
':vtg1'?: number
613+
null?: null
614+
}
615+
616+
const isGot = (object: any): object is Got => {
617+
return ':vtg1' in object && 'null' in object
618+
}
619+
620+
mockPool.intercept({ path: EXECUTE_PATH, method: 'POST' }).reply(200, (opts) => {
621+
expect(opts.headers['Authorization']).toMatch(/Basic /)
622+
const bodyObj = JSON.parse(opts.body.toString())
623+
expect(bodyObj.session).toEqual(null)
624+
return mockResponse
625+
})
626+
627+
const connection = connect(config)
628+
const got = await connection.execute<Got>('SELECT 1, null from dual;')
629+
630+
expect(got).toEqual(want)
631+
expect(isGot(got.rows[0])).toBe(true)
632+
633+
mockPool.intercept({ path: EXECUTE_PATH, method: 'POST' }).reply(200, (opts) => {
634+
expect(opts.headers['Authorization']).toMatch(/Basic /)
635+
const bodyObj = JSON.parse(opts.body.toString())
636+
expect(bodyObj.session).toEqual(mockSession)
637+
return mockResponse
638+
})
639+
640+
const got2 = await connection.execute<Got>('SELECT 1, null from dual;')
641+
642+
expect(got2).toEqual(want)
643+
expect(isGot(got2.rows[0])).toBe(true)
644+
})
645+
646+
test('allows setting the type of rows in the query results', async () => {
647+
const mockResponse = {
648+
session: mockSession,
649+
result: {
650+
fields: [{ name: ':vtg1', type: 'INT32' }, { name: 'null' }],
651+
rows: [{ lengths: ['1', '-1'], values: 'MQ==' }]
652+
},
653+
timing: 1
654+
}
655+
656+
const want = {
597657
headers: [':vtg1', 'null'],
598658
types: { ':vtg1': 'INT32', null: 'NULL' },
599659
fields: [

‎src/index.ts‎

Lines changed: 51 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -4,8 +4,6 @@ export { hex } from './text.js'
44
import { decode } from './text.js'
55
import { Version } from './version.js'
66

7-
type Row<T = Record<string, any> | any[]> = T
8-
97
interface VitessError {
108
message: string
119
code: string
@@ -24,10 +22,10 @@ export class DatabaseError extends Error {
2422

2523
type Types = Record<string, string>
2624

27-
export interface ExecutedQuery<T = any> {
25+
export interface ExecutedQuery<T> {
2826
headers: string[]
2927
types: Types
30-
rows: Row<T>[]
28+
rows: T[]
3129
fields: Field[]
3230
size: number
3331
statement: string
@@ -100,18 +98,9 @@ interface QueryResult {
10098
rows?: QueryResultRow[]
10199
}
102100

103-
type ExecuteAs = 'array' | 'object'
104-
105-
type ExecuteOptions = {
106-
as?: ExecuteAs
107-
cast?: Cast
108-
}
109-
110101
type ExecuteArgs = object | any[] | null
111102

112-
const defaultExecuteOptions: ExecuteOptions = {
113-
as: 'object'
114-
}
103+
type ExecuteAs = 'array' | 'object'
115104

116105
export class Client {
117106
private config: Config
@@ -124,12 +113,22 @@ export class Client {
124113
return this.connection().transaction(fn)
125114
}
126115

127-
async execute(
116+
async execute<T = Record<string, any>>(
117+
query: string,
118+
args?: ExecuteArgs,
119+
options?: { as?: 'object'; cast?: Cast }
120+
): Promise<ExecutedQuery<T>>
121+
async execute<T = any[]>(
122+
query: string,
123+
args: ExecuteArgs,
124+
options: { as: 'array'; cast?: Cast }
125+
): Promise<ExecutedQuery<T>>
126+
async execute<T = any>(
128127
query: string,
129128
args: ExecuteArgs = null,
130-
options: ExecuteOptions = defaultExecuteOptions
131-
): Promise<ExecutedQuery> {
132-
return this.connection().execute(query, args, options)
129+
options: any = { as: 'object' }
130+
): Promise<ExecutedQuery<T>> {
131+
return this.connection().execute<T>(query, args, options)
133132
}
134133

135134
connection(): Connection {
@@ -146,12 +145,22 @@ class Tx {
146145
this.conn = conn
147146
}
148147

149-
async execute(
148+
async execute<T = Record<string, any>>(
149+
query: string,
150+
args?: ExecuteArgs,
151+
options?: { as?: 'object'; cast?: Cast }
152+
): Promise<ExecutedQuery<T>>
153+
async execute<T = any[]>(
154+
query: string,
155+
args: ExecuteArgs,
156+
options: { as: 'array'; cast?: Cast }
157+
): Promise<ExecutedQuery<T>>
158+
async execute<T = any>(
150159
query: string,
151160
args: ExecuteArgs = null,
152-
options: ExecuteOptions = defaultExecuteOptions
153-
): Promise<ExecutedQuery> {
154-
return this.conn.execute(query, args, options)
161+
options: any = { as: 'object' }
162+
): Promise<ExecutedQuery<any>> {
163+
return this.conn.execute<T>(query, args, options)
155164
}
156165
}
157166

@@ -209,11 +218,21 @@ export class Connection {
209218
await this.createSession()
210219
}
211220

212-
async execute(
221+
async execute<T = Record<string, any>>(
222+
query: string,
223+
args?: ExecuteArgs,
224+
options?: { as?: 'object'; cast?: Cast }
225+
): Promise<ExecutedQuery<T>>
226+
async execute<T = any[]>(
227+
query: string,
228+
args: ExecuteArgs,
229+
options: { as: 'array'; cast?: Cast }
230+
): Promise<ExecutedQuery<T>>
231+
async execute<T = any>(
213232
query: string,
214233
args: ExecuteArgs = null,
215-
options: ExecuteOptions = defaultExecuteOptions
216-
): Promise<ExecutedQuery> {
234+
options: any = { as: 'object' }
235+
): Promise<ExecutedQuery<T>> {
217236
const url = new URL('/psdb.v1alpha1.Database/Execute', this.url)
218237

219238
const formatter = this.config.format || format
@@ -244,7 +263,7 @@ export class Connection {
244263
}
245264

246265
const castFn = options.cast || this.config.cast || cast
247-
const rows = result ? parse(result, castFn, options.as || 'object') : []
266+
const rows = result ? parse<T>(result, castFn, options.as || 'object') : []
248267
const headers = fields.map((f) => f.name)
249268

250269
const typeByName = (acc, { name, type }) => ({ ...acc, [name]: type })
@@ -307,28 +326,28 @@ export function connect(config: Config): Connection {
307326
return new Connection(config)
308327
}
309328

310-
function parseArrayRow(fields: Field[], rawRow: QueryResultRow, cast: Cast): Row {
329+
function parseArrayRow<T = any[]>(fields: Field[], rawRow: QueryResultRow, cast: Cast): T {
311330
const row = decodeRow(rawRow)
312331

313332
return fields.map((field, ix) => {
314333
return cast(field, row[ix])
315-
})
334+
}) as T
316335
}
317336

318-
function parseObjectRow(fields: Field[], rawRow: QueryResultRow, cast: Cast): Row {
337+
function parseObjectRow<T = Record<string, any>>(fields: Field[], rawRow: QueryResultRow, cast: Cast): T {
319338
const row = decodeRow(rawRow)
320339

321340
return fields.reduce((acc, field, ix) => {
322341
acc[field.name] = cast(field, row[ix])
323342
return acc
324-
}, {} as Row)
343+
}, {} as T)
325344
}
326345

327-
function parse(result: QueryResult, cast: Cast, returnAs: ExecuteAs): Row[] {
346+
function parse<T>(result: QueryResult, cast: Cast, returnAs: ExecuteAs): T[] {
328347
const fields = result.fields
329348
const rows = result.rows ?? []
330349
return rows.map((row) =>
331-
returnAs === 'array' ? parseArrayRow(fields, row, cast) : parseObjectRow(fields, row, cast)
350+
returnAs === 'array' ? parseArrayRow<T>(fields, row, cast) : parseObjectRow<T>(fields, row, cast)
332351
)
333352
}
334353

0 commit comments

Comments
 (0)