feat: make feature flags a bit more granular

src/handler.test.ts

@@ -49,13 +49,18 @@ vi.mock('./plugin', () => ({
     })),
 }))
 
-vi.mock('./utils', () => ({
-    createResponse: vi.fn((result, error, status) => ({
-        result,
-        error,
-        status,
-    })),
-}))
+vi.mock('./utils', async () => {
+    const { getFeatureFromConfig } = await import('./utils')
+
+    return {
+        createResponse: vi.fn((result, error, status) => ({
+            result,
+            error,
+            status,
+        })),
+        getFeatureFromConfig,
+    }
+})
 
 let instance: StarbaseDB
 let mockDataSource: DataSource

src/handler.ts

@@ -5,7 +5,12 @@ import { validator } from 'hono/validator'
 import { DataSource } from './types'
 import { LiteREST } from './literest'
 import { executeQuery, executeTransaction } from './operation'
-import { createResponse, QueryRequest, QueryTransactionRequest } from './utils'
+import {
+    createResponse,
+    QueryRequest,
+    QueryTransactionRequest,
+    getFeatureFromConfig,
+} from './utils'
 import { dumpDatabaseRoute } from './export/dump'
 import { exportTableToJsonRoute } from './export/json'
 import { exportTableToCsvRoute } from './export/csv'
@@ -26,6 +31,10 @@ export interface StarbaseDBConfiguration {
         websocket?: boolean
         export?: boolean
         import?: boolean
+        studio?: boolean
+        cron?: boolean
+        cdc?: boolean
+        interface?: boolean
     }
 }
 
