@powersync/lib-service-postgres 0.1.0

package/src/db/connection/ConnectionSlot.ts

@@ -0,0 +1,165 @@
+import * as framework from '@powersync/lib-services-framework';
+import * as pgwire from '@powersync/service-jpgwire';
+
+export interface NotificationListener extends framework.DisposableListener {
+  notification?: (payload: pgwire.PgNotification) => void;
+}
+
+export interface ConnectionSlotListener extends NotificationListener {
+  connectionAvailable?: () => void;
+  connectionError?: (exception: any) => void;
+  connectionCreated?: (connection: pgwire.PgConnection) => Promise<void>;
+}
+
+export type ConnectionLease = {
+  connection: pgwire.PgConnection;
+  release: () => void;
+};
+
+export type ConnectionSlotOptions = {
+  config: pgwire.NormalizedConnectionConfig;
+  notificationChannels?: string[];
+};
+
+export const MAX_CONNECTION_ATTEMPTS = 5;
+
+export class ConnectionSlot extends framework.DisposableObserver<ConnectionSlotListener> {
+  isAvailable: boolean;
+  isPoking: boolean;
+
+  closed: boolean;
+
+  protected connection: pgwire.PgConnection | null;
+  protected connectingPromise: Promise<pgwire.PgConnection> | null;
+
+  constructor(protected options: ConnectionSlotOptions) {
+    super();
+    this.isAvailable = false;
+    this.connection = null;
+    this.isPoking = false;
+    this.connectingPromise = null;
+    this.closed = false;
+  }
+
+  get isConnected() {
+    return !!this.connection;
+  }
+
+  protected async connect() {
+    this.connectingPromise = pgwire.connectPgWire(this.options.config, { type: 'standard' });
+    const connection = await this.connectingPromise;
+    this.connectingPromise = null;
+    await this.iterateAsyncListeners(async (l) => l.connectionCreated?.(connection));
+    if (this.hasNotificationListener()) {
+      await this.configureConnectionNotifications(connection);
+    }
+    return connection;
+  }
+
+  async [Symbol.asyncDispose]() {
+    this.closed = true;
+    const connection = this.connection ?? (await this.connectingPromise);
+    await connection?.end();
+    return super[Symbol.dispose]();
+  }
+
+  protected async configureConnectionNotifications(connection: pgwire.PgConnection) {
+    if (connection.onnotification == this.handleNotification || this.closed == true) {
+      // Already configured
+      return;
+    }
+
+    connection.onnotification = this.handleNotification;
+
+    for (const channelName of this.options.notificationChannels ?? []) {
+      await connection.query({
+        statement: `LISTEN ${channelName}`
+      });
+    }
+  }
+
+  registerListener(listener: Partial<ConnectionSlotListener>): () => void {
+    const dispose = super.registerListener(listener);
+    if (this.connection && this.hasNotificationListener()) {
+      this.configureConnectionNotifications(this.connection);
+    }
+    return () => {
+      dispose();
+      if (this.connection && !this.hasNotificationListener()) {
+        this.connection.onnotification = () => {};
+      }
+    };
+  }
+
+  protected handleNotification = (payload: pgwire.PgNotification) => {
+    if (!this.options.notificationChannels?.includes(payload.channel)) {
+      return;
+    }
+    this.iterateListeners((l) => l.notification?.(payload));
+  };
+
+  protected hasNotificationListener() {
+    return !!Object.values(this.listeners).find((l) => !!l.notification);
+  }
+
+  /**
+   * Test the connection if it can be reached.
+   */
+  async poke() {
+    if (this.isPoking || (this.isConnected && this.isAvailable == false) || this.closed) {
+      return;
+    }
+    this.isPoking = true;
+    for (let retryCounter = 0; retryCounter <= MAX_CONNECTION_ATTEMPTS; retryCounter++) {
+      try {
+        const connection = this.connection ?? (await this.connect());
+
+        await connection.query({
+          statement: 'SELECT 1'
+        });
+
+        if (!this.connection) {
+          this.connection = connection;
+          this.setAvailable();
+        } else if (this.isAvailable) {
+          this.iterateListeners((cb) => cb.connectionAvailable?.());
+        }
+
+        // Connection is alive and healthy
+        break;
+      } catch (ex) {
+        // Should be valid for all cases
+        this.isAvailable = false;
+        if (this.connection) {
+          this.connection.onnotification = () => {};
+          this.connection.destroy();
+          this.connection = null;
+        }
+        if (retryCounter >= MAX_CONNECTION_ATTEMPTS) {
+          this.iterateListeners((cb) => cb.connectionError?.(ex));
+        }
+      }
+    }
+    this.isPoking = false;
+  }
+
+  protected setAvailable() {
+    this.isAvailable = true;
+    this.iterateListeners((l) => l.connectionAvailable?.());
+  }
+
+  lock(): ConnectionLease | null {
+    if (!this.isAvailable || !this.connection || this.closed) {
+      return null;
+    }
+
+    this.isAvailable = false;
+
+    return {
+      connection: this.connection,
+      release: () => {
+        this.setAvailable();
+      }
+    };
+  }
+}

