@powersync/service-module-postgres 0.0.0-dev-20240918092408

package/src/api/PostgresRouteAPIAdapter.ts

@@ -0,0 +1,307 @@
+import { api, ParseSyncRulesOptions } from '@powersync/service-core';
+import * as pgwire from '@powersync/service-jpgwire';
+
+import * as sync_rules from '@powersync/service-sync-rules';
+import * as service_types from '@powersync/service-types';
+import * as replication_utils from '../replication/replication-utils.js';
+import * as types from '../types/types.js';
+import * as pg_utils from '../utils/pgwire_utils.js';
+import { getDebugTableInfo } from '../replication/replication-utils.js';
+import { PUBLICATION_NAME } from '../replication/WalStream.js';
+
+export class PostgresRouteAPIAdapter implements api.RouteAPI {
+  protected pool: pgwire.PgClient;
+
+  connectionTag: string;
+  // TODO this should probably be configurable one day
+  publicationName = PUBLICATION_NAME;
+
+  constructor(protected config: types.ResolvedConnectionConfig) {
+    this.pool = pgwire.connectPgWirePool(config, {
+      idleTimeout: 30_000
+    });
+    this.connectionTag = config.tag ?? sync_rules.DEFAULT_TAG;
+  }
+
+  getParseSyncRulesOptions(): ParseSyncRulesOptions {
+    return {
+      defaultSchema: 'public'
+    };
+  }
+
+  async shutdown(): Promise<void> {
+    await this.pool.end();
+  }
+
+  async getSourceConfig(): Promise<service_types.configFile.DataSourceConfig> {
+    return this.config;
+  }
+
+  async getConnectionStatus(): Promise<service_types.ConnectionStatusV2> {
+    const base = {
+      id: this.config.id,
+      uri: types.baseUri(this.config)
+    };
+
+    try {
+      await pg_utils.retriedQuery(this.pool, `SELECT 'PowerSync connection test'`);
+    } catch (e) {
+      return {
+        ...base,
+        connected: false,
+        errors: [{ level: 'fatal', message: e.message }]
+      };
+    }
+
+    try {
+      await replication_utils.checkSourceConfiguration(this.pool, this.publicationName);
+    } catch (e) {
+      return {
+        ...base,
+        connected: true,
+        errors: [{ level: 'fatal', message: e.message }]
+      };
+    }
+
+    return {
+      ...base,
+      connected: true,
+      errors: []
+    };
+  }
+
+  async executeQuery(query: string, params: any[]): Promise<service_types.internal_routes.ExecuteSqlResponse> {
+    if (!this.config.debug_api) {
+      return service_types.internal_routes.ExecuteSqlResponse.encode({
+        results: {
+          columns: [],
+          rows: []
+        },
+        success: false,
+        error: 'SQL querying is not enabled'
+      });
+    }
+
+    try {
+      const result = await this.pool.query({
+        statement: query,
+        params: params.map(pg_utils.autoParameter)
+      });
+
+      return service_types.internal_routes.ExecuteSqlResponse.encode({
+        success: true,
+        results: {
+          columns: result.columns.map((c) => c.name),
+          rows: result.rows.map((row) => {
+            return row.map((value) => {
+              const sqlValue = sync_rules.toSyncRulesValue(value);
+              if (typeof sqlValue == 'bigint') {
+                return Number(value);
+              } else if (sync_rules.isJsonValue(sqlValue)) {
+                return sqlValue;
+              } else {
+                return null;
+              }
+            });
+          })
+        }
+      });
+    } catch (e) {
+      return service_types.internal_routes.ExecuteSqlResponse.encode({
+        results: {
+          columns: [],
+          rows: []
+        },
+        success: false,
+        error: e.message
+      });
+    }
+  }
+
+  async getDebugTablesInfo(
+    tablePatterns: sync_rules.TablePattern[],
+    sqlSyncRules: sync_rules.SqlSyncRules
+  ): Promise<api.PatternResult[]> {
+    let result: api.PatternResult[] = [];
+
+    for (let tablePattern of tablePatterns) {
+      const schema = tablePattern.schema;
+
+      let patternResult: api.PatternResult = {
+        schema: schema,
+        pattern: tablePattern.tablePattern,
+        wildcard: tablePattern.isWildcard
+      };
+      result.push(patternResult);
+
+      if (tablePattern.isWildcard) {
+        patternResult.tables = [];
+        const prefix = tablePattern.tablePrefix;
+        const results = await pg_utils.retriedQuery(this.pool, {
+          statement: `SELECT c.oid AS relid, c.relname AS table_name
+        FROM pg_class c
+        JOIN pg_namespace n ON n.oid = c.relnamespace
+        WHERE n.nspname = $1
+        AND c.relkind = 'r'
+        AND c.relname LIKE $2`,
+          params: [
+            { type: 'varchar', value: schema },
+            { type: 'varchar', value: tablePattern.tablePattern }
+          ]
+        });
+
+        for (let row of pgwire.pgwireRows(results)) {
+          const name = row.table_name as string;
+          const relationId = row.relid as number;
+          if (!name.startsWith(prefix)) {
+            continue;
+          }
+          const details = await this.getDebugTableInfo(tablePattern, name, relationId, sqlSyncRules);
+          patternResult.tables.push(details);
+        }
+      } else {
+        const results = await pg_utils.retriedQuery(this.pool, {
+          statement: `SELECT c.oid AS relid, c.relname AS table_name
+        FROM pg_class c
+        JOIN pg_namespace n ON n.oid = c.relnamespace
+        WHERE n.nspname = $1
+        AND c.relkind = 'r'
+        AND c.relname = $2`,
+          params: [
+            { type: 'varchar', value: schema },
+            { type: 'varchar', value: tablePattern.tablePattern }
+          ]
+        });
+        if (results.rows.length == 0) {
+          // Table not found
+          patternResult.table = await this.getDebugTableInfo(tablePattern, tablePattern.name, null, sqlSyncRules);
+        } else {
+          const row = pgwire.pgwireRows(results)[0];
+          const name = row.table_name as string;
+          const relationId = row.relid as number;
+          patternResult.table = await this.getDebugTableInfo(tablePattern, name, relationId, sqlSyncRules);
+        }
+      }
+    }
+    return result;
+  }
+
+  protected async getDebugTableInfo(
+    tablePattern: sync_rules.TablePattern,
+    name: string,
+    relationId: number | null,
+    syncRules: sync_rules.SqlSyncRules
+  ): Promise<service_types.TableInfo> {
+    return getDebugTableInfo({
+      db: this.pool,
+      name: name,
+      publicationName: this.publicationName,
+      connectionTag: this.connectionTag,
+      tablePattern: tablePattern,
+      relationId: relationId,
+      syncRules: syncRules
+    });
+  }
+
+  async getReplicationLag(syncRulesId: string): Promise<number> {
+    const results = await pg_utils.retriedQuery(this.pool, {
+      statement: `SELECT
+  slot_name,
+  confirmed_flush_lsn, 
+  pg_current_wal_lsn(), 
+  (pg_current_wal_lsn() - confirmed_flush_lsn) AS lsn_distance
+FROM pg_replication_slots WHERE slot_name = $1 LIMIT 1;`,
+      params: [{ type: 'varchar', value: syncRulesId }]
+    });
+    const [row] = pgwire.pgwireRows(results);
+    if (row) {
+      return Number(row.lsn_distance);
+    }
+
+    throw new Error(`Could not determine replication lag for slot ${syncRulesId}`);
+  }
+
+  async getReplicationHead(): Promise<string> {
+    const [{ lsn }] = pgwire.pgwireRows(
+      await pg_utils.retriedQuery(this.pool, `SELECT pg_logical_emit_message(false, 'powersync', 'ping') as lsn`)
+    );
+    return String(lsn);
+  }
+
+  async getConnectionSchema(): Promise<service_types.DatabaseSchema[]> {
+    // https://github.com/Borvik/vscode-postgres/blob/88ec5ed061a0c9bced6c5d4ec122d0759c3f3247/src/language/server.ts
+    const results = await pg_utils.retriedQuery(
+      this.pool,
+      `SELECT
+tbl.schemaname,
+tbl.tablename,
+tbl.quoted_name,
+json_agg(a ORDER BY attnum) as columns
+FROM
+(
+  SELECT
+    n.nspname as schemaname,
+    c.relname as tablename,
+    (quote_ident(n.nspname) || '.' || quote_ident(c.relname)) as quoted_name
+  FROM
+    pg_catalog.pg_class c
+    JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace
+  WHERE
+    c.relkind = 'r'
+    AND n.nspname not in ('information_schema', 'pg_catalog', 'pg_toast')
+    AND n.nspname not like 'pg_temp_%'
+    AND n.nspname not like 'pg_toast_temp_%'
+    AND c.relnatts > 0
+    AND has_schema_privilege(n.oid, 'USAGE') = true
+    AND has_table_privilege(quote_ident(n.nspname) || '.' || quote_ident(c.relname), 'SELECT, INSERT, UPDATE, DELETE, TRUNCATE, REFERENCES, TRIGGER') = true
+) as tbl
+LEFT JOIN (
+  SELECT
+    attrelid,
+    attname,
+    format_type(atttypid, atttypmod) as data_type,
+    (SELECT typname FROM pg_catalog.pg_type WHERE oid = atttypid) as pg_type,
+    attnum,
+    attisdropped
+  FROM
+    pg_attribute
+) as a ON (
+  a.attrelid = tbl.quoted_name::regclass
+  AND a.attnum > 0
+  AND NOT a.attisdropped
+  AND has_column_privilege(tbl.quoted_name, a.attname, 'SELECT, INSERT, UPDATE, REFERENCES')
+)
+GROUP BY schemaname, tablename, quoted_name`
+    );
+    const rows = pgwire.pgwireRows(results);
+
+    let schemas: Record<string, any> = {};
+
+    for (let row of rows) {
+      const schema = (schemas[row.schemaname] ??= {
+        name: row.schemaname,
+        tables: []
+      });
+      const table = {
+        name: row.tablename,
+        columns: [] as any[]
+      };
+      schema.tables.push(table);
+
+      const columnInfo = JSON.parse(row.columns);
+      for (let column of columnInfo) {
+        let pg_type = column.pg_type as string;
+        if (pg_type.startsWith('_')) {
+          pg_type = `${pg_type.substring(1)}[]`;
+        }
+        table.columns.push({
+          name: column.attname,
+          type: column.data_type,
+          pg_type: pg_type
+        });
+      }
+    }
+
+    return Object.values(schemas);
+  }
+}

