feat: async tenant migrations

Author: fenos

.gitignore

@@ -8,4 +8,5 @@ static/api.json
 data/
 bin/
 coverage/
-.idea/
+.idea/
+migrations/tenants-migration-hash.txt

docker-compose-infra.yml

@@ -48,7 +48,7 @@ services:
       PGBOUNCER_STATS_USERS: postgres
 
   supavisor:
-    image: supabase/supavisor:1.1.6
+    image: supabase/supavisor:1.1.23
     depends_on:
       multitenant_db:
         condition: service_healthy
@@ -78,7 +78,7 @@ services:
     command: sh -c "/app/bin/migrate && /app/bin/server"
 
   supavisor_setup:
-    image: supabase/supavisor:1.1.6
+    image: supabase/supavisor:1.1.23
     command: |
       curl -X PUT \
         "http://supavisor:4000/api/tenants/bjhaohmqunupljrqypxz" \

migrations/multitenant/0006-tenants-add-created-at-migrations-version.sql

@@ -0,0 +1,5 @@
+ALTER TABLE tenants ADD COLUMN IF NOT EXISTS cursor_id SERIAL;
+ALTER TABLE tenants ADD COLUMN IF NOT EXISTS created_at TIMESTAMP DEFAULT current_timestamp;
+ALTER TABLE tenants ADD COLUMN IF NOT EXISTS migrations_version text null DEFAULT null;
+
+create index if not exists tenants_migration_version_idx on tenants(cursor_id, migrations_version);

src/config.ts

@@ -1,5 +1,6 @@
 import dotenv from 'dotenv'
 import jwt from 'jsonwebtoken'
+import path from 'path'
 
 export type StorageBackendType = 'file' | 's3'
 
