feat(driver-adapters-executor): add types

prisma · Jan 23, 2024 · 94e4600 · 94e4600
1 parent 0e77726
commit 94e4600
Showing 1 changed file with 25 additions and 19 deletions.
diff --git a/query-engine/driver-adapters/executor/src/recording.ts b/query-engine/driver-adapters/executor/src/recording.ts
@@ -4,7 +4,8 @@ import {
   type Result,
   type ResultSet,
 } from "@prisma/driver-adapter-utils";
-import { RetryHandler } from "undici";
+
+type Recordings = ReturnType<typeof createInMemoryRecordings>;
 
 export function recording(adapter: DriverAdapter) {
   const recordings = createInMemoryRecordings();
@@ -15,7 +16,7 @@ export function recording(adapter: DriverAdapter) {
   };
 }
 
-function recorder(adapter: DriverAdapter, recordings) {
+function recorder(adapter: DriverAdapter, recordings: Recordings) {
   return {
     provider: adapter.provider,
     startTransaction: () => {
@@ -25,20 +26,19 @@ function recorder(adapter: DriverAdapter, recordings) {
       return adapter.getConnectionInfo!();
     },
     queryRaw: async (params) => {
-      const result = adapter.queryRaw(params);
+      const result = await adapter.queryRaw(params);
       recordings.addQueryResults(params, result);
       return result;
     },
-
     executeRaw: async (params) => {
-      const result = adapter.executeRaw(params);
+      const result = await adapter.executeRaw(params);
       recordings.addCommandResults(params, result);
       return result;
     },
   };
 }
 
-function replayer(adapter: DriverAdapter, recordings) {
+function replayer(adapter: DriverAdapter, recordings: Recordings) {
   return {
     provider: adapter.provider,
     recordings: recordings,
@@ -58,36 +58,42 @@ function replayer(adapter: DriverAdapter, recordings) {
 }
 
 function createInMemoryRecordings() {
-  const queryResults = {};
-  const commandResults = {};
+  const queryResults: Map<string, Result<ResultSet>> = new Map();
+  const commandResults: Map<string, Result<number>> = new Map();
 
-  const queryToKey = (query) => JSON.stringify(query);
+  const queryToKey = (params: Query) => {
+    return JSON.stringify(params);
+  };
 
   return {
-    addQueryResults: (params, result) => {
+    addQueryResults: (params: Query, result: Result<ResultSet>) => {
       const key = queryToKey(params);
-      queryResults[key] = result;
+      queryResults.set(key, result);
     },
 
-    getQueryResults: (params) => {
+    getQueryResults: (params: Query) => {
       const key = queryToKey(params);
-      if (!(key in queryResults)) {
+
+      if (!queryResults.has(key)) {
         throw new Error(`Query not recorded: ${key}`);
       }
-      return queryResults[key];
+
+      return queryResults.get(key)!;
     },
 
-    addCommandResults: (params, result) => {
+    addCommandResults: (params: Query, result: Result<number>) => {
       const key = queryToKey(params);
-      commandResults[key] = result;
+      commandResults.set(key, result);
     },
 
-    getCommandResults: (params) => {
+    getCommandResults: (params: Query) => {
       const key = queryToKey(params);
-      if (!(key in commandResults)) {
+
+      if (!commandResults.has(key)) {
         throw new Error(`Command not recorded: ${key}`);
       }
-      return commandResults[key];
+
+      return commandResults.get(key)!;
     },
   };
 }