feat(orm): add $diagnostics() for cache stats and slow query tracking (#2481)

Author: ymc9

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

@@ -10,9 +10,11 @@ import {
     Transaction,
     type KyselyProps,
 } from 'kysely';
+import z from 'zod';
 import type { ProcedureDef, SchemaDef } from '../schema';
 import type { AnyKysely } from '../utils/kysely-utils';
 import type { UnwrapTuplePromises } from '../utils/type-utils';
+import { formatError } from '../utils/zod-utils';
 import type {
     AuthType,
     ClientConstructor,
@@ -31,6 +33,7 @@ import { FindOperationHandler } from './crud/operations/find';
 import { GroupByOperationHandler } from './crud/operations/group-by';
 import { UpdateOperationHandler } from './crud/operations/update';
 import { InputValidator } from './crud/validator';
+import type { Diagnostics, QueryInfo } from './diagnostics';
 import { createConfigError, createNotFoundError, createNotSupportedError } from './errors';
 import { ZenStackDriver } from './executor/zenstack-driver';
 import { ZenStackQueryExecutor } from './executor/zenstack-query-executor';
@@ -60,6 +63,7 @@ export class ClientImpl {
     readonly kyselyProps: KyselyProps;
     private auth: AuthType<SchemaDef> | undefined;
     inputValidator: InputValidator<SchemaDef>;
+    readonly slowQueries: QueryInfo[] = [];
 
     constructor(
         private readonly schema: SchemaDef,
@@ -75,10 +79,7 @@ export class ClientImpl {
             ...this.$options.functions,
         };
 
-        if (!baseClient && !options.skipValidationForComputedFields) {
-            // validate computed fields configuration once for the root client
-            this.validateComputedFieldsConfig();
-        }
+        this.validateOptions(baseClient, options);
 
         // here we use kysely's props constructor so we can pass a custom query executor
         if (baseClient) {
@@ -96,6 +97,7 @@ export class ClientImpl {
             };
             this.kyselyRaw = baseClient.kyselyRaw;
             this.auth = baseClient.auth;
+            this.slowQueries = baseClient.slowQueries;
         } else {
             const driver = new ZenStackDriver(options.dialect.createDriver(), new Log(this.$options.log ?? []));
             const compiler = options.dialect.createQueryCompiler();
@@ -127,37 +129,30 @@ export class ClientImpl {
         return createClientProxy(this);
     }
 
-    get $qb() {
-        return this.kysely;
-    }
-
-    get $qbRaw() {
-        return this.kyselyRaw;
-    }
-
-    get $zod() {
-        return this.inputValidator.zodFactory;
-    }
-
-    get isTransaction() {
-        return this.kysely.isTransaction;
-    }
+    private validateOptions(baseClient: ClientImpl | undefined, options: ClientOptions<SchemaDef>) {
+        if (!baseClient && !options.skipValidationForComputedFields) {
+            // validate computed fields configuration once for the root client
+            this.validateComputedFieldsConfig(options);
+        }
 
-    /**
-     * Create a new client with a new query executor.
-     */
-    withExecutor(executor: QueryExecutor) {
-        return new ClientImpl(this.schema, this.$options, this, executor);
+        if (options.diagnostics) {
+            const diagnosticsSchema = z.object({
+                slowQueryThresholdMs: z.number().nonnegative().optional(),
+                slowQueryMaxRecords: z.int().nonnegative().or(z.literal(Infinity)).optional(),
+            });
+            const parseResult = diagnosticsSchema.safeParse(options.diagnostics);
+            if (!parseResult.success) {
+                throw createConfigError(`Invalid diagnostics configuration: ${formatError(parseResult.error)}`);
+            }
+        }
     }
 
     /**
      * Validates that all computed fields in the schema have corresponding configurations.
      */
-    private validateComputedFieldsConfig() {
+    private validateComputedFieldsConfig(options: ClientOptions<SchemaDef>) {
         const computedFieldsConfig =
-            'computedFields' in this.$options
-                ? (this.$options.computedFields as Record<string, any> | undefined)
-                : undefined;
+            'computedFields' in options ? (options.computedFields as Record<string, any> | undefined) : undefined;
 
         for (const [modelName, modelDef] of Object.entries(this.$schema.models)) {
             if (modelDef.computedFields) {
@@ -183,6 +178,29 @@ export class ClientImpl {
         }
     }
 
+    get $qb() {
+        return this.kysely;
+    }
+
+    get $qbRaw() {
+        return this.kyselyRaw;
+    }
+
+    get $zod() {
+        return this.inputValidator.zodFactory;
+    }
+
+    get isTransaction() {
+        return this.kysely.isTransaction;
+    }
+
+    /**
+     * Create a new client with a new query executor.
+     */
+    withExecutor(executor: QueryExecutor) {
+        return new ClientImpl(this.schema, this.$options, this, executor);
+    }
+
     // overload for interactive transaction
     $transaction<T>(
         callback: (tx: ClientContract<SchemaDef>) => Promise<T>,
@@ -422,6 +440,13 @@ export class ClientImpl {
         return this.$setOptions(newOptions);
     }
 
+    async $diagnostics(): Promise<Diagnostics> {
+        return {
+            zodCache: this.inputValidator.zodFactory.cacheStats,
+            slowQueries: this.slowQueries.map((q) => ({ ...q })),
+        };
+    }
+
     $executeRaw(query: TemplateStringsArray, ...values: any[]) {
         return createZenStackPromise(async () => {
             const result = await sql(query, ...values).execute(this.kysely);

packages/orm/src/client/contract.ts

@@ -39,6 +39,7 @@ import type {
     UpdateManyArgs,
     UpsertArgs,
 } from './crud-types';
+import type { Diagnostics } from './diagnostics';
 import type { ClientOptions, QueryOptions } from './options';
 import type { ExtClientMembersBase, ExtQueryArgsBase, RuntimePlugin } from './plugin';
 import type { ZenStackPromise } from './promise';
@@ -211,6 +212,11 @@ export type ClientContract<
      * @private
      */
     $pushSchema(): Promise<void>;
+
+    /**
+     * Returns diagnostics information such as cache and slow query statistics.
+     */
+    $diagnostics(): Promise<Diagnostics>;
 } & {
     [Key in GetSlicedModels<Schema, Options> as Uncapitalize<Key>]: ModelOperations<Schema, Key, Options, ExtQueryArgs>;
 } & ProcedureOperations<Schema, Options> &

packages/orm/src/client/diagnostics.ts

@@ -0,0 +1,49 @@
+/**
+ * Zod schema cache statistics.
+ */
+export interface ZodCacheStats {
+    /**
+     * Number of cached Zod schemas.
+     */
+    size: number;
+
+    /**
+     * Keys of the cached Zod schemas.
+     */
+    keys: string[];
+}
+
+/**
+ * Information about a query, used for diagnostics.
+ */
+export interface QueryInfo {
+    /**
+     * Time when the query started.
+     */
+    startedAt: Date;
+
+    /**
+     * Duration of the query in milliseconds.
+     */
+    durationMs: number;
+
+    /**
+     * SQL statement of the query.
+     */
+    sql: string;
+}
+
+/**
+ * ZenStackClient diagnostics.
+ */
+export interface Diagnostics {
+    /**
+     * Statistics about the Zod schemas (used for query args validation) cache.
+     */
+    zodCache: ZodCacheStats;
+
+    /**
+     * Slow queries.
+     */
+    slowQueries: QueryInfo[];
+}

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

@@ -31,7 +31,7 @@ import {
 } from 'kysely';
 import { match } from 'ts-pattern';
 import type { ModelDef, SchemaDef, TypeDefDef } from '../../schema';
-import { type ClientImpl } from '../client-impl';
+import type { ClientImpl } from '../client-impl';
 import { TransactionIsolationLevel, type ClientContract } from '../contract';
 import { getCrudDialect } from '../crud/dialects';
 import type { BaseCrudDialect } from '../crud/dialects/base-dialect';
@@ -70,6 +70,8 @@ type CallAfterMutationHooksArgs = {
     afterMutationEntities?: Record<string, unknown>[];
 };
 
+const DEFAULT_MAX_SLOW_RECORDS = 100;
+
 export class ZenStackQueryExecutor extends DefaultQueryExecutor {
     // #region constructor, fields and props
 
@@ -673,8 +675,17 @@ In such cases, ZenStack cannot reliably determine the IDs of the mutated entitie
             compiledQuery = { ...compiledQuery, parameters: parameters };
         }
 
+        const trackSlowQuery = this.options.diagnostics !== undefined;
+        const startTimestamp = trackSlowQuery ? performance.now() : undefined;
+        const startedAt = trackSlowQuery ? new Date() : undefined;
+
         try {
             const result = await connection.executeQuery<any>(compiledQuery);
+
+            if (startTimestamp !== undefined) {
+                this.trackSlowQuery(compiledQuery, startTimestamp, startedAt!);
+            }
+
             return this.ensureProperQueryResult(compiledQuery.query, result);
         } catch (err) {
             throw createDBQueryError(
@@ -686,6 +697,38 @@ In such cases, ZenStack cannot reliably determine the IDs of the mutated entitie
         }
     }
 
+    private trackSlowQuery(compiledQuery: CompiledQuery, startTimestamp: number, startedAt: Date) {
+        const durationMs = performance.now() - startTimestamp;
+        const thresholdMs = this.options.diagnostics?.slowQueryThresholdMs;
+        if (thresholdMs === undefined || durationMs < thresholdMs) {
+            return;
+        }
+
+        const slowQueries = this.client.slowQueries;
+        const maxRecords = this.options.diagnostics?.slowQueryMaxRecords ?? DEFAULT_MAX_SLOW_RECORDS;
+        if (maxRecords <= 0) {
+            return;
+        }
+
+        const queryInfo = { startedAt, durationMs, sql: compiledQuery.sql };
+
+        if (slowQueries.length >= maxRecords) {
+            // find and remove the entry with the lowest duration
+            let minIndex = 0;
+            for (let i = 1; i < slowQueries.length; i++) {
+                if (slowQueries[i]!.durationMs < slowQueries[minIndex]!.durationMs) {
+                    minIndex = i;
+                }
+            }
+            // only replace if the new query is slower than the minimum
+            if (durationMs > slowQueries[minIndex]!.durationMs) {
+                slowQueries[minIndex] = queryInfo;
+            }
+        } else {
+            slowQueries.push(queryInfo);
+        }
+    }
+
     private ensureProperQueryResult(query: RootOperationNode, result: QueryResult<any>) {
         let finalResult = result;
 

packages/orm/src/client/options.ts

@@ -178,7 +178,8 @@ export type ClientOptions<Schema extends SchemaDef> = QueryOptions<Schema> & {
     plugins?: AnyPlugin[];
 
     /**
-     * Logging configuration.
+     * Logging configuration. Extends Kysely's log config with a `'warning'` level
+     * for ZenStack-specific diagnostics (e.g., slow query warnings).
      */
     log?: KyselyConfig['log'];
 
@@ -211,9 +212,25 @@ export type ClientOptions<Schema extends SchemaDef> = QueryOptions<Schema> & {
     useCompactAliasNames?: boolean;
 
     /**
-     * Whether to skip validation for computed fields.
+     * Whether to skip validation for whether all computed fields are properly defined.
      */
     skipValidationForComputedFields?: boolean;
+
+    /**
+     * Diagnostics related options.
+     */
+    diagnostics?: {
+        /**
+         * Threshold in milliseconds for determining slow queries. If not specified, no query will be considered slow.
+         */
+        slowQueryThresholdMs?: number;
+
+        /**
+         * Maximum number of slow query records to keep in memory. Defaults to `100`. When the number is exceeded, the
+         * entry with the lowest duration will be removed. Set to `Infinity` to keep unlimited records.
+         */
+        slowQueryMaxRecords?: number;
+    };
 } & (HasComputedFields<Schema> extends true
         ? {
               /**