package/src/auth/SupabaseKeyCollector.ts

@@ -0,0 +1,70 @@
+import { auth } from '@powersync/service-core';
+import * as pgwire from '@powersync/service-jpgwire';
+import * as jose from 'jose';
+
+import * as types from '../types/types.js';
+import * as pgwire_utils from '../utils/pgwire_utils.js';
+
+/**
+ * Fetches key from the Supabase database.
+ *
+ * Unfortunately, despite the JWTs containing a kid, we have no way to lookup that kid
+ * before receiving a valid token.
+ */
+export class SupabaseKeyCollector implements auth.KeyCollector {
+  private pool: pgwire.PgClient;
+
+  private keyOptions: auth.KeyOptions = {
+    requiresAudience: ['authenticated'],
+    maxLifetimeSeconds: 86400 * 7 + 1200 // 1 week + 20 minutes margin
+  };
+
+  constructor(connectionConfig: types.ResolvedConnectionConfig) {
+    this.pool = pgwire.connectPgWirePool(connectionConfig, {
+      // To avoid overloading the source database with open connections,
+      // limit to a single connection, and close the connection shortly
+      // after using it.
+      idleTimeout: 5_000,
+      maxSize: 1
+    });
+  }
+
+  shutdown() {
+    return this.pool.end();
+  }
+
+  async getKeys() {
+    let row: { jwt_secret: string };
+    try {
+      const rows = pgwire.pgwireRows(
+        await pgwire_utils.retriedQuery(this.pool, `SELECT current_setting('app.settings.jwt_secret') as jwt_secret`)
+      );
+      row = rows[0] as any;
+    } catch (e) {
+      if (e.message?.includes('unrecognized configuration parameter')) {
+        throw new jose.errors.JOSEError(`Generate a new JWT secret on Supabase. Cause: ${e.message}`);
+      } else {
+        throw e;
+      }
+    }
+    const secret = row?.jwt_secret as string | undefined;
+    if (secret == null) {
+      return {
+        keys: [],
+        errors: [new jose.errors.JWKSNoMatchingKey()]
+      };
+    } else {
+      const key: jose.JWK = {
+        kty: 'oct',
+        alg: 'HS256',
+        // While the secret is valid base64, the base64-encoded form is the secret value.
+        k: Buffer.from(secret, 'utf8').toString('base64url')
+      };
+      const imported = await auth.KeySpec.importKey(key, this.keyOptions);
+      return {
+        keys: [imported],
+        errors: []
+      };
+    }
+  }
+}