@@ -283,9 +292,9 @@ export class StarbaseDB {
      */
     private getFeature(
         key: keyof NonNullable<StarbaseDBConfiguration['features']>,
-        defaultValue = true
+        defaultValue?: boolean
     ): boolean {
-        return this.config.features?.[key] ?? !!defaultValue
+        return getFeatureFromConfig(this.config.features)(key, defaultValue)
     }
 
     async queryRoute(request: Request, isRaw: boolean): Promise<Response> {

src/index.ts

@@ -1,4 +1,4 @@
-import { createResponse } from './utils'
+import { createResponse, getFeatureFromConfig } from './utils'
 import { StarbaseDB, StarbaseDBConfiguration } from './handler'
 import { DataSource, RegionLocationHint } from './types'
 import { createRemoteJWKSet, jwtVerify } from 'jose'
@@ -32,6 +32,14 @@ export interface Env {
 
     ENABLE_ALLOWLIST?: boolean
     ENABLE_RLS?: boolean
+    ENABLE_REST?: boolean
+    ENABLE_WEBSOCKET?: boolean
+    ENABLE_EXPORT?: boolean
+    ENABLE_IMPORT?: boolean
+    ENABLE_CRON?: boolean
+    ENABLE_CDC?: boolean
+    ENABLE_INTERFACE?: boolean
+    ENABLE_STUDIO?: boolean
 
     // External database source details
     OUTERBASE_API_KEY?: string
@@ -172,43 +180,63 @@ export default {
                 features: {
                     allowlist: env.ENABLE_ALLOWLIST,
                     rls: env.ENABLE_RLS,
+                    rest: env.ENABLE_REST,
+                    websocket: env.ENABLE_WEBSOCKET,
+                    export: env.ENABLE_EXPORT,
+                    import: env.ENABLE_IMPORT,
+                    cron: env.ENABLE_CRON,
+                    cdc: env.ENABLE_CDC,
+                    interface: env.ENABLE_INTERFACE,
+                    studio: env.ENABLE_STUDIO,
                 },
             }
 
-            const webSocketPlugin = new WebSocketPlugin()
-            const cronPlugin = new CronPlugin()
-            const cdcPlugin = new ChangeDataCapturePlugin({
+            const getFeature = getFeatureFromConfig(config.features)
+
+            /**
+             * Plugins
+             */
+            const webSocketPlugin = getFeature('websocket') ? new WebSocketPlugin() : undefined
+            const studioPlugin = getFeature('studio') ? new StudioPlugin({
+                username: env.STUDIO_USER,
+                password: env.STUDIO_PASS,
+                apiKey: env.ADMIN_AUTHORIZATION_TOKEN,
+            }) : undefined
+            const sqlMacrosPlugin = new SqlMacrosPlugin({
+                preventSelectStar: false,
+            })
+            const queryLogPlugin = new QueryLogPlugin({ ctx })
+            const cdcPlugin = getFeature('cdc') ? new ChangeDataCapturePlugin({
                 stub,
                 broadcastAllEvents: false,
                 events: [],
-            })
-
-            cdcPlugin.onEvent(async ({ action, schema, table, data }) => {
-                // Include change data capture code here
-            }, ctx)
-
-            cronPlugin.onEvent(async ({ name, cron_tab, payload }) => {
-                // Include cron event code here
-            }, ctx)
-
-            const interfacePlugin = new InterfacePlugin()
-
-            const plugins = [
+            }) : undefined
+            const cronPlugin = getFeature('cron') ? new CronPlugin() : undefined
+            const statsPlugin = new StatsPlugin()
+            const interfacePlugin = getFeature('interface') ? new InterfacePlugin() : undefined
+
+            const plugins: StarbasePlugin[] = [
                 webSocketPlugin,
-                new StudioPlugin({
-                    username: env.STUDIO_USER,
-                    password: env.STUDIO_PASS,
-                    apiKey: env.ADMIN_AUTHORIZATION_TOKEN,
-                }),
-                new SqlMacrosPlugin({
-                    preventSelectStar: false,
-                }),
-                new QueryLogPlugin({ ctx }),
+                studioPlugin,
+                sqlMacrosPlugin,
+                queryLogPlugin,
                 cdcPlugin,
                 cronPlugin,
-                new StatsPlugin(),
+                statsPlugin,
                 interfacePlugin,
-            ] satisfies StarbasePlugin[]
+            ].filter(plugin => !!plugin)
+
+            if (getFeature('cdc')) {
+                cdcPlugin?.onEvent(async ({ action, schema, table, data }) => {
+                    // Include change data capture code here
+                }, ctx)
+            }
+
+            if (getFeature('cron')) {
+                cronPlugin?.onEvent(async ({ name, cron_tab, payload }) => {
+                    // Include cron event code here
+                }, ctx)
+            }
 
             const starbase = new StarbaseDB({
                 dataSource,
@@ -228,7 +256,10 @@ export default {
             // next authentication checks happen. If a page is meant to have any
             // sort of authentication, it can provide Basic Auth itself or expose
             // itself in another plugin.
-            if (interfacePlugin.matchesRoute(url.pathname)) {
+            if (
+                getFeature('interface') &&
+                interfacePlugin?.matchesRoute(url.pathname)
+            ) {
                 return await starbase.handle(request, ctx)
             }
 

src/operation.ts

@@ -24,7 +24,7 @@ import { isQueryAllowed } from './allowlist'
 import { applyRLS } from './rls'
 import type { SqlConnection } from '@outerbase/sdk/dist/connections/sql-base'
 import { StarbasePlugin } from './plugin'
-
+import { getFeatureFromConfig } from './utils'
 export type OperationQueueItem = {
     queries: { sql: string; params?: any[] }[]
     isTransaction: boolean
@@ -204,18 +204,20 @@ export async function executeQuery(opts: {
         return []
     }
 
+    const getFeature = getFeatureFromConfig(config.features)
+
     // If the allowlist feature is enabled, we should verify the query is allowed before proceeding.
     await isQueryAllowed({
         sql: sql,
-        isEnabled: config?.features?.allowlist ?? false,
+        isEnabled: getFeature('allowlist', false),
         dataSource,
         config,
     })
 
     // If the row level security feature is enabled, we should apply our policies to this SQL statement.
     sql = await applyRLS({
         sql,
-        isEnabled: config?.features?.rls ?? true,
+        isEnabled: getFeature('rls', true),
         dataSource,
         config,
     })

src/utils.ts

@@ -1,4 +1,5 @@
 import { corsHeaders } from './cors'
+import { StarbaseDBConfiguration } from './handler'
 
 export type QueryTransactionRequest = {
     transaction?: QueryRequest[]
@@ -22,3 +23,14 @@ export function createResponse(
         },
     })
 }
+
+export function getFeatureFromConfig(
+    features: StarbaseDBConfiguration['features']
+) {
+    return function getFeature(
+        key: keyof NonNullable<StarbaseDBConfiguration['features']>,
+        defaultValue = true
+    ): boolean {
+        return features?.[key] ?? !!defaultValue
+    }
+}

worker-configuration.d.ts

@@ -8,6 +8,14 @@ interface Env {
     STUDIO_PASS: '123456'
     ENABLE_ALLOWLIST: 0
     ENABLE_RLS: 0
+    ENABLE_CRON: 0
+    ENABLE_CDC: 0
+    ENABLE_INTERFACE: 0
+    ENABLE_STUDIO: 0
+    ENABLE_REST: 1
+    ENABLE_WEBSOCKET: 1
+    ENABLE_EXPORT: 1
+    ENABLE_IMPORT: 1
     AUTH_ALGORITHM: 'RS256'
     AUTH_JWKS_ENDPOINT: ''
     DATABASE_DURABLE_OBJECT: DurableObjectNamespace<

wrangler.toml

@@ -51,6 +51,14 @@ REGION = "auto"
 # Toggle to enable default features 
 ENABLE_ALLOWLIST = 0
 ENABLE_RLS = 0
+ENABLE_CRON = 1
+ENABLE_CDC = 1
+ENABLE_INTERFACE = 1
+ENABLE_STUDIO = 1
+ENABLE_REST = 1
+ENABLE_WEBSOCKET = 1
+ENABLE_EXPORT = 1
+ENABLE_IMPORT = 1
 
 # External database source details
 # This enables Starbase to connect to an external data source