package/src/db/connection/DatabaseClient.ts

@@ -0,0 +1,261 @@
+import * as lib_postgres from '@powersync/lib-service-postgres';
+import * as pgwire from '@powersync/service-jpgwire';
+import pDefer, { DeferredPromise } from 'p-defer';
+import { AbstractPostgresConnection, sql } from './AbstractPostgresConnection.js';
+import { ConnectionLease, ConnectionSlot, NotificationListener } from './ConnectionSlot.js';
+import { WrappedConnection } from './WrappedConnection.js';
+
+export type DatabaseClientOptions = {
+  config: lib_postgres.NormalizedBasePostgresConnectionConfig;
+  /**
+   * Optional schema which will be used as the default search path
+   */
+  schema?: string;
+  /**
+   * Notification channels to listen to.
+   */
+  notificationChannels?: string[];
+};
+
+export type DatabaseClientListener = NotificationListener & {
+  connectionCreated?: (connection: pgwire.PgConnection) => Promise<void>;
+};
+
+export const TRANSACTION_CONNECTION_COUNT = 5;
+
+/**
+ * This provides access to Postgres via the PGWire library.
+ * A connection pool is used for individual query executions while
+ * a custom pool of connections is available for transactions or other operations
+ * which require being executed on the same connection.
+ */
+export class DatabaseClient extends AbstractPostgresConnection<DatabaseClientListener> {
+  closed: boolean;
+
+  protected pool: pgwire.PgClient;
+  protected connections: ConnectionSlot[];
+
+  protected initialized: Promise<void>;
+  protected queue: DeferredPromise<ConnectionLease>[];
+
+  constructor(protected options: DatabaseClientOptions) {
+    super();
+    this.closed = false;
+    this.pool = pgwire.connectPgWirePool(options.config);
+    this.connections = Array.from({ length: TRANSACTION_CONNECTION_COUNT }, () => {
+      const slot = new ConnectionSlot({ config: options.config, notificationChannels: options.notificationChannels });
+      slot.registerListener({
+        connectionAvailable: () => this.processConnectionQueue(),
+        connectionError: (ex) => this.handleConnectionError(ex),
+        connectionCreated: (connection) => this.iterateAsyncListeners(async (l) => l.connectionCreated?.(connection))
+      });
+      return slot;
+    });
+    this.queue = [];
+    this.initialized = this.initialize();
+  }
+
+  protected get baseConnection() {
+    return this.pool;
+  }
+
+  protected get schemaStatement() {
+    const { schema } = this.options;
+    if (!schema) {
+      return;
+    }
+    return {
+      statement: `SET search_path TO ${schema};`
+    };
+  }
+
+  registerListener(listener: Partial<DatabaseClientListener>): () => void {
+    let disposeNotification: (() => void) | null = null;
+    if ('notification' in listener) {
+      // Pass this on to the first connection slot
+      // It will only actively listen on the connection once a listener has been registered
+      disposeNotification = this.connections[0].registerListener({
+        notification: listener.notification
+      });
+      this.pokeSlots();
+
+      delete listener['notification'];
+    }
+
+    const superDispose = super.registerListener(listener);
+    return () => {
+      disposeNotification?.();
+      superDispose();
+    };
+  }
+
+  query(script: string, options?: pgwire.PgSimpleQueryOptions): Promise<pgwire.PgResult>;
+  query(...args: pgwire.Statement[]): Promise<pgwire.PgResult>;
+  async query(...args: any[]): Promise<pgwire.PgResult> {
+    await this.initialized;
+    /**
+     * There is no direct way to set the default schema with pgwire.
+     * This hack uses multiple statements in order to always ensure the
+     * appropriate connection (in the pool) uses the correct schema.
+     */
+    const { schemaStatement } = this;
+    if (typeof args[0] == 'object' && schemaStatement) {
+      args.unshift(schemaStatement);
+    } else if (typeof args[0] == 'string' && schemaStatement) {
+      args[0] = `${schemaStatement.statement}; ${args[0]}`;
+    }
+
+    // Retry pool queries. Note that we can't retry queries in a transaction
+    // since a failed query will end the transaction.
+    return lib_postgres.retriedQuery(this.pool, ...args);
+  }
+
+  async *stream(...args: pgwire.Statement[]): AsyncIterableIterator<pgwire.PgChunk> {
+    await this.initialized;
+    const { schemaStatement } = this;
+    if (schemaStatement) {
+      args.unshift(schemaStatement);
+    }
+    yield* super.stream(...args);
+  }
+
+  async lockConnection<T>(callback: (db: WrappedConnection) => Promise<T>): Promise<T> {
+    const { connection, release } = await this.requestConnection();
+
+    await this.setSchema(connection);
+
+    try {
+      return await callback(new WrappedConnection(connection));
+    } finally {
+      release();
+    }
+  }
+
+  async transaction<T>(tx: (db: WrappedConnection) => Promise<T>): Promise<T> {
+    return this.lockConnection(async (db) => {
+      try {
+        await db.query(sql`BEGIN`);
+        const result = await tx(db);
+        await db.query(sql`COMMIT`);
+        return result;
+      } catch (ex) {
+        await db.query(sql`ROLLBACK`);
+        throw ex;
+      }
+    });
+  }
+
+  /**
+   * Use the `powersync` schema as the default when resolving table names
+   */
+  protected async setSchema(client: pgwire.PgClient) {
+    const { schemaStatement } = this;
+    if (!schemaStatement) {
+      return;
+    }
+    await client.query(schemaStatement);
+  }
+
+  protected async initialize() {
+    const { schema } = this.options;
+    if (schema) {
+      // First check if it exists
+      const exists = await this.pool.query(sql`
+        SELECT
+          schema_name
+        FROM
+          information_schema.schemata
+        WHERE
+          schema_name = ${{ type: 'varchar', value: schema }};
+      `);
+
+      if (exists.rows.length) {
+        return;
+      }
+      // Create the schema if it doesn't exist
+      await this.pool.query({ statement: `CREATE SCHEMA IF NOT EXISTS ${this.options.schema}` });
+    }
+  }
+
+  protected async requestConnection(): Promise<ConnectionLease> {
+    if (this.closed) {
+      throw new Error('Database client is closed');
+    }
+
+    await this.initialized;
+
+    // Queue the operation
+    const deferred = pDefer<ConnectionLease>();
+    this.queue.push(deferred);
+
+    this.pokeSlots();
+
+    return deferred.promise;
+  }
+
+  protected pokeSlots() {
+    // Poke the slots to check if they are alive
+    for (const slot of this.connections) {
+      // No need to await this. Errors are reported asynchronously
+      slot.poke();
+    }
+  }
+
+  protected leaseConnectionSlot(): ConnectionLease | null {
+    const availableSlots = this.connections.filter((s) => s.isAvailable);
+    for (const slot of availableSlots) {
+      const lease = slot.lock();
+      if (lease) {
+        return lease;
+      }
+      // Possibly some contention detected, keep trying
+    }
+    return null;
+  }
+
+  protected processConnectionQueue() {
+    if (this.closed && this.queue.length) {
+      for (const q of this.queue) {
+        q.reject(new Error('Database has closed while waiting for a connection'));
+      }
+      this.queue = [];
+    }
+
+    if (this.queue.length) {
+      const lease = this.leaseConnectionSlot();
+      if (lease) {
+        const deferred = this.queue.shift()!;
+        deferred.resolve(lease);
+      }
+    }
+  }
+
+  /**
+   * Reports connection errors which might occur from bad configuration or
+   * a server which is no longer available.
+   * This fails all pending requests.
+   */
+  protected handleConnectionError(exception: any) {
+    for (const q of this.queue) {
+      q.reject(exception);
+    }
+    this.queue = [];
+  }
+
+  async [Symbol.asyncDispose]() {
+    await this.initialized;
+    this.closed = true;
+
+    for (const c of this.connections) {
+      await c[Symbol.asyncDispose]();
+    }
+
+    await this.pool.end();
+
+    // Reject all remaining items
+    for (const q of this.queue) {
+      q.reject(new Error(`Database is disposed`));
+    }
+    this.queue = [];
+  }
+}

