feat(cbjs): introduce hook onQueryStart and onQueryEnd

Author: JesusTheHun

packages/cbjs/src/cluster.ts

@@ -43,6 +43,7 @@ import {
   PingOptions,
   PingResult,
 } from './diagnosticstypes.js';
+import { CouchbaseError } from './errors.js';
 import { EventingFunctionManager } from './eventingfunctionmanager.js';
 import { QueryExecutor } from './queryexecutor.js';
 import { QueryIndexManager } from './queryindexmanager.js';
@@ -159,13 +160,39 @@ export type DnsConfig = {
   dnsSrvTimeout?: number;
 };
 
+export type Hooks<T extends CouchbaseClusterTypes> = {
+  /**
+   * Any error thrown in a hook will be swallowed and this function will be called
+   * @param err The error thrown by the hook function.
+   */
+  onHookError: (err: unknown) => void;
+  onQueryStart: (opts: {
+    statement: string;
+    options: QueryOptions<T, boolean>;
+  }) => unknown;
+  onQueryEnd: (
+    opts:
+      | {
+          statement: string;
+          options: QueryOptions<T, boolean>;
+          result: QueryResult<any, boolean>;
+        }
+      | {
+          statement: string;
+          options: QueryOptions<T, boolean>;
+          error: Error;
+        },
+    returnOfOnQueryStart: unknown
+  ) => void;
+};
+
 /**
  * Specifies the options which can be specified when connecting
  * to a cluster.
  *
  * @category Core
  */