@@ -29,6 +30,8 @@ type StorageConfigType = {
   dbRefreshMigrationHashesOnMismatch: boolean
   dbSuperUser: string
   dbSearchPath: string
+  dbMigrationHash?: string
+  dbDisableTenantMigrations: boolean
   databaseURL: string
   databaseSSLRootCert?: string
   databasePoolURL?: string
@@ -130,13 +133,15 @@ export function getConfig(options?: { reload?: boolean }): StorageConfigType {
 
   dotenv.config()
 
+  const isMultitenant = getOptionalConfigFromEnv('MULTI_TENANT', 'IS_MULTITENANT') === 'true'
+
   config = {
     // Tenant
     tenantId:
       getOptionalConfigFromEnv('PROJECT_REF') ||
       getOptionalConfigFromEnv('TENANT_ID') ||
       'storage-single-tenant',
-    isMultitenant: getOptionalConfigFromEnv('MULTI_TENANT', 'IS_MULTITENANT') === 'true',
+    isMultitenant,
 
     // Server
     region: getOptionalConfigFromEnv('SERVER_REGION', 'REGION') || 'not-specified',
@@ -230,6 +235,8 @@ export function getConfig(options?: { reload?: boolean }): StorageConfigType {
       getOptionalConfigFromEnv('DB_ALLOW_MIGRATION_REFRESH') === 'false'
     ),
     dbSuperUser: getOptionalConfigFromEnv('DB_SUPER_USER') || 'postgres',
+    dbMigrationHash: getOptionalConfigFromEnv('DB_MIGRATION_HASH'),
+    dbDisableTenantMigrations: getOptionalConfigFromEnv('DB_DISABLE_TENANT_MIGRATIONS') === 'true',
 
     // Database - Connection
     dbSearchPath: getOptionalConfigFromEnv('DATABASE_SEARCH_PATH', 'DB_SEARCH_PATH') || '',

src/database/migrate.ts

@@ -7,8 +7,10 @@ import { validateMigrationHashes } from 'postgres-migrations/dist/validation'
 import { runMigration } from 'postgres-migrations/dist/run-migration'
 import SQL from 'sql-template-strings'
 import { searchPath } from './connection'
+import { updateTenantMigrationVersion } from './tenant'
 
 const {
+  isMultitenant,
   multitenantDatabaseUrl,
   databaseSSLRootCert,
   dbAnonRole,

src/database/tenant.ts

@@ -1,10 +1,10 @@
 import { getConfig } from '../config'
 import { decrypt, verifyJWT } from '../auth'
-import { runMigrationsOnTenant } from './migrate'
 import { knex } from './multitenant-db'
 import { StorageBackendError } from '../storage'
 import { JwtPayload } from 'jsonwebtoken'
 import { PubSubAdapter } from '../pubsub'
+import { RunMigrationsEvent } from '../queue/events/run-migrations'
 
 interface TenantConfig {
   anonKey?: string
@@ -26,7 +26,14 @@ export interface Features {
   }
 }
 
-const { isMultitenant, dbServiceRole, serviceKey, jwtSecret } = getConfig()
+const {
+  isMultitenant,
+  dbServiceRole,
+  serviceKey,
+  jwtSecret,
+  dbMigrationHash,
+  dbDisableTenantMigrations,
+} = getConfig()
 
 const tenantConfigCache = new Map<string, TenantConfig>()
 
@@ -45,29 +52,77 @@ const singleTenantServiceKey:
   : undefined
 
 /**
- * Runs migrations in a specific tenant
- * @param tenantId
- * @param databaseUrl
- * @param logOnError
+ * List all tenants that have not run migrations yet
  */
-export async function runMigrations(
-  tenantId: string,
-  databaseUrl: string,
-  logOnError = false
-): Promise<void> {
+export async function* listTenantsToMigrate() {
+  let lastCursor = 0
+
+  while (true) {
+    const data = await knex
+      .table('tenants')
+      .select('id', 'cursor_id')
+      .where('cursor_id', '>', lastCursor)
+      .where((builder) => {
+        builder
+          .where('migrations_version', '=', dbMigrationHash || '')
+          .orWhere('migrations_version', null)
+      })
+      .orderBy('cursor_id', 'desc')
+      .limit(100)
+
+    yield data.map((tenant) => tenant.id)
+
+    if (data.length === 0) {
+      break
+    }
+
+    lastCursor = data[data.length - 1].cursor_id
+  }
+}
+
+/**
+ * Runs migrations for all tenants
+ */
+export async function runMigrations() {
+  if (dbDisableTenantMigrations) {
+    return
+  }
+  const result = await knex.raw(`SELECT pg_try_advisory_lock(?);`, ['-8575985245963000605'])
+  const lockAcquired = result.rows.shift()?.pg_try_advisory_lock || false
+
+  if (!lockAcquired) {
+    return
+  }
+
   try {
-    await runMigrationsOnTenant(databaseUrl)
-    console.log(`${tenantId} migrations ran successfully`)
-  } catch (error: any) {
-    if (logOnError) {
-      console.error(`${tenantId} migration error:`, error.message)
-      return
-    } else {
-      throw error
+    const tenants = listTenantsToMigrate()
+    for await (const tenantBatch of tenants) {
+      await Promise.allSettled(
+        tenantBatch.map((tenant) => {
+          return RunMigrationsEvent.send({
+            tenantId: tenant,
+            singletonKey: tenant,
+            tenant: {
+              ref: tenant,
+            },
+          })
+        })
+      )
     }
+  } finally {
+    try {
+      await knex.raw(`SELECT pg_advisory_unlock(?);`, ['-8575985245963000605'])
+    } catch (e) {}
   }
 }
 
+export function updateTenantMigrationVersion(tenantIds: string[]) {
+  return knex
+    .table('tenants')
+    .whereIn('id', tenantIds)
+    .update({ migrations_version: dbMigrationHash })
+}
+
 /**
  * Deletes tenants config from the in-memory cache
  * @param tenantId
@@ -124,7 +179,7 @@ export async function getTenantConfig(tenantId: string): Promise<TenantConfig> {
       },
     },
   }
-  await cacheTenantConfigAndRunMigrations(tenantId, config)
+  tenantConfigCache.set(tenantId, config)
   return config
 }
 
@@ -195,12 +250,3 @@ export async function listenForTenantUpdate(pubSub: PubSubAdapter): Promise<void
     tenantConfigCache.delete(tenantId)
   })
 }
-
-async function cacheTenantConfigAndRunMigrations(
-  tenantId: string,
-  config: TenantConfig,
-  logOnError = false
-): Promise<void> {
-  await runMigrations(tenantId, config.databaseUrl, logOnError)
-  tenantConfigCache.set(tenantId, config)
-}

src/http/routes/tenant/index.ts

@@ -3,8 +3,9 @@ import { FromSchema } from 'json-schema-to-ts'
 import apiKey from '../../plugins/apikey'
 import { decrypt, encrypt } from '../../../auth'
 import { knex } from '../../../database/multitenant-db'
-import { deleteTenantConfig, runMigrations } from '../../../database/tenant'
+import { deleteTenantConfig } from '../../../database/tenant'
 import { dbSuperUser, storage } from '../../plugins'
+import { runMigrationsOnTenant } from '../../../database/migrate'
 
 const patchSchema = {
   body: {
@@ -150,7 +151,7 @@ export default async function routes(fastify: FastifyInstance) {
       maxConnections,
     } = request.body
 
-    await runMigrations(tenantId, databaseUrl)
+    await runMigrationsOnTenant(databaseUrl)
     await knex('tenants').insert({
       id: tenantId,
       anon_key: encrypt(anonKey),
@@ -181,7 +182,7 @@ export default async function routes(fastify: FastifyInstance) {
       } = request.body
       const { tenantId } = request.params
       if (databaseUrl) {
-        await runMigrations(tenantId, databaseUrl)
+        await runMigrationsOnTenant(databaseUrl)
       }
       console.log(databasePoolUrl, databasePoolUrl === null)
       await knex('tenants')
@@ -216,7 +217,7 @@ export default async function routes(fastify: FastifyInstance) {
       maxConnections,
     } = request.body
     const { tenantId } = request.params
-    await runMigrations(tenantId, databaseUrl)
+    await runMigrationsOnTenant(databaseUrl)
 
     const tenantInfo: tenantDBInterface = {
       id: tenantId,

src/queue/events/base-event.ts

@@ -1,5 +1,5 @@
 import { Queue } from '../queue'
-import { Job, SendOptions, WorkOptions } from 'pg-boss'
+import { BatchWorkOptions, Job, SendOptions, WorkOptions } from 'pg-boss'
 import { getServiceKeyUser } from '../../database/tenant'
 import { getPostgresConnection } from '../../database'
 import { Storage } from '../../storage'
@@ -11,19 +11,33 @@ import { logger } from '../../monitoring'
 
 export interface BasePayload {
   $version: string
+  singletonKey?: string
   reqId?: string
   tenant: {
     ref: string
     host: string
   }
 }
 
-export type StaticThis<T> = { new (...args: any): T }
-
 const { pgQueueEnable, storageBackendType, storageS3Endpoint } = getConfig()
 const storageS3Protocol = storageS3Endpoint?.includes('http://') ? 'http' : 'https'
 const httpAgent = createAgent(storageS3Protocol)
 
+type StaticThis<T extends BaseEvent<any>> = BaseEventConstructor<T>
+
+interface BaseEventConstructor<Base extends BaseEvent<any>> {
+  version: string
+
+  new (...args: any): Base
+
+  send(
+    this: StaticThis<Base>,
+    payload: Omit<Base['payload'], '$version'>
+  ): Promise<string | void | null>
+
+  eventName(): string
+}
+
 export abstract class BaseEvent<T extends Omit<BasePayload, '$version'>> {
   public static readonly version: string = 'v1'
   protected static queueName = ''
@@ -46,7 +60,7 @@ export abstract class BaseEvent<T extends Omit<BasePayload, '$version'>> {
     return undefined
   }
 
-  static getWorkerOptions(): WorkOptions {
+  static getWorkerOptions(): WorkOptions | BatchWorkOptions {
     return {}
   }
 
@@ -55,7 +69,7 @@ export abstract class BaseEvent<T extends Omit<BasePayload, '$version'>> {
     payload: Omit<T['payload'], '$version'>
   ) {
     if (!payload.$version) {
-      ;(payload as any).$version = (this as any).version
+      ;(payload as T['payload']).$version = this.version
     }
     const that = new this(payload)
     return that.send()
@@ -67,13 +81,13 @@ export abstract class BaseEvent<T extends Omit<BasePayload, '$version'>> {
   ) {
     // eslint-disable-next-line @typescript-eslint/no-var-requires
     const { Webhook } = require('./webhook')
-    const eventType = (this as any).eventName()
+    const eventType = this.eventName()
 
     try {
       await Webhook.send({
         event: {
           type: eventType,
-          $version: (this as any).version,
+          $version: this.version,
           applyTime: Date.now(),
           payload,
         },
@@ -85,7 +99,7 @@ export abstract class BaseEvent<T extends Omit<BasePayload, '$version'>> {
           error: e,
           event: {
             type: eventType,
-            $version: (this as any).version,
+            $version: this.version,
             applyTime: Date.now(),
             payload: JSON.stringify(payload),
           },
@@ -96,7 +110,7 @@ export abstract class BaseEvent<T extends Omit<BasePayload, '$version'>> {
     }
   }
 
-  static handle(job: Job<BaseEvent<any>['payload']>) {
+  static handle(job: Job<BaseEvent<any>['payload']> | Job<BaseEvent<any>['payload']>[]) {
     throw new Error('not implemented')
   }
 
@@ -122,6 +136,10 @@ export abstract class BaseEvent<T extends Omit<BasePayload, '$version'>> {
     return new Storage(storageBackend, db)
   }
 
+  singletonKey(payload: T) {
+    return
+  }
+
   async send() {
     const constructor = this.constructor as typeof BaseEvent
 
@@ -144,7 +162,10 @@ export abstract class BaseEvent<T extends Omit<BasePayload, '$version'>> {
         ...this.payload,
         $version: constructor.version,
       },
-      options: constructor.getQueueOptions(),
+      options: {
+        ...constructor.getQueueOptions(),
+        singletonKey: this.payload.singletonKey,
+      },
     })
 
     timer({