package/src/index.ts

@@ -0,0 +1,5 @@
+import { PostgresModule } from './module/PostgresModule.js';
+
+export const module = new PostgresModule();
+
+export default module;

package/src/module/PostgresModule.ts

@@ -0,0 +1,122 @@
+import {
+  api,
+  auth,
+  ConfigurationFileSyncRulesProvider,
+  replication,
+  system,
+  TearDownOptions
+} from '@powersync/service-core';
+import * as jpgwire from '@powersync/service-jpgwire';
+import * as types from '../types/types.js';
+import { PostgresRouteAPIAdapter } from '../api/PostgresRouteAPIAdapter.js';
+import { SupabaseKeyCollector } from '../auth/SupabaseKeyCollector.js';
+import { WalStreamReplicator } from '../replication/WalStreamReplicator.js';
+import { ConnectionManagerFactory } from '../replication/ConnectionManagerFactory.js';
+import { PostgresErrorRateLimiter } from '../replication/PostgresErrorRateLimiter.js';
+import { cleanUpReplicationSlot } from '../replication/replication-utils.js';
+import { PgManager } from '../replication/PgManager.js';
+
+export class PostgresModule extends replication.ReplicationModule<types.PostgresConnectionConfig> {
+  constructor() {
+    super({
+      name: 'Postgres',
+      type: types.POSTGRES_CONNECTION_TYPE,
+      configSchema: types.PostgresConnectionConfig
+    });
+  }
+
+  async initialize(context: system.ServiceContextContainer): Promise<void> {
+    await super.initialize(context);
+
+    // Record replicated bytes using global jpgwire metrics.
+    if (context.configuration.base_config.client_auth?.supabase) {
+      this.registerSupabaseAuth(context);
+    }
+
+    if (context.metrics) {
+      jpgwire.setMetricsRecorder({
+        addBytesRead(bytes) {
+          context.metrics!.data_replicated_bytes.add(bytes);
+        }
+      });
+    }
+  }
+
+  protected createRouteAPIAdapter(): api.RouteAPI {
+    return new PostgresRouteAPIAdapter(this.resolveConfig(this.decodedConfig!));
+  }
+
+  protected createReplicator(context: system.ServiceContext): replication.AbstractReplicator {
+    const normalisedConfig = this.resolveConfig(this.decodedConfig!);
+    const syncRuleProvider = new ConfigurationFileSyncRulesProvider(context.configuration.sync_rules);
+    const connectionFactory = new ConnectionManagerFactory(normalisedConfig);
+
+    return new WalStreamReplicator({
+      id: this.getDefaultId(normalisedConfig.database),
+      syncRuleProvider: syncRuleProvider,
+      storageEngine: context.storageEngine,
+      connectionFactory: connectionFactory,
+      rateLimiter: new PostgresErrorRateLimiter()
+    });
+  }
+
+  /**
+   * Combines base config with normalized connection settings
+   */
+  private resolveConfig(config: types.PostgresConnectionConfig): types.ResolvedConnectionConfig {
+    return {
+      ...config,
+      ...types.normalizeConnectionConfig(config)
+    };
+  }
+
+  async teardown(options: TearDownOptions): Promise<void> {
+    const normalisedConfig = this.resolveConfig(this.decodedConfig!);
+    const connectionManager = new PgManager(normalisedConfig, {
+      idleTimeout: 30_000,
+      maxSize: 1
+    });
+
+    try {
+      if (options.syncRules) {
+        // TODO: In the future, once we have more replication types, we will need to check if these syncRules are for Postgres
+        for (let syncRules of options.syncRules) {
+          try {
+            await cleanUpReplicationSlot(syncRules.slot_name, connectionManager.pool);
+          } catch (e) {
+            // Not really much we can do here for failures, most likely the database is no longer accessible
+            this.logger.warn(`Failed to fully clean up Postgres replication slot: ${syncRules.slot_name}`, e);
+          }
+        }
+      }
+    } finally {
+      await connectionManager.end();
+    }
+  }
+
+  // TODO: This should rather be done by registering the key collector in some kind of auth engine
+  private registerSupabaseAuth(context: system.ServiceContextContainer) {
+    const { configuration } = context;
+    // Register the Supabase key collector(s)
+    configuration.connections
+      ?.map((baseConfig) => {
+        if (baseConfig.type != types.POSTGRES_CONNECTION_TYPE) {
+          return;
+        }
+        try {
+          return this.resolveConfig(types.PostgresConnectionConfig.decode(baseConfig as any));
+        } catch (ex) {
+          this.logger.warn('Failed to decode configuration.', ex);
+        }
+      })
+      .filter((c) => !!c)
+      .forEach((config) => {
+        const keyCollector = new SupabaseKeyCollector(config!);
+        context.lifeCycleEngine.withLifecycle(keyCollector, {
+          // Close the internal pool
+          stop: (collector) => collector.shutdown()
+        });
+        configuration.client_keystore.collector.add(new auth.CachedKeyCollector(keyCollector));
+      });
+  }
+}