package/src/db/connection/WrappedConnection.ts

@@ -0,0 +1,11 @@
+import * as pgwire from '@powersync/service-jpgwire';
+import { AbstractPostgresConnection } from './AbstractPostgresConnection.js';
+
+/**
+ * Provides helper functionality to transaction contexts given an existing PGWire connection
+ */
+export class WrappedConnection extends AbstractPostgresConnection {
+  constructor(protected baseConnection: pgwire.PgConnection) {
+    super();
+  }
+}

package/src/db/db-index.ts

@@ -0,0 +1,4 @@
+export * from './connection/AbstractPostgresConnection.js';
+export * from './connection/ConnectionSlot.js';
+export * from './connection/DatabaseClient.js';
+export * from './connection/WrappedConnection.js';

package/src/index.ts

@@ -0,0 +1,11 @@
+export * from './db/db-index.js';
+export * as db from './db/db-index.js';
+
+export * from './locks/locks-index.js';
+export * as locks from './locks/locks-index.js';
+
+export * from './types/types.js';
+export * as types from './types/types.js';
+
+export * from './utils/utils-index.js';
+export * as utils from './utils/utils-index.js';

package/src/locks/PostgresLockManager.ts

@@ -0,0 +1,128 @@
+import * as framework from '@powersync/lib-services-framework';
+import { v4 as uuidv4 } from 'uuid';
+import { DatabaseClient, sql } from '../db/db-index.js';
+
+const DEFAULT_LOCK_TIMEOUT = 60_000; // 1 minute
+
+export interface PostgresLockManagerParams extends framework.locks.LockManagerParams {
+  db: DatabaseClient;
+}
+
+export class PostgresLockManager extends framework.locks.AbstractLockManager {
+  constructor(protected params: PostgresLockManagerParams) {
+    super(params);
+  }
+
+  protected get db() {
+    return this.params.db;
+  }
+
+  get timeout() {
+    return this.params.timeout ?? DEFAULT_LOCK_TIMEOUT;
+  }
+
+  get name() {
+    return this.params.name;
+  }
+
+  async init() {
+    /**
+     * Locks are required for migrations, which means this table can't be
+     * created inside a migration. This ensures the locks table is present.
+     */
+    await this.db.query(sql`
+      CREATE TABLE IF NOT EXISTS locks (
+        name TEXT PRIMARY KEY,
+        lock_id UUID NOT NULL,
+        ts TIMESTAMPTZ NOT NULL
+      );
+    `);
+  }
+
+  protected async acquireHandle(): Promise<framework.LockHandle | null> {
+    const id = await this._acquireId();
+    if (!id) {
+      return null;
+    }
+    return {
+      refresh: () => this.refreshHandle(id),
+      release: () => this.releaseHandle(id)
+    };
+  }
+
+  protected async _acquireId(): Promise<string | null> {
+    const now = new Date();
+    const nowISO = now.toISOString();
+    const expiredTs = new Date(now.getTime() - this.timeout).toISOString();
+    const lockId = uuidv4();
+
+    try {
+      // Attempt to acquire or refresh the lock
+      const res = await this.db.queryRows<{ lock_id: string }>(sql`
+        INSERT INTO
+          locks (name, lock_id, ts)
+        VALUES
+          (
+            ${{ type: 'varchar', value: this.name }},
+            ${{ type: 'uuid', value: lockId }},
+            ${{ type: 1184, value: nowISO }}
+          )
+        ON CONFLICT (name) DO UPDATE
+        SET
+          lock_id = CASE
+            WHEN locks.ts <= ${{ type: 1184, value: expiredTs }} THEN ${{ type: 'uuid', value: lockId }}
+            ELSE locks.lock_id
+          END,
+          ts = CASE
+            WHEN locks.ts <= ${{ type: 1184, value: expiredTs }} THEN ${{
+          type: 1184,
+          value: nowISO
+        }}
+            ELSE locks.ts
+          END
+        RETURNING
+          lock_id;
+      `);
+
+      if (res.length == 0 || res[0].lock_id !== lockId) {
+        // Lock is active and could not be acquired
+        return null;
+      }
+
+      return lockId;
+    } catch (err) {
+      console.error('Error acquiring lock:', err);
+      throw err;
+    }
+  }
+
+  protected async refreshHandle(lockId: string) {
+    const res = await this.db.query(sql`
+      UPDATE locks
+      SET
+        ts = ${{ type: 1184, value: new Date().toISOString() }}
+      WHERE
+        lock_id = ${{ type: 'uuid', value: lockId }}
+      RETURNING
+        lock_id;
+    `);
+
+    if (res.rows.length === 0) {
+      throw new Error('Lock not found, could not refresh');
+    }
+  }
+
+  protected async releaseHandle(lockId: string) {
+    const res = await this.db.query(sql`
+      DELETE FROM locks
+      WHERE
+        lock_id = ${{ type: 'uuid', value: lockId }}
+      RETURNING
+        lock_id;
+    `);
+
+    if (res.rows.length == 0) {
+      throw new Error('Lock not found, could not release');
+    }
+  }
+}

package/src/locks/locks-index.ts

@@ -0,0 +1 @@
+export * from './PostgresLockManager.js';