From e06596e51ca2d6f1d0ac50338dcba69f4a232d00 Mon Sep 17 00:00:00 2001 From: Ondrej Sedlak Date: Fri, 13 Nov 2026 23:34:09 +0100 Subject: [PATCH] feat: add replication plugin to pull external data into internal sqlite Adds a ReplicationPlugin that periodically pulls rows from the configured external data source into the internal Durable Object SQLite database, with per-table watermark columns for append-only polling, batch limits, admin-triggered sync/status routes, and interval-based auto sync. --- plugins/replication/README.md | 49 ++++ plugins/replication/index.test.ts | 217 ++++++++++++++++++ plugins/replication/index.ts | 369 ++++++++++++++++++++++++++++++ plugins/replication/meta.json | 13 ++ src/index.ts | 22 ++ wrangler.toml | 5 + 6 files changed, 675 insertions(+) create mode 100644 plugins/replication/README.md create mode 100644 plugins/replication/index.test.ts create mode 100644 plugins/replication/index.ts create mode 100644 plugins/replication/meta.json diff --git a/plugins/replication/README.md b/plugins/replication/README.md new file mode 100644 index 0000000..cb68892 --- /dev/null +++ b/plugins/replication/README.md @@ -0,0 +1,49 @@ +## Replication Plugin + +Pulls data from the configured external data source (e.g. Postgres on Supabase) +into the internal Durable Object SQLite database, so the StarbaseDB instance +acts as a close-to-edge replica. + +### Example Usage + +```ts +import { ReplicationPlugin } from '../plugins/replication' + +// .... +// .... + +const replicationPlugin = new ReplicationPlugin({ + // How often (in milliseconds) an automatic sync cycle should run. + // Omit to only sync when triggered manually. + intervalMs: 60_000, + tables: [ + { + table: 'users', + schema: 'public', + // Append-only polling: only rows with a greater `id` than the + // last synced watermark are pulled on each cycle. + watermarkColumn: 'id', + batchSize: 500, + }, + { + table: 'products', + schema: 'public', + }, + ], +}) + +const plugins = [ + // ... + replicationPlugin, +] satisfies StarbasePlugin[] +``` + +### Behavior + +- A `tmp_replication_state` table is created in the internal database to track + the last synced watermark (`last_value`) and timestamp (`last_sync_at`) per table. +- On each sync cycle, rows are pulled from the external source with + `SELECT * FROM "schema"."table" WHERE "watermark" > ? ORDER BY "watermark" ASC LIMIT ?` + and written locally with `INSERT OR REPLACE`. +- Admin users can trigger a sync manually with `POST /replication/sync` and + inspect replication state with `GET /replication/status`. diff --git a/plugins/replication/index.test.ts b/plugins/replication/index.test.ts new file mode 100644 index 0000000..6130e4c --- /dev/null +++ b/plugins/replication/index.test.ts @@ -0,0 +1,217 @@ +import { describe, it, expect, vi, beforeEach } from 'vitest' +import { ReplicationPlugin, ReplicationTableConfig } from './index' +import { StarbaseDBConfiguration } from '../../src/handler' +import { DataSource, QueryResult } from '../../src/types' + +function createMockDataSource() { + const store = new Map() + const calls: { sql: string; params?: unknown[] }[] = [] + + const rpc = { + executeQuery: vi.fn(async ({ sql, params }: { sql: string; params?: unknown[] }) => { + calls.push({ sql, params }) + const normalized = sql.trim().toUpperCase() + + if (normalized.startsWith('CREATE TABLE')) { + return [] + } + if ( + normalized.startsWith('SELECT') && + normalized.includes('WHERE TABLE_NAME = ?') + ) { + const tableName = params?.[0] as string + return store.get(`state:${tableName}`) ?? [] + } + if (normalized.startsWith('SELECT')) { + const rows: any[] = [] + for (const [key, value] of store) { + if (key.startsWith('state:')) rows.push(...value) + } + return rows + } + if (normalized.startsWith('INSERT OR REPLACE INTO TMP_REPLICATION_STATE')) { + store.set(`state:${params?.[0]}`, [ + { + table_name: params?.[0], + schema_name: params?.[1], + watermark_column: params?.[2], + last_value: params?.[3], + last_sync_at: params?.[4], + }, + ]) + return [] + } + if (normalized.startsWith('INSERT OR REPLACE INTO')) { + const tableMatch = sql.match(/INTO "([^"]+)"/) + const table = tableMatch?.[1] ?? 'unknown' + const key = `rows:${table}` + const existing = store.get(key) ?? [] + existing.push(params) + store.set(key, existing) + return [] + } + return [] + }), + } + + const dataSource = { + rpc, + source: 'internal', + } as unknown as DataSource + + return { dataSource, store, calls } +} + +const baseTables: ReplicationTableConfig[] = [ + { table: 'users', schema: 'public', watermarkColumn: 'id' }, +] + +let mockConfig: StarbaseDBConfiguration + +beforeEach(() => { + vi.clearAllMocks() + mockConfig = { role: 'admin' } as StarbaseDBConfiguration +}) + +describe('ReplicationPlugin - syncNow', () => { + it('pulls rows from the external source and writes them locally', async () => { + const { dataSource, store } = createMockDataSource() + const externalRows: QueryResult[] = [ + { id: 1, name: 'Alice' }, + { id: 2, name: 'Bob' }, + ] + const queryExternal = vi.fn().mockResolvedValue(externalRows) + + const plugin = new ReplicationPlugin({ + tables: baseTables, + queryExternal: queryExternal as any, + }) + ;(plugin as any).dataSource = dataSource + ;(plugin as any).config = mockConfig + + const results = await plugin.syncNow() + + expect(results).toHaveLength(1) + expect(results[0]).toEqual({ + table: 'users', + rowsReplicated: 2, + lastValue: '2', + }) + + // External query should target the qualified table with a LIMIT. + const externalCall = queryExternal.mock.calls[0][0] + expect(externalCall.sql).toContain('FROM "public"."users"') + expect(externalCall.sql).toContain('ORDER BY "id" ASC') + expect(externalCall.sql).toContain('LIMIT ?') + expect(externalCall.params).toEqual([500]) + + // Rows should have been inserted locally. + expect(store.get('rows:users')).toHaveLength(2) + + // State should have been persisted with the max watermark. + const state = store.get('state:users') + expect(state[0].last_value).toBe('2') + expect(state[0].watermark_column).toBe('id') + }) + + it('appends only rows newer than the last watermark on subsequent syncs', async () => { + const { dataSource, store } = createMockDataSource() + store.set('state:users', [ + { + table_name: 'users', + schema_name: 'public', + watermark_column: 'id', + last_value: '2', + last_sync_at: Date.now(), + }, + ]) + + const queryExternal = vi + .fn() + .mockResolvedValue([{ id: 3, name: 'Carol' }]) + + const plugin = new ReplicationPlugin({ + tables: baseTables, + queryExternal: queryExternal as any, + }) + ;(plugin as any).dataSource = dataSource + ;(plugin as any).config = mockConfig + + const results = await plugin.syncNow() + + const externalCall = queryExternal.mock.calls[0][0] + expect(externalCall.sql).toContain('WHERE "id" > ?') + expect(externalCall.params).toEqual(['2', 500]) + expect(results[0].lastValue).toBe('3') + }) + + it('rejects unsafe table identifiers', async () => { + const { dataSource } = createMockDataSource() + const queryExternal = vi.fn().mockResolvedValue([]) + + const plugin = new ReplicationPlugin({ + tables: [{ table: 'users; DROP TABLE users' }], + queryExternal: queryExternal as any, + }) + ;(plugin as any).dataSource = dataSource + ;(plugin as any).config = mockConfig + + const results = await plugin.syncNow() + expect(results).toEqual([]) + expect(queryExternal).not.toHaveBeenCalled() + }) + + it('handles empty external result sets without inserting rows', async () => { + const { dataSource, store } = createMockDataSource() + const queryExternal = vi.fn().mockResolvedValue([]) + + const plugin = new ReplicationPlugin({ + tables: baseTables, + queryExternal: queryExternal as any, + }) + ;(plugin as any).dataSource = dataSource + ;(plugin as any).config = mockConfig + + const results = await plugin.syncNow() + + expect(results[0].rowsReplicated).toBe(0) + expect(results[0].lastValue).toBeNull() + expect(store.get('rows:users')).toBeUndefined() + }) + + it('throws when the plugin has not been initialized', async () => { + const plugin = new ReplicationPlugin({ + tables: baseTables, + queryExternal: vi.fn() as any, + }) + + await expect(plugin.syncNow()).rejects.toThrow( + 'ReplicationPlugin not properly initialized' + ) + }) +}) + +describe('ReplicationPlugin - getStatus', () => { + it('returns the persisted replication state', async () => { + const { dataSource, store } = createMockDataSource() + store.set('state:users', [ + { + table_name: 'users', + schema_name: 'public', + watermark_column: 'id', + last_value: '7', + last_sync_at: 123, + }, + ]) + + const plugin = new ReplicationPlugin({ + tables: baseTables, + queryExternal: vi.fn() as any, + }) + ;(plugin as any).dataSource = dataSource + + const status = await plugin.getStatus() + expect(status).toHaveLength(1) + expect(status[0].last_value).toBe('7') + }) +}) diff --git a/plugins/replication/index.ts b/plugins/replication/index.ts new file mode 100644 index 0000000..a324780 --- /dev/null +++ b/plugins/replication/index.ts @@ -0,0 +1,369 @@ +import { StarbaseApp, StarbaseDBConfiguration } from '../../src/handler' +import { StarbasePlugin } from '../../src/plugin' +import { DataSource, QueryResult } from '../../src/types' +import { createResponse } from '../../src/utils' +import { executeExternalQuery } from '../../src/operation' + +/** + * Configuration for a single table that should be replicated from the + * external data source into the internal Durable Object SQLite database. + */ +export interface ReplicationTableConfig { + // Name of the table on the external data source (and locally). + table: string + // Schema of the table on the external data source (e.g. `public` for Postgres). + schema?: string + // Optional column used for append-only polling (e.g. `id` or `created_at`). + // When provided, only rows with a greater value than the last synced + // watermark are pulled on each cycle. + watermarkColumn?: string + // Maximum number of rows pulled per sync cycle. Defaults to 500. + batchSize?: number +} + +export interface ReplicationPluginOptions { + // Tables that should be replicated from the external source. + tables: ReplicationTableConfig[] + // Interval (in milliseconds) between automatic sync cycles. When omitted, + // syncs only run when triggered manually via the `/replication/sync` route + // or by calling `syncNow()`. + intervalMs?: number + // Injectable external query executor, primarily for testing. Defaults to + // running the query against the configured external data source. + queryExternal?: (opts: { + sql: string + params?: unknown[] + dataSource: DataSource + config: StarbaseDBConfiguration + }) => Promise +} + +interface ReplicationStateRow { + table_name: string + schema_name: string | null + watermark_column: string | null + last_value: string | null + last_sync_at: number | null +} + +export interface ReplicationSyncResult { + table: string + rowsReplicated: number + lastValue: string | null +} + +const DEFAULT_BATCH_SIZE = 500 + +const SQL_QUERIES = { + CREATE_STATE_TABLE: ` + CREATE TABLE IF NOT EXISTS tmp_replication_state ( + table_name TEXT PRIMARY KEY, + schema_name TEXT, + watermark_column TEXT, + last_value TEXT, + last_sync_at INTEGER + ) + `, + GET_STATE: ` + SELECT table_name, schema_name, watermark_column, last_value, last_sync_at + FROM tmp_replication_state + WHERE table_name = ? + `, + UPSERT_STATE: ` + INSERT OR REPLACE INTO tmp_replication_state + (table_name, schema_name, watermark_column, last_value, last_sync_at) + VALUES (?, ?, ?, ?, ?) + `, + GET_ALL_STATE: ` + SELECT table_name, schema_name, watermark_column, last_value, last_sync_at + FROM tmp_replication_state + `, +} + +// Identifiers are interpolated into SQL statements, so restrict them to a +// safe subset of characters to prevent SQL injection. +const IDENTIFIER_PATTERN = /^[A-Za-z_][A-Za-z0-9_]*$/ + +function assertIdentifier(value: string, kind: string): string { + if (!IDENTIFIER_PATTERN.test(value)) { + throw new Error(`Invalid ${kind} identifier: ${value}`) + } + return value +} + +function quoteIdentifier(value: string): string { + return `"${assertIdentifier(value, 'table/column').replace(/"/g, '""')}"` +} + +export class ReplicationPlugin extends StarbasePlugin { + public pathPrefix: string = '/replication' + private tables: ReplicationTableConfig[] + private intervalMs?: number + private queryExternal: NonNullable + private dataSource?: DataSource + private config?: StarbaseDBConfiguration + private syncing: boolean = false + + constructor(opts: ReplicationPluginOptions) { + super('starbasedb:replication', { + requiresAuth: true, + }) + this.tables = opts.tables + this.intervalMs = opts.intervalMs + this.queryExternal = + opts.queryExternal ?? + (async ({ sql, params, dataSource, config }) => + (await executeExternalQuery({ + sql, + params, + dataSource, + config, + })) as QueryResult[]) + } + + override async register(app: StarbaseApp) { + app.use(async (c, next) => { + this.dataSource = c?.get('dataSource') + this.config = c?.get('config') + await this.init() + await this.maybeAutoSync() + await next() + }) + + // Manually trigger a replication sync cycle (admin only). + app.post(`${this.pathPrefix}/sync`, async (c) => { + if (this.config?.role !== 'admin') { + return createResponse(undefined, 'Unauthorized request', 401) + } + + const results = await this.syncNow() + return createResponse({ results }, undefined, 200) + }) + + // Inspect the replication state of each configured table (admin only). + app.get(`${this.pathPrefix}/status`, async (c) => { + if (this.config?.role !== 'admin') { + return createResponse(undefined, 'Unauthorized request', 401) + } + + const state = await this.getStatus() + return createResponse({ state }, undefined, 200) + }) + } + + private async init() { + if (!this.dataSource) return + + await this.dataSource.rpc.executeQuery({ + sql: SQL_QUERIES.CREATE_STATE_TABLE, + params: [], + }) + } + + /** + * On each request, check whether the configured interval has elapsed since + * the last sync and, if so, kick off a background sync cycle. + */ + private async maybeAutoSync() { + if (!this.dataSource || !this.intervalMs || this.tables.length === 0) + return + if (this.syncing) return + + try { + const state = (await this.dataSource.rpc.executeQuery({ + sql: SQL_QUERIES.GET_ALL_STATE, + params: [], + })) as ReplicationStateRow[] + + const now = Date.now() + const due = this.tables.some((t) => { + const row = state.find((s) => s.table_name === t.table) + return !row?.last_sync_at || now - row.last_sync_at >= this.intervalMs! + }) + + if (due) { + const promise = this.syncNow() + if (this.dataSource.executionContext) { + this.dataSource.executionContext.waitUntil(promise) + } + } + } catch (error) { + console.error('Error checking replication schedule:', error) + } + } + + /** + * Run a single replication cycle across all configured tables. + */ + public async syncNow(): Promise { + if (!this.dataSource) { + throw new Error('ReplicationPlugin not properly initialized') + } + if (this.syncing) { + return [] + } + + this.syncing = true + try { + const results: ReplicationSyncResult[] = [] + for (const table of this.tables) { + try { + results.push(await this.syncTable(table)) + } catch (error) { + console.error( + `Error replicating table ${table.table}:`, + error + ) + } + } + return results + } finally { + this.syncing = false + } + } + + private async syncTable( + table: ReplicationTableConfig + ): Promise { + const dataSource = this.dataSource! + const config = this.config! + const batchSize = table.batchSize ?? DEFAULT_BATCH_SIZE + + const tableName = assertIdentifier(table.table, 'table') + const schemaName = table.schema + ? assertIdentifier(table.schema, 'schema') + : undefined + const watermarkColumn = table.watermarkColumn + ? assertIdentifier(table.watermarkColumn, 'watermark column') + : undefined + + const qualifiedName = schemaName + ? `${quoteIdentifier(schemaName)}.${quoteIdentifier(tableName)}` + : quoteIdentifier(tableName) + + // Load the last known watermark for this table. + const stateRows = (await dataSource.rpc.executeQuery({ + sql: SQL_QUERIES.GET_STATE, + params: [tableName], + })) as ReplicationStateRow[] + const state = stateRows[0] + const lastValue = state?.last_value ?? null + + // Pull the next batch of rows from the external source. + let externalSql = `SELECT * FROM ${qualifiedName}` + const params: unknown[] = [] + if (watermarkColumn && lastValue !== null) { + externalSql += ` WHERE ${quoteIdentifier(watermarkColumn)} > ?` + params.push(lastValue) + } + if (watermarkColumn) { + externalSql += ` ORDER BY ${quoteIdentifier(watermarkColumn)} ASC` + } + externalSql += ` LIMIT ?` + params.push(batchSize) + + const rows = await this.queryExternal({ + sql: externalSql, + params, + dataSource, + config, + }) + + let newLastValue = lastValue + if (rows.length > 0) { + await this.ensureLocalTable(tableName, rows[0]) + await this.insertRows(tableName, rows) + + if (watermarkColumn) { + for (const row of rows) { + const value = row[watermarkColumn] + if (value !== null && value !== undefined) { + newLastValue = this.maxWatermark(newLastValue, value) + } + } + } + } + + await dataSource.rpc.executeQuery({ + sql: SQL_QUERIES.UPSERT_STATE, + params: [ + tableName, + schemaName ?? null, + watermarkColumn ?? null, + newLastValue !== null ? String(newLastValue) : null, + Date.now(), + ], + }) + + return { + table: tableName, + rowsReplicated: rows.length, + lastValue: newLastValue !== null ? String(newLastValue) : null, + } + } + + /** + * Create the local replica table if it does not exist yet. Columns are + * declared without a type so SQLite preserves the original value types + * (important for numeric watermark comparisons). + */ + private async ensureLocalTable( + tableName: string, + sampleRow: QueryResult + ) { + const columns = Object.keys(sampleRow) + if (columns.length === 0) return + + const columnDefs = columns.map((c) => quoteIdentifier(c)).join(', ') + await this.dataSource!.rpc.executeQuery({ + sql: `CREATE TABLE IF NOT EXISTS ${quoteIdentifier(tableName)} (${columnDefs})`, + params: [], + }) + } + + private async insertRows(tableName: string, rows: QueryResult[]) { + const columns = Object.keys(rows[0]) + const placeholders = columns.map(() => '?').join(', ') + const sql = `INSERT OR REPLACE INTO ${quoteIdentifier(tableName)} (${columns + .map((c) => quoteIdentifier(c)) + .join(', ')}) VALUES (${placeholders})` + + for (const row of rows) { + await this.dataSource!.rpc.executeQuery({ + sql, + params: columns.map((c) => row[c] as unknown), + }) + } + } + + /** + * Return the greater of two watermark values, comparing numerically when + * both values are numeric and lexically otherwise. + */ + private maxWatermark( + current: string | null, + value: SqlStorageValue + ): string { + const next = String(value) + if (current === null) return next + + const currentNum = Number(current) + const nextNum = Number(next) + if (!Number.isNaN(currentNum) && !Number.isNaN(nextNum)) { + return nextNum > currentNum ? next : current + } + + return next > current ? next : current + } + + public async getStatus(): Promise { + if (!this.dataSource) { + throw new Error('ReplicationPlugin not properly initialized') + } + + return (await this.dataSource.rpc.executeQuery({ + sql: SQL_QUERIES.GET_ALL_STATE, + params: [], + })) as ReplicationStateRow[] + } +} diff --git a/plugins/replication/meta.json b/plugins/replication/meta.json new file mode 100644 index 0000000..bc45f59 --- /dev/null +++ b/plugins/replication/meta.json @@ -0,0 +1,13 @@ +{ + "version": "1.0.0", + "resources": { + "tables": {}, + "secrets": {}, + "variables": {} + }, + "dependencies": { + "tables": {}, + "secrets": {}, + "variables": {} + } +} diff --git a/src/index.ts b/src/index.ts index 4d08932..cdb978f 100644 --- a/src/index.ts +++ b/src/index.ts @@ -12,6 +12,10 @@ import { QueryLogPlugin } from '../plugins/query-log' import { StatsPlugin } from '../plugins/stats' import { CronPlugin } from '../plugins/cron' import { InterfacePlugin } from '../plugins/interface' +import { + ReplicationPlugin, + ReplicationTableConfig, +} from '../plugins/replication' export { StarbaseDBDurableObject } from './do' @@ -54,6 +58,10 @@ export interface Env { AUTH_ALGORITHM?: string AUTH_JWKS_ENDPOINT?: string + // Replication plugin configuration (JSON array of table configs) + REPLICATION_TABLES?: string + REPLICATION_INTERVAL_MS?: number + HYPERDRIVE: Hyperdrive // ## DO NOT REMOVE: TEMPLATE INTERFACE ## @@ -211,6 +219,19 @@ export default { const interfacePlugin = new InterfacePlugin() + let replicationTables: ReplicationTableConfig[] = [] + try { + replicationTables = env.REPLICATION_TABLES + ? JSON.parse(env.REPLICATION_TABLES) + : [] + } catch (error) { + console.error('Failed to parse REPLICATION_TABLES:', error) + } + const replicationPlugin = new ReplicationPlugin({ + tables: replicationTables, + intervalMs: env.REPLICATION_INTERVAL_MS, + }) + const plugins = [ webSocketPlugin, new StudioPlugin({ @@ -226,6 +247,7 @@ export default { cronPlugin, new StatsPlugin(), interfacePlugin, + replicationPlugin, ] satisfies StarbasePlugin[] const starbase = new StarbaseDB({ diff --git a/wrangler.toml b/wrangler.toml index 395c4ac..36ae208 100644 --- a/wrangler.toml +++ b/wrangler.toml @@ -75,6 +75,11 @@ ENABLE_RLS = 0 AUTH_ALGORITHM = "RS256" AUTH_JWKS_ENDPOINT = "" +# Replication plugin: pull data from the external source into the internal +# Durable Object SQLite on an interval. +# REPLICATION_TABLES = '[{"table":"users","schema":"public","watermarkColumn":"id","batchSize":500}]' +# REPLICATION_INTERVAL_MS = 60000 + # [[hyperdrive]] # binding = "HYPERDRIVE" # id = ""