package/src/replication/ConnectionManagerFactory.ts

@@ -0,0 +1,28 @@
+import { PgManager } from './PgManager.js';
+import { NormalizedPostgresConnectionConfig } from '../types/types.js';
+import { PgPoolOptions } from '@powersync/service-jpgwire';
+import { logger } from '@powersync/lib-services-framework';
+
+export class ConnectionManagerFactory {
+  private readonly connectionManagers: PgManager[];
+  private readonly dbConnectionConfig: NormalizedPostgresConnectionConfig;
+
+  constructor(dbConnectionConfig: NormalizedPostgresConnectionConfig) {
+    this.dbConnectionConfig = dbConnectionConfig;
+    this.connectionManagers = [];
+  }
+
+  create(poolOptions: PgPoolOptions) {
+    const manager = new PgManager(this.dbConnectionConfig, poolOptions);
+    this.connectionManagers.push(manager);
+    return manager;
+  }
+
+  async shutdown() {
+    logger.info('Shutting down Postgres connection Managers...');
+    for (const manager of this.connectionManagers) {
+      await manager.end();
+    }
+    logger.info('Postgres connection Managers shutdown completed.');
+  }
+}

package/src/replication/PgManager.ts

@@ -0,0 +1,70 @@
+import * as pgwire from '@powersync/service-jpgwire';
+import { NormalizedPostgresConnectionConfig } from '../types/types.js';
+
+export class PgManager {
+  /**
+   * Do not use this for any transactions.
+   */
+  public readonly pool: pgwire.PgClient;
+
+  private connectionPromises: Promise<pgwire.PgConnection>[] = [];
+
+  constructor(public options: NormalizedPostgresConnectionConfig, public poolOptions: pgwire.PgPoolOptions) {
+    // The pool is lazy - no connections are opened until a query is performed.
+    this.pool = pgwire.connectPgWirePool(this.options, poolOptions);
+  }
+
+  public get connectionTag() {
+    return this.options.tag;
+  }
+
+  /**
+   * Create a new replication connection.
+   */
+  async replicationConnection(): Promise<pgwire.PgConnection> {
+    const p = pgwire.connectPgWire(this.options, { type: 'replication' });
+    this.connectionPromises.push(p);
+    return await p;
+  }
+
+  /**
+   * Create a new standard connection, used for initial snapshot.
+   *
+   * This connection must not be shared between multiple async contexts.
+   */
+  async snapshotConnection(): Promise<pgwire.PgConnection> {
+    const p = pgwire.connectPgWire(this.options, { type: 'standard' });
+    this.connectionPromises.push(p);
+    return await p;
+  }
+
+  async end(): Promise<void> {
+    for (let result of await Promise.allSettled([
+      this.pool.end(),
+      ...this.connectionPromises.map(async (promise) => {
+        const connection = await promise;
+        return await connection.end();
+      })
+    ])) {
+      // Throw the first error, if any
+      if (result.status == 'rejected') {
+        throw result.reason;
+      }
+    }
+  }
+
+  async destroy() {
+    this.pool.destroy();
+    for (let result of await Promise.allSettled([
+      ...this.connectionPromises.map(async (promise) => {
+        const connection = await promise;
+        return connection.destroy();
+      })
+    ])) {
+      // Throw the first error, if any
+      if (result.status == 'rejected') {
+        throw result.reason;
+      }
+    }
+  }
+}

