From 3b5468499a6ccae4579d16069dde01140f459fe5 Mon Sep 17 00:00:00 2001
From: FoxxMD <FoxxMD@users.noreply.github.com>
Date: Fri, 24 Apr 2026 18:21:09 +0000
Subject: [PATCH] feat(database): Use asyncstorage for query logging

---
 .../common/database/drizzle/drizzleUtils.ts   | 28 +++-----
 .../common/database/drizzle/logContext.ts     | 71 +++++++++++++++++++
 2 files changed, 79 insertions(+), 20 deletions(-)
 create mode 100644 src/backend/common/database/drizzle/logContext.ts

diff --git a/src/backend/common/database/drizzle/drizzleUtils.ts b/src/backend/common/database/drizzle/drizzleUtils.ts
index 54906dca..f8975a8b 100644
--- a/src/backend/common/database/drizzle/drizzleUtils.ts
+++ b/src/backend/common/database/drizzle/drizzleUtils.ts
@@ -10,6 +10,7 @@ import { childLogger, Logger, LogLevel } from '@foxxmd/logging';
 import { loggerNoop } from '../../MaybeLogger.js';
 import { projectDir } from '../../index.js';
 import { relations } from './schema/schema.js';
+import { addToContext, executeQuery } from './logContext.js';
 
 export async function shouldBackupDb(dbPath: string, opts: {logger?: Logger, migrationsFolder?: string} = {}): Promise<[boolean, string[]]> {
   const {
@@ -70,10 +71,9 @@ export async function shouldBackupDb(dbPath: string, opts: {logger?: Logger, mig
 export const getDb = (dbName: string = 'ms', opts: { logger?: Logger, workingDirectory?: string } = {}) => {
   const {
     workingDirectory,
-    logger = loggerNoop
+    logger = loggerNoop,
   } = opts;
   const dbPath = getDbPath(dbName, workingDirectory);
-  logger.info(`Using database at ${dbPath}`);
   return drizzle(dbPath, {relations: relations, logger: createDrizzleLogger(logger)});
 }
 
@@ -87,7 +87,8 @@ export const migrateDb = async (db: ReturnType<typeof drizzle>, opts: {logger?:
   const logger = childLogger(parentLogger, 'Migrations');
 
   try {
-    await migrate(db, { migrationsFolder: migrationsFolder ?? path.resolve(projectDir, 'src/backend/common/database/drizzle/migrations') });
+    logger.info('Starting migrations...');
+    await executeQuery('migrations', async () => migrate(db, { migrationsFolder: migrationsFolder ?? path.resolve(projectDir, 'src/backend/common/database/drizzle/migrations') }), logger, process.env.LOG_MIGRATION === 'true' ? true : 'error');
     logger.info('Migrations complete');
   } catch (e) {
     throw new Error('Failed to migrate database', { cause: e });
@@ -105,24 +106,11 @@ export const performDbMigrationWithBackup = async (dbName: string = 'ms', opts:
   await migrateDb(db, opts);
 }
 
-export const createDrizzleLogger = (parentLogger: Logger, opts: {level?: LogLevel, query?: boolean} = {}): LogWriter & DrizzleLogger => {
-  const {
-    level = 'trace',
-    query = false,
-  } = opts;
-
-  const logger = childLogger(parentLogger, 'Drizzle');
-
-  let queryFunc: (query: string, params: unknown[]) => void = (_, __) => {};
-  if(query) {
-    queryFunc = (query: string, params: unknown[]) => logger[level]({params}, `SQL Query => ${query}`);
-  }
-
+export const createDrizzleLogger = (parentLogger: Logger, opts: {level?: LogLevel} = {}): DrizzleLogger => {
   return {
-    write(message: string) {
-      logger[level](message);
-    },
-    logQuery: queryFunc
+    logQuery: (query: string, params: unknown[]) => {
+      addToContext({sql: query, params})
+    }
   }
 }
 
diff --git a/src/backend/common/database/drizzle/logContext.ts b/src/backend/common/database/drizzle/logContext.ts
new file mode 100644
index 00000000..3901f6c0
--- /dev/null
+++ b/src/backend/common/database/drizzle/logContext.ts
@@ -0,0 +1,71 @@
+import { Logger } from '@foxxmd/logging'
+import { AsyncLocalStorage } from 'async_hooks'
+
+// based on https://numeric.substack.com/p/upgrading-drizzleorm-logging-with
+interface QueryContext {
+    queryKey: string
+    startTime: number
+    queries: { sql?: string, params?: unknown[] }[]
+}
+
+const queryStorage = new AsyncLocalStorage<QueryContext>()
+
+function wrapQuery<T>(queryKey: string, fn: () => Promise<T>): Promise<T> {
+    return queryStorage.run(
+        {
+            queryKey,
+            startTime: Date.now(),
+            queries: []
+        },
+        fn
+    )
+}
+
+function getContext(): QueryContext | undefined {
+    return queryStorage.getStore()
+}
+
+export function addToContext(data: { sql?: string, params?: unknown[] }): void {
+    const context = getContext()
+    if (context) {
+        context.queries.push(data);
+    }
+}
+
+/**
+ * Log all queries made by drizzle during the execution of a promise
+ * 
+ * use second parameter to configure when logging occurs
+ * * true    => log everything (default)
+ * * false   => log nothing, skips asyncstorage entirely
+ * * 'error' => only log if promise throws
+ * 
+ */
+export async function executeQuery<T = any>(queryKey: string, queryPromise: () => Promise<T>, logger: Logger, when: boolean | 'error' = true) {
+    if(when === false) {
+        try {
+            return await queryPromise();
+        } catch (e) {
+            throw e;
+        }
+    }
+    return wrapQuery(queryKey, async () => {
+        try {
+            const results = await queryPromise()
+
+            if (when !== 'error') {
+                // Query is done - grab everything from context
+                const context = getContext()
+                const executionTime = context ? Date.now() - context.startTime : 0
+                logger.info({ labels: ['DB Query', queryKey], queries: context?.queries }, `Execution Complete in ${executionTime}ms`);
+            }
+
+            return results
+        } catch (error) {
+            const context = getContext()
+            const executionTime = context ? Date.now() - context.startTime : 0;
+            logger.warn({ labels: ['DB Query', queryKey], queries: context?.queries }, `Execution failed in ${executionTime}ms`);
+            throw error
+        }
+    })
+}
\ No newline at end of file
-- 
2.51.2