Skip to content

Commit ae407ac

Browse files
ymc9claude
andcommitted
feat(orm): add $diagnostics() method for cache stats and slow query tracking
Introduces a new `$diagnostics()` method on ZenStackClient that returns Zod schema cache statistics and slow query information, helping users monitor and debug ORM performance. - Add `diagnostics` option to `ClientOptions` with `slowQueryThresholdMs` and `slowQueryMaxRecords` settings - Track slow queries in `ZenStackQueryExecutor` when diagnostics is enabled - Share slow query collection across derived clients (via $setAuth, $setOptions, $use, transactions, etc.) - Cap slow query records with an eviction policy that keeps the slowest queries (default max: 100) - Validate diagnostics config with Zod in ClientImpl constructor - Add `Diagnostics`, `QueryInfo`, and `ZodCacheStats` types - Add e2e tests covering all diagnostics features Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
1 parent d4fbb38 commit ae407ac

File tree

7 files changed

+382
-49
lines changed

7 files changed

+382
-49
lines changed

packages/orm/src/client/client-impl.ts

Lines changed: 53 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -10,9 +10,11 @@ import {
1010
Transaction,
1111
type KyselyProps,
1212
} from 'kysely';
13+
import z from 'zod';
1314
import type { ProcedureDef, SchemaDef } from '../schema';
1415
import type { AnyKysely } from '../utils/kysely-utils';
1516
import type { UnwrapTuplePromises } from '../utils/type-utils';
17+
import { formatError } from '../utils/zod-utils';
1618
import type {
1719
AuthType,
1820
ClientConstructor,
@@ -31,6 +33,7 @@ import { FindOperationHandler } from './crud/operations/find';
3133
import { GroupByOperationHandler } from './crud/operations/group-by';
3234
import { UpdateOperationHandler } from './crud/operations/update';
3335
import { InputValidator } from './crud/validator';
36+
import type { Diagnostics, QueryInfo } from './diagnostics';
3437
import { createConfigError, createNotFoundError, createNotSupportedError } from './errors';
3538
import { ZenStackDriver } from './executor/zenstack-driver';
3639
import { ZenStackQueryExecutor } from './executor/zenstack-query-executor';
@@ -60,6 +63,7 @@ export class ClientImpl {
6063
readonly kyselyProps: KyselyProps;
6164
private auth: AuthType<SchemaDef> | undefined;
6265
inputValidator: InputValidator<SchemaDef>;
66+
readonly slowQueries: QueryInfo[] = [];
6367

6468
constructor(
6569
private readonly schema: SchemaDef,
@@ -75,10 +79,7 @@ export class ClientImpl {
7579
...this.$options.functions,
7680
};
7781

78-
if (!baseClient && !options.skipValidationForComputedFields) {
79-
// validate computed fields configuration once for the root client
80-
this.validateComputedFieldsConfig();
81-
}
82+
this.validateOptions(baseClient, options);
8283

8384
// here we use kysely's props constructor so we can pass a custom query executor
8485
if (baseClient) {
@@ -96,6 +97,7 @@ export class ClientImpl {
9697
};
9798
this.kyselyRaw = baseClient.kyselyRaw;
9899
this.auth = baseClient.auth;
100+
this.slowQueries = baseClient.slowQueries;
99101
} else {
100102
const driver = new ZenStackDriver(options.dialect.createDriver(), new Log(this.$options.log ?? []));
101103
const compiler = options.dialect.createQueryCompiler();
@@ -125,37 +127,30 @@ export class ClientImpl {
125127
return createClientProxy(this);
126128
}
127129

128-
get $qb() {
129-
return this.kysely;
130-
}
131-
132-
get $qbRaw() {
133-
return this.kyselyRaw;
134-
}
135-
136-
get $zod() {
137-
return this.inputValidator.zodFactory;
138-
}
139-
140-
get isTransaction() {
141-
return this.kysely.isTransaction;
142-
}
130+
private validateOptions(baseClient: ClientImpl | undefined, options: ClientOptions<SchemaDef>) {
131+
if (!baseClient && !options.skipValidationForComputedFields) {
132+
// validate computed fields configuration once for the root client
133+
this.validateComputedFieldsConfig(options);
134+
}
143135

144-
/**
145-
* Create a new client with a new query executor.
146-
*/
147-
withExecutor(executor: QueryExecutor) {
148-
return new ClientImpl(this.schema, this.$options, this, executor);
136+
if (options.diagnostics) {
137+
const diagnosticsSchema = z.object({
138+
slowQueryThresholdMs: z.int().nonnegative().optional(),
139+
slowQueryMaxRecords: z.int().nonnegative().optional(),
140+
});
141+
const parseResult = diagnosticsSchema.safeParse(options.diagnostics);
142+
if (!parseResult.success) {
143+
throw createConfigError(`Invalid diagnostics configuration: ${formatError(parseResult.error)}`);
144+
}
145+
}
149146
}
150147

151148
/**
152149
* Validates that all computed fields in the schema have corresponding configurations.
153150
*/
154-
private validateComputedFieldsConfig() {
151+
private validateComputedFieldsConfig(options: ClientOptions<SchemaDef>) {
155152
const computedFieldsConfig =
156-
'computedFields' in this.$options
157-
? (this.$options.computedFields as Record<string, any> | undefined)
158-
: undefined;
153+
'computedFields' in options ? (options.computedFields as Record<string, any> | undefined) : undefined;
159154

160155
for (const [modelName, modelDef] of Object.entries(this.$schema.models)) {
161156
if (modelDef.computedFields) {
@@ -181,6 +176,29 @@ export class ClientImpl {
181176
}
182177
}
183178

179+
get $qb() {
180+
return this.kysely;
181+
}
182+
183+
get $qbRaw() {
184+
return this.kyselyRaw;
185+
}
186+
187+
get $zod() {
188+
return this.inputValidator.zodFactory;
189+
}
190+
191+
get isTransaction() {
192+
return this.kysely.isTransaction;
193+
}
194+
195+
/**
196+
* Create a new client with a new query executor.
197+
*/
198+
withExecutor(executor: QueryExecutor) {
199+
return new ClientImpl(this.schema, this.$options, this, executor);
200+
}
201+
184202
// overload for interactive transaction
185203
$transaction<T>(
186204
callback: (tx: ClientContract<SchemaDef>) => Promise<T>,
@@ -412,6 +430,13 @@ export class ClientImpl {
412430
return this.$setOptions(newOptions);
413431
}
414432

433+
async $diagnostics(): Promise<Diagnostics> {
434+
return {
435+
zodCache: this.inputValidator.zodFactory.cacheStats,
436+
slowQueries: [...this.slowQueries],
437+
};
438+
}
439+
415440
$executeRaw(query: TemplateStringsArray, ...values: any[]) {
416441
return createZenStackPromise(async () => {
417442
const result = await sql(query, ...values).execute(this.kysely);

packages/orm/src/client/contract.ts

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,7 @@ import type {
3939
UpdateManyArgs,
4040
UpsertArgs,
4141
} from './crud-types';
42+
import type { Diagnostics } from './diagnostics';
4243
import type { ClientOptions, QueryOptions } from './options';
4344
import type { ExtClientMembersBase, ExtQueryArgsBase, RuntimePlugin } from './plugin';
4445
import type { ZenStackPromise } from './promise';
@@ -212,6 +213,11 @@ export type ClientContract<
212213
* @private
213214
*/
214215
$pushSchema(): Promise<void>;
216+
217+
/**
218+
* Returns diagnostics information such as cache and slow query statistics.
219+
*/
220+
$diagnostics(): Promise<Diagnostics>;
215221
} & {
216222
[Key in GetSlicedModels<Schema, Options> as Uncapitalize<Key>]: ModelOperations<Schema, Key, Options, ExtQueryArgs>;
217223
} & ProcedureOperations<Schema, Options> &
Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,44 @@
1+
/**
2+
* Zod schema cache statistics.
3+
*/
4+
export interface ZodCacheStats {
5+
/**
6+
* Number of cached Zod schemas.
7+
*/
8+
size: number;
9+
10+
/**
11+
* Keys of the cached Zod schemas.
12+
*/
13+
keys: string[];
14+
}
15+
16+
/**
17+
* Information about a query, used for diagnostics.
18+
*/
19+
export interface QueryInfo {
20+
/**
21+
* Duration of the query in milliseconds.
22+
*/
23+
durationMs: number;
24+
25+
/**
26+
* SQL statement of the query.
27+
*/
28+
sql: string;
29+
}
30+
31+
/**
32+
* ZenStackClient diagnostics.
33+
*/
34+
export interface Diagnostics {
35+
/**
36+
* Statistics about the Zod schemas (used for query args validation) cache.
37+
*/
38+
zodCache: ZodCacheStats;
39+
40+
/**
41+
* Slow queries.
42+
*/
43+
slowQueries: QueryInfo[];
44+
}

packages/orm/src/client/executor/zenstack-query-executor.ts

Lines changed: 41 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,7 @@ import {
3131
} from 'kysely';
3232
import { match } from 'ts-pattern';
3333
import type { ModelDef, SchemaDef, TypeDefDef } from '../../schema';
34-
import { type ClientImpl } from '../client-impl';
34+
import type { ClientImpl } from '../client-impl';
3535
import { TransactionIsolationLevel, type ClientContract } from '../contract';
3636
import { getCrudDialect } from '../crud/dialects';
3737
import type { BaseCrudDialect } from '../crud/dialects/base-dialect';
@@ -70,6 +70,8 @@ type CallAfterMutationHooksArgs = {
7070
afterMutationEntities?: Record<string, unknown>[];
7171
};
7272

73+
const DEFAULT_MAX_SLOW_RECORDS = 100;
74+
7375
export class ZenStackQueryExecutor extends DefaultQueryExecutor {
7476
// #region constructor, fields and props
7577

@@ -673,8 +675,16 @@ In such cases, ZenStack cannot reliably determine the IDs of the mutated entitie
673675
compiledQuery = { ...compiledQuery, parameters: parameters };
674676
}
675677

678+
const trackSlowQuery = this.options.diagnostics !== undefined;
679+
const startTime = trackSlowQuery ? performance.now() : undefined;
680+
676681
try {
677682
const result = await connection.executeQuery<any>(compiledQuery);
683+
684+
if (startTime !== undefined) {
685+
this.trackSlowQuery(compiledQuery, startTime);
686+
}
687+
678688
return this.ensureProperQueryResult(compiledQuery.query, result);
679689
} catch (err) {
680690
throw createDBQueryError(
@@ -686,6 +696,36 @@ In such cases, ZenStack cannot reliably determine the IDs of the mutated entitie
686696
}
687697
}
688698

699+
private trackSlowQuery(compiledQuery: CompiledQuery, startTime: number) {
700+
const durationMs = performance.now() - startTime;
701+
const thresholdMs = this.options.diagnostics?.slowQueryThresholdMs;
702+
if (thresholdMs === undefined || durationMs < thresholdMs) {
703+
return;
704+
}
705+
706+
const slowQueries = this.client.slowQueries;
707+
const maxRecords = this.options.diagnostics?.slowQueryMaxRecords ?? DEFAULT_MAX_SLOW_RECORDS;
708+
if (maxRecords <= 0) {
709+
return;
710+
}
711+
712+
if (slowQueries.length >= maxRecords) {
713+
// find and remove the entry with the lowest duration
714+
let minIndex = 0;
715+
for (let i = 1; i < slowQueries.length; i++) {
716+
if (slowQueries[i]!.durationMs < slowQueries[minIndex]!.durationMs) {
717+
minIndex = i;
718+
}
719+
}
720+
// only replace if the new query is slower than the minimum
721+
if (durationMs > slowQueries[minIndex]!.durationMs) {
722+
slowQueries[minIndex] = { durationMs, sql: compiledQuery.sql };
723+
}
724+
} else {
725+
slowQueries.push({ durationMs, sql: compiledQuery.sql });
726+
}
727+
}
728+
689729
private ensureProperQueryResult(query: RootOperationNode, result: QueryResult<any>) {
690730
let finalResult = result;
691731

packages/orm/src/client/options.ts

Lines changed: 19 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -178,7 +178,8 @@ export type ClientOptions<Schema extends SchemaDef> = QueryOptions<Schema> & {
178178
plugins?: AnyPlugin[];
179179

180180
/**
181-
* Logging configuration.
181+
* Logging configuration. Extends Kysely's log config with a `'warning'` level
182+
* for ZenStack-specific diagnostics (e.g., slow query warnings).
182183
*/
183184
log?: KyselyConfig['log'];
184185

@@ -208,9 +209,25 @@ export type ClientOptions<Schema extends SchemaDef> = QueryOptions<Schema> & {
208209
useCompactAliasNames?: boolean;
209210

210211
/**
211-
* Whether to skip validation for computed fields.
212+
* Whether to skip validation for whether all computed fields are properly defined.
212213
*/
213214
skipValidationForComputedFields?: boolean;
215+
216+
/**
217+
* Diagnostics related options.
218+
*/
219+
diagnostics?: {
220+
/**
221+
* Threshold in milliseconds for determining slow queries. If not specified, no query will be considered slow.
222+
*/
223+
slowQueryThresholdMs?: number;
224+
225+
/**
226+
* Maximum number of slow query records to keep in memory. Defaults to `100`. When the number is exceeded, the
227+
* entry with the lowest duration will be removed. Set to `Infinity` to keep unlimited records.
228+
*/
229+
slowQueryMaxRecords?: number;
230+
};
214231
} & (HasComputedFields<Schema> extends true
215232
? {
216233
/**

0 commit comments

Comments
 (0)