package/src/replication/PgRelation.ts

@@ -0,0 +1,31 @@
+import { storage } from '@powersync/service-core';
+import { PgoutputRelation } from '@powersync/service-jpgwire';
+
+export type ReplicationIdentity = 'default' | 'nothing' | 'full' | 'index';
+
+export function getReplicaIdColumns(relation: PgoutputRelation): storage.ColumnDescriptor[] {
+  if (relation.replicaIdentity == 'nothing') {
+    return [];
+  } else {
+    return relation.columns
+      .filter((c) => (c.flags & 0b1) != 0)
+      .map((c) => ({ name: c.name, typeId: c.typeOid } satisfies storage.ColumnDescriptor));
+  }
+}
+export function getRelId(source: PgoutputRelation): number {
+  // Source types are wrong here
+  const relId = (source as any).relationOid as number;
+  if (!relId) {
+    throw new Error(`No relation id!`);
+  }
+  return relId;
+}
+
+export function getPgOutputRelation(source: PgoutputRelation): storage.SourceEntityDescriptor {
+  return {
+    name: source.name,
+    schema: source.schema,
+    objectId: getRelId(source),
+    replicationColumns: getReplicaIdColumns(source)
+  } satisfies storage.SourceEntityDescriptor;
+}

package/src/replication/PostgresErrorRateLimiter.ts

