beforeQuery and afterQuery from registry for plugins

Author: Brayden

src/handler.ts

@@ -87,6 +87,7 @@ export class StarbaseDB {
             plugins: this.plugins,
         })
         await registry.init()
+        this.dataSource.registry = registry
 
         this.app.post('/query/raw', async (c) =>
             this.queryRoute(c.req.raw, true)

src/index.ts

@@ -177,8 +177,6 @@ export default {
                 }),
             ] satisfies StarbasePlugin[]
 
-            dataSource.plugins = plugins
-
             const starbase = new StarbaseDB({
                 dataSource,
                 config,

src/operation.ts

@@ -57,15 +57,11 @@ async function beforeQuery(opts: {
 }): Promise<{ sql: string; params?: unknown[] }> {
     let { sql, params, dataSource, config } = opts
 
-    if (dataSource?.plugins?.length) {
-        await Promise.all(
-            dataSource.plugins.map(async (plugin: StarbasePlugin) => {
-                const { sql: _sql, params: _params } =
-                    await plugin.beforeQuery(opts)
-                sql = _sql
-                params = _params
-            })
-        )
+    if (dataSource?.registry) {
+        const { sql: _sql, params: _params } =
+            await dataSource?.registry?.beforeQuery(opts)
+        sql = _sql
+        params = _params
     }
 
     return { sql, params }
@@ -81,13 +77,10 @@ async function afterQuery(opts: {
     let { result, isRaw, dataSource } = opts
     result = isRaw ? transformRawResults(result, 'from') : result
 
-    if (dataSource?.plugins?.length) {
-        await Promise.all(
-            dataSource.plugins.map(async (plugin: StarbasePlugin) => {
-                result = await plugin.afterQuery(opts)
-            })
-        )
-    }
+    result = await dataSource?.registry?.afterQuery({
+        ...opts,
+        result,
+    })
 
     return isRaw ? transformRawResults(result, 'to') : result
 }
@@ -100,7 +93,7 @@ function transformRawResults(
         // Convert our result from the `raw` output to a traditional object
         result = {
             ...result,
-            rows: result.rows.map((row: any) =>
+            rows: result?.rows?.map((row: any) =>
                 result.columns.reduce(
                     (obj: any, column: string, index: number) => {
                         obj[column] = row[index]

src/plugin.ts

@@ -69,4 +69,52 @@ export class StarbasePluginRegistry {
             throw e
         }
     }
+
+    public async beforeQuery(opts: {
+        sql: string
+        params?: unknown[]
+        dataSource?: DataSource
+        config?: StarbaseDBConfiguration
+    }): Promise<{ sql: string; params?: unknown[] }> {
+        let { sql, params } = opts
+
+        if (this.plugins?.length) {
+            await Promise.all(
+                this.plugins.map(async (plugin: StarbasePlugin) => {
+                    const { sql: _sql, params: _params } =
+                        await plugin.beforeQuery(opts)
+                    sql = _sql
+                    params = _params
+                })
+            )
+        }
+
+        return {
+            sql,
+            params,
+        }
+    }
+
+    public async afterQuery(opts: {
+        sql: string
+        result: any
+        isRaw: boolean
+        dataSource?: DataSource
+        config?: StarbaseDBConfiguration
+    }): Promise<any> {
+        let { result } = opts
+
+        if (this.plugins?.length) {
+            await Promise.all(
+                this.plugins.map(async (plugin: StarbasePlugin) => {
+                    result = await plugin.afterQuery({
+                        ...opts,
+                        result,
+                    })
+                })
+            )
+        }
+
+        return result
+    }
 }

src/types.ts

@@ -1,5 +1,5 @@
 import { StarbaseDBDurableObject } from './do'
-import { StarbasePlugin } from './plugin'
+import { StarbasePlugin, StarbasePluginRegistry } from './plugin'
 
 export type QueryResult = Record<string, SqlStorageValue>
 
@@ -56,7 +56,7 @@ export type DataSource = {
     context?: Record<string, unknown>
     cache?: boolean
     cacheTTL?: number
-    plugins?: StarbasePlugin[]
+    registry?: StarbasePluginRegistry
 }
 
 export enum RegionLocationHint {