From 2b27416af01272cb1f49398cf8976ebfb88e8eca Mon Sep 17 00:00:00 2001
From: Andres Aya <andresayac@gmail.com>
Date: Thu, 20 Apr 2023 08:14:11 -0500
Subject: [PATCH] refactor: store w transaction

---
 src/Types/Auth.ts       |  4 ++--
 src/Utils/auth-utils.ts | 47 ++++++++++++++++++-----------------------
 2 files changed, 23 insertions(+), 28 deletions(-)

diff --git a/src/Types/Auth.ts b/src/Types/Auth.ts
index 72bde7c..6c06333 100644
--- a/src/Types/Auth.ts
+++ b/src/Types/Auth.ts
@@ -82,7 +82,7 @@ export type SignalKeyStore = {
 
 export type SignalKeyStoreWithTransaction = SignalKeyStore & {
     isInTransaction: () => boolean
-    transaction(exec: () => Promise<void>): Promise<void>
+    transaction<T>(exec: () => Promise<T>): Promise<T>
 }
 
 export type TransactionCapabilityOptions = {
@@ -92,7 +92,7 @@ export type TransactionCapabilityOptions = {
 
 export type SignalAuthState = {
     creds: SignalCreds
-    keys: SignalKeyStore
+    keys: SignalKeyStore | SignalKeyStoreWithTransaction
 }
 
 export type AuthenticationState = {
diff --git a/src/Utils/auth-utils.ts b/src/Utils/auth-utils.ts
index a4f2e73..b8198f9 100644
--- a/src/Utils/auth-utils.ts
+++ b/src/Utils/auth-utils.ts
@@ -93,32 +93,25 @@ export const addTransactionCapability = (
 	let transactionCache: SignalDataSet = { }
 	let mutations: SignalDataSet = { }
 
-	/**
-	 * prefetches some data and stores in memory,
-	 * useful if these data points will be used together often
-	 * */
-	const prefetch = async<T extends keyof SignalDataTypeMap>(type: T, ids: string[]) => {
-		const dict = transactionCache[type]
-		const idsRequiringFetch = dict
-			? ids.filter(item => typeof dict[item] !== 'undefined')
-			: ids
-		// only fetch if there are any items to fetch
-		if(idsRequiringFetch.length) {
-			dbQueriesInTransaction += 1
-			const result = await state.get(type, idsRequiringFetch)
-
-			transactionCache[type] ||= {}
-			transactionCache[type] = Object.assign(
-				transactionCache[type]!,
-				result
-			)
-		}
-	}
-
 	return {
 		get: async(type, ids) => {
 			if(inTransaction) {
-				await prefetch(type, ids)
+				const dict = transactionCache[type]
+				const idsRequiringFetch = dict
+					? ids.filter(item => typeof dict[item] !== 'undefined')
+					: ids
+				// only fetch if there are any items to fetch
+				if(idsRequiringFetch.length) {
+					dbQueriesInTransaction += 1
+					const result = await state.get(type, idsRequiringFetch)
+
+					transactionCache[type] ||= {}
+					Object.assign(
+						transactionCache[type]!,
+						result
+					)
+				}
+
 				return ids.reduce(
 					(dict, id) => {
 						const value = transactionCache[type]?.[id]
@@ -148,16 +141,17 @@ export const addTransactionCapability = (
 			}
 		},
 		isInTransaction: () => inTransaction,
-		transaction: async(work) => {
+		async transaction(work) {
+			let result: Awaited<ReturnType<typeof work>>
 			// if we're already in a transaction,
 			// just execute what needs to be executed -- no commit required
 			if(inTransaction) {
-				await work()
+				result = await work()
 			} else {
 				logger.trace('entering transaction')
 				inTransaction = true
 				try {
-					await work()
+					result = await work()
 					if(Object.keys(mutations).length) {
 						logger.trace('committing transaction')
 						// retry mechanism to ensure we've some recovery
@@ -184,6 +178,7 @@ export const addTransactionCapability = (
 					dbQueriesInTransaction = 0
 				}
 			}
+			return result
 		}
 	}
 }