Merge pull request #80 from outerbase/bwilmoth/internal-stats

feat: Statistics and information about instance

Author: Brayden

plugins/stats/index.ts

@@ -0,0 +1,45 @@
+import { DataSource } from '../../dist'
+import { StarbaseApp, StarbaseDBConfiguration } from '../../src/handler'
+import { StarbasePlugin } from '../../src/plugin'
+
+export class StatsPlugin extends StarbasePlugin {
+    // Prefix route
+    prefix: string = '/_internal/stats'
+    // Configuration details about the request and user
+    private config?: StarbaseDBConfiguration
+    // Data source to run internal RPC queries
+    dataSource?: DataSource
+
+    constructor() {
+        super('starbasedb:stats', {
+            requiresAuth: true,
+        })
+    }
+
+    override async register(app: StarbaseApp) {
+        app.use(async (c, next) => {
+            this.config = c?.get('config')
+            this.dataSource = c?.get('dataSource')
+            await next()
+        })
+
+        app.get(this.prefix, async (c, next) => {
+            // Only admin authorized users are permitted to subscribe to CDC events.
+            if (this.config?.role !== 'admin') {
+                return new Response('Unauthorized request', { status: 400 })
+            }
+
+            // Get stats from internal source
+            const stats = await this.dataSource?.rpc.getStatistics()
+            const additionalStats = {
+                ...stats,
+                plugins: this.dataSource?.registry?.currentPlugins(),
+            }
+            return new Response(JSON.stringify(additionalStats), {
+                headers: {
+                    'Content-Type': 'application/json',
+                },
+            })
+        })
+    }
+}

src/do.ts

@@ -63,10 +63,35 @@ export class StarbaseDBDurableObject extends DurableObject {
 
     init() {
         return {
+            getStatistics: this.getStatistics.bind(this),
             executeQuery: this.executeQuery.bind(this),
         }
     }
 
+    public async getStatistics(): Promise<{
+        databaseSize: number
+        activeConnections: number
+        recentQueries: number
+    }> {
+        const sql = `SELECT COUNT(*) as count 
+            FROM tmp_query_log 
+            WHERE created_at >= datetime('now', '-24 hours')`
+        const result = (await this.executeQuery({
+            sql,
+            isRaw: false,
+        })) as Record<string, SqlStorageValue>[]
+        const row = result.length ? result[0] : { count: 0 }
+
+        return {
+            // Size in bytes
+            databaseSize: this.sql.databaseSize,
+            // Count of persistent web socket connections
+            activeConnections: this.connections.keys.length,
+            // Assuming the `QueryLogPlugin` is in use, count is of the last 24 hours
+            recentQueries: Number(row.count),
+        }
+    }
+
     async fetch(request: Request) {
         const url = new URL(request.url)
 

src/index.ts

@@ -9,6 +9,7 @@ import { StudioPlugin } from '../plugins/studio'
 import { SqlMacrosPlugin } from '../plugins/sql-macros'
 import { ChangeDataCapturePlugin } from '../plugins/cdc'
 import { QueryLogPlugin } from '../plugins/query-log'
+import { StatsPlugin } from '../plugins/stats'
 
 export { StarbaseDBDurableObject } from './do'
 
@@ -190,6 +191,7 @@ export default {
                 }),
                 new QueryLogPlugin({ ctx }),
                 cdcPlugin,
+                new StatsPlugin(),
             ] satisfies StarbasePlugin[]
 
             const starbase = new StarbaseDB({

src/plugin.ts

@@ -70,6 +70,10 @@ export class StarbasePluginRegistry {
         }
     }
 
+    public currentPlugins(): string[] {
+        return this.plugins.map((plugin) => plugin.name)
+    }
+
     public async beforeQuery(opts: {
         sql: string
         params?: unknown[]

worker-configuration.d.ts

@@ -8,7 +8,6 @@ interface Env {
     STUDIO_PASS: '123456'
     ENABLE_ALLOWLIST: 0
     ENABLE_RLS: 0
-    EXTERNAL_DB_TYPE: 'postgresql'
     AUTH_ALGORITHM: 'RS256'
     AUTH_JWKS_ENDPOINT: ''
     DATABASE_DURABLE_OBJECT: DurableObjectNamespace<

wrangler.toml

@@ -49,7 +49,7 @@ ENABLE_RLS = 0
 # External database source details
 # This enables Starbase to connect to an external data source
 # OUTERBASE_API_KEY = ""
-EXTERNAL_DB_TYPE = "postgresql"
+# EXTERNAL_DB_TYPE = "postgresql"
 # EXTERNAL_DB_HOST = ""
 # EXTERNAL_DB_PORT = 0
 # EXTERNAL_DB_USER = ""