@@ -0,0 +1,44 @@
+import { setTimeout } from 'timers/promises';
+import { ErrorRateLimiter } from '@powersync/service-core';
+
+export class PostgresErrorRateLimiter implements ErrorRateLimiter {
+  nextAllowed: number = Date.now();
+
+  async waitUntilAllowed(options?: { signal?: AbortSignal | undefined } | undefined): Promise<void> {
+    const delay = Math.max(0, this.nextAllowed - Date.now());
+    // Minimum delay between connections, even without errors
+    this.setDelay(500);
+    await setTimeout(delay, undefined, { signal: options?.signal });
+  }
+
+  mayPing(): boolean {
+    return Date.now() >= this.nextAllowed;
+  }
+
+  reportError(e: any): void {
+    const message = (e.message as string) ?? '';
+    if (message.includes('password authentication failed')) {
+      // Wait 15 minutes, to avoid triggering Supabase's fail2ban
+      this.setDelay(900_000);
+    } else if (message.includes('ENOTFOUND')) {
+      // DNS lookup issue - incorrect URI or deleted instance
+      this.setDelay(120_000);
+    } else if (message.includes('ECONNREFUSED')) {
+      // Could be fail2ban or similar
+      this.setDelay(120_000);
+    } else if (
+      message.includes('Unable to do postgres query on ended pool') ||
+      message.includes('Postgres unexpectedly closed connection')
+    ) {
+      // Connection timed out - ignore / immediately retry
+      // We don't explicitly set the delay to 0, since there could have been another error that
+      // we need to respect.
+    } else {
+      this.setDelay(30_000);
+    }
+  }
+
+  private setDelay(delay: number) {
+    this.nextAllowed = Math.max(this.nextAllowed, Date.now() + delay);
+  }
+}