-export type ConnectOptions = {
+export type ConnectOptions<T extends CouchbaseClusterTypes = any> = {
   /**
    * Specifies a username to use for an implicitly created IPasswordAuthenticator
    * used for authentication with the cluster.
@@ -207,6 +234,11 @@ export type ConnectOptions = {
    */
   queryResultParser?: (value: string) => any;
 
+  /**
+   * Specifies hook functions that allow you to perform actions before and after some Couchbase operations.
+   */
+  hooks?: Partial<Hooks<T>>;
+
   /**
    * Specifies the options for transactions.
    */
@@ -253,6 +285,7 @@ export class Cluster<in out T extends CouchbaseClusterTypes = DefaultClusterType
   private _conn: CppConnection;
   private _transcoder: Transcoder;
   private _queryResultParser: (value: string) => any;
+  private _hooks?: Partial<Hooks<T>>;
   private _txnConfig: TransactionsConfig;
   private _transactions?: Transactions<T>;
   private readonly _openBuckets: Map<BucketName<T>, Promise<void>>;
@@ -279,6 +312,13 @@ export class Cluster<in out T extends CouchbaseClusterTypes = DefaultClusterType
     return this._queryResultParser;
   }
 
+  /**
+  @internal
+  */
+  get hooks(): Partial<Hooks<T>> | undefined {
+    return this._hooks;
+  }
+
   /**
   @internal
   */
@@ -353,7 +393,7 @@ export class Cluster<in out T extends CouchbaseClusterTypes = DefaultClusterType
   @internal
   @deprecated Use the static sdk-level {@link connect} method instead.
   */
-  constructor(connStr: string, options?: ConnectOptions) {
+  constructor(connStr: string, options?: ConnectOptions<T>) {
     if (!options) {
       options = {};
     }
@@ -369,7 +409,7 @@ export class Cluster<in out T extends CouchbaseClusterTypes = DefaultClusterType
     this._trustStorePath = options.security.trustStorePath ?? '';
 
     if (options.configProfile) {
-      connectionProfiles.applyProfile(options.configProfile, options);
+      connectionProfiles.applyProfile(options.configProfile, options as never);
     }
     this._kvTimeout = options.timeouts.kvTimeout ?? 2500;
     this._kvDurableTimeout = options.timeouts.kvDurableTimeout ?? 10000;
@@ -384,6 +424,7 @@ export class Cluster<in out T extends CouchbaseClusterTypes = DefaultClusterType
 
     this._transcoder = options.transcoder ?? new DefaultTranscoder();
     this._queryResultParser = options.queryResultParser ?? JSON.parse;
+    this._hooks = options.hooks;
 
     if (options.transactions) {
       this._txnConfig = options.transactions;
@@ -434,7 +475,7 @@ export class Cluster<in out T extends CouchbaseClusterTypes = DefaultClusterType
   */
   static async connect<T extends CouchbaseClusterTypes = DefaultClusterTypes>(
     connStr: string,
-    options?: ConnectOptions,
+    options?: ConnectOptions<T>,
     callback?: NodeCallback<Cluster<T>>
   ): Promise<Cluster<T>> {
     return await PromiseHelper.wrapAsync(async () => {

packages/cbjs/src/couchbase.ts

@@ -32,7 +32,7 @@ import { NodeCallback } from './utilities.js';
  */
 export async function connect<T extends CouchbaseClusterTypes = DefaultClusterTypes>(
   connStr: `couchbase://` | (string & NonNullable<unknown>),
-  options?: ConnectOptions,
+  options?: ConnectOptions<T>,
   callback?: NodeCallback<Cluster<T>>
 ): Promise<Cluster<T>> {
   return await Cluster.connect<T>(connStr, options, callback);

packages/cbjs/src/queryexecutor.ts

@@ -147,7 +147,18 @@ export class QueryExecutor<T extends CouchbaseClusterTypes = CouchbaseClusterTyp
     const timeout = options.timeout ?? this._cluster.queryTimeout;
     const rowParser = options.queryResultParser ?? this._cluster.queryResultParser;
 
-    return QueryExecutor.execute((callback) => {
+    let hookReturnValue: unknown;
+
+    try {
+      hookReturnValue = this._cluster.hooks?.onQueryStart?.({
+        statement: query,
+        options,
+      });
+    } catch (err) {
+      this._cluster.hooks?.onHookError?.(err);
+    }
+
+    const result = QueryExecutor.execute((callback) => {
       this._cluster.conn.query(
         {
           statement: query,
@@ -192,5 +203,39 @@ export class QueryExecutor<T extends CouchbaseClusterTypes = CouchbaseClusterTyp
         callback
       );
     }, rowParser);
+
+    if (this._cluster.hooks?.onQueryEnd) {
+      void result
+        .then((r) => {
+          try {
+            this._cluster.hooks?.onQueryEnd?.(
+              {
+                statement: query,
+                options,
+                result: r,
+              },
+              hookReturnValue
+            );
+          } catch (err) {
+            this._cluster.hooks?.onHookError?.(err);
+          }
+        })
+        .catch((err) => {
+          try {
+            this._cluster.hooks?.onQueryEnd?.(
+              {
+                statement: query,
+                options,
+                error: err,
+              },
+              hookReturnValue
+            );
+          } catch (err) {
+            this._cluster.hooks?.onHookError?.(err);
+          }
+        });
+    }
+
+    return result;
   }
 }

tests/cbjs/tests/query.request.spec.ts

@@ -15,11 +15,11 @@
  * limitations under the License.
  */
 import JSONBigint from 'json-bigint';
-import { describe } from 'vitest';
+import { describe, vi } from 'vitest';
 
 import { QueryMetaData, QueryProfileMode, QueryResult, QueryStatus } from '@cbjsdev/cbjs';
 import { ServerFeatures } from '@cbjsdev/http-client';
-import { getConnectionParams, invariant, keyspacePath } from '@cbjsdev/shared';
+import { getConnectionParams, keyspacePath } from '@cbjsdev/shared';
 import { createCouchbaseTest, getDefaultServerTestContext } from '@cbjsdev/vitest';
 
 import { useSampleData } from '../fixtures/useSampleData.js';
@@ -239,6 +239,113 @@ describe.runIf(serverSupportsFeatures(ServerFeatures.Query))(
       }
     );
 
+    describe('hooks', () => {
+      test('should call onQueryStart and onQueryEnd', async ({
+        expect,
+        serverTestContext,
+        useDocumentKey,
+      }) => {
+        const mockOnQueryStart = vi.fn().mockReturnValue('mockStartReturn');
+        const mockOnQueryEnd = vi.fn();
+
+        const params = getConnectionParams();
+        const newConnection = await serverTestContext.newConnection(params, {
+          hooks: {
+            onQueryStart: mockOnQueryStart,
+            onQueryEnd: mockOnQueryEnd,
+          },
+        });
+
+        const collection = newConnection
+          .bucket(serverTestContext.bucket.name)
+          .scope(serverTestContext.scope.name)
+          .collection(serverTestContext.collection.name);
+
+        const docKey = useDocumentKey();
+        await collection.insert(docKey, {
+          title: 'cbjs',
+        });
+
+        const statement = `SELECT META().cas AS cas FROM ${serverTestContext.getKeyspacePath()} USE KEYS $1`;
+        const options = {
+          parameters: [docKey],
+        };
+        const queryResult = await newConnection.query<{ cas: bigint }>(
+          statement,
+          options
+        );
+
+        expect(mockOnQueryStart).toHaveBeenCalledWith({
+          statement,
+          options,
+        });
+
+        expect(mockOnQueryEnd).toHaveBeenCalledWith(
+          {
+            statement,
+            options,
+            result: queryResult,
+          },
+          'mockStartReturn'
+        );
+      });
+
+      test('should call onQueryEnd with the error when the query throws', async ({
+        expect,
+        serverTestContext,
+      }) => {
+        expect.hasAssertions();
+        const mockOnQueryEnd = vi.fn();
+
+        const params = getConnectionParams();
+        const newConnection = await serverTestContext.newConnection(params, {
+          hooks: {
+            onQueryEnd: mockOnQueryEnd,
+          },
+        });
+
+        const statement = `SELECT INVALID QUERY`;
+
+        try {
+          await newConnection.query(statement);
+        } catch (err) {
+          expect(mockOnQueryEnd).toHaveBeenCalledWith(
+            {
+              statement,
+              options: {},
+              error: err,
+            },
+            undefined
+          );
+        }
+      });
+
+      test('should call onHookError when either onQueryStart or onQueryEnd throws', async ({
+        expect,
+        serverTestContext,
+      }) => {
+        const mockOnHookError = vi.fn();
+
+        const params = getConnectionParams();
+        const newConnection = await serverTestContext.newConnection(params, {
+          hooks: {
+            onHookError: mockOnHookError,
+            onQueryStart: () => {
+              throw new Error('onQueryStart failed');
+            },
+            onQueryEnd: () => {
+              throw new Error('onQueryEnd failed');
+            },
+          },
+        });
+
+        const statement = `SELECT 1 as foo`;
+        await newConnection.query(statement);
+
+        expect(mockOnHookError).toHaveBeenCalledTimes(2);
+      });
+    });
+
     describe('row parser', () => {
       test('should use the cluster custom row parser for cluster level query', async ({
         expect,