@powersync/service-module-postgres-storage 0.0.0-dev-20250116115804

package/src/storage/PostgresTestStorageFactoryGenerator.ts

@@ -0,0 +1,61 @@
+import { framework, PowerSyncMigrationManager, ServiceContext, TestStorageOptions } from '@powersync/service-core';
+import { PostgresMigrationAgent } from '../migrations/PostgresMigrationAgent.js';
+import { normalizePostgresStorageConfig, PostgresStorageConfigDecoded } from '../types/types.js';
+import { PostgresBucketStorageFactory } from './PostgresBucketStorageFactory.js';
+
+export type PostgresTestStorageOptions = {
+  url: string;
+  /**
+   * Vitest can cause issues when loading .ts files for migrations.
+   * This allows for providing a custom PostgresMigrationAgent.
+   */
+  migrationAgent?: (config: PostgresStorageConfigDecoded) => PostgresMigrationAgent;
+};
+
+export const PostgresTestStorageFactoryGenerator = (factoryOptions: PostgresTestStorageOptions) => {
+  return async (options?: TestStorageOptions) => {
+    try {
+      const migrationManager: PowerSyncMigrationManager = new framework.MigrationManager();
+
+      const BASE_CONFIG = {
+        type: 'postgresql' as const,
+        uri: factoryOptions.url,
+        sslmode: 'disable' as const
+      };
+
+      const TEST_CONNECTION_OPTIONS = normalizePostgresStorageConfig(BASE_CONFIG);
+
+      await using migrationAgent = factoryOptions.migrationAgent
+        ? factoryOptions.migrationAgent(BASE_CONFIG)
+        : new PostgresMigrationAgent(BASE_CONFIG);
+      migrationManager.registerMigrationAgent(migrationAgent);
+
+      const mockServiceContext = { configuration: { storage: BASE_CONFIG } } as unknown as ServiceContext;
+
+      if (!options?.doNotClear) {
+        await migrationManager.migrate({
+          direction: framework.migrations.Direction.Down,
+          migrationContext: {
+            service_context: mockServiceContext
+          }
+        });
+      }
+
+      await migrationManager.migrate({
+        direction: framework.migrations.Direction.Up,
+        migrationContext: {
+          service_context: mockServiceContext
+        }
+      });
+
+      return new PostgresBucketStorageFactory({
+        config: TEST_CONNECTION_OPTIONS,
+        slot_name_prefix: 'test_'
+      });
+    } catch (ex) {
+      // Vitest does not display these errors nicely when using the `await using` syntx
+      console.error(ex, ex.cause);
+      throw ex;
+    }
+  };
+};

package/src/storage/batch/OperationBatch.ts

@@ -0,0 +1,101 @@
+/**
+ * TODO share this implementation better in the core package.
+ * There are some subtle differences in this implementation.
+ */
+
+import { storage, utils } from '@powersync/service-core';
+import { RequiredOperationBatchLimits } from '../../types/types.js';
+
+/**
+ * Batch of input operations.
+ *
+ * We accumulate operations up to MAX_RECORD_BATCH_SIZE,
+ * then further split into sub-batches if MAX_CURRENT_DATA_BATCH_SIZE is exceeded.
+ */
+export class OperationBatch {
+  batch: RecordOperation[] = [];
+  currentSize: number = 0;
+
+  readonly maxBatchCount: number;
+  readonly maxRecordSize: number;
+  readonly maxCurrentDataBatchSize: number;
+
+  get length() {
+    return this.batch.length;
+  }
+
+  constructor(protected options: RequiredOperationBatchLimits) {
+    this.maxBatchCount = options.max_record_count;
+    this.maxRecordSize = options.max_estimated_size;
+    this.maxCurrentDataBatchSize = options.max_current_data_batch_size;
+  }
+
+  push(op: RecordOperation) {
+    this.batch.push(op);
+    this.currentSize += op.estimatedSize;
+  }
+
+  shouldFlush() {
+    return this.batch.length >= this.maxBatchCount || this.currentSize > this.maxCurrentDataBatchSize;
+  }
+
+  /**
+   *
+   * @param sizes Map of source key to estimated size of the current_data document, or undefined if current_data is not persisted.
+   *
+   */
+  *batched(sizes: Map<string, number> | undefined): Generator<RecordOperation[]> {
+    if (sizes == null) {
+      yield this.batch;
+      return;
+    }
+    let currentBatch: RecordOperation[] = [];
+    let currentBatchSize = 0;
+    for (let op of this.batch) {
+      const key = op.internalBeforeKey;
+      const size = sizes.get(key) ?? 0;
+      if (currentBatchSize + size > this.maxCurrentDataBatchSize && currentBatch.length > 0) {
+        yield currentBatch;
+        currentBatch = [];
+        currentBatchSize = 0;
+      }
+      currentBatchSize += size;
+      currentBatch.push(op);
+    }
+    if (currentBatch.length > 0) {
+      yield currentBatch;
+    }
+  }
+}
+
+export class RecordOperation {
+  public readonly afterId: storage.ReplicaId | null;
+  public readonly beforeId: storage.ReplicaId;
+  public readonly internalBeforeKey: string;
+  public readonly internalAfterKey: string | null;
+  public readonly estimatedSize: number;
+
+  constructor(public readonly record: storage.SaveOptions) {
+    const afterId = record.afterReplicaId ?? null;
+    const beforeId = record.beforeReplicaId ?? record.afterReplicaId;
+    this.afterId = afterId;
+    this.beforeId = beforeId;
+    this.internalBeforeKey = cacheKey(record.sourceTable.id, beforeId);
+    this.internalAfterKey = afterId ? cacheKey(record.sourceTable.id, afterId) : null;
+    this.estimatedSize = utils.estimateRowSize(record.before) + utils.estimateRowSize(record.after);
+  }
+}
+
+/**
+ * In-memory cache key - must not be persisted.
+ */
+export function cacheKey(sourceTableId: string, id: storage.ReplicaId) {
+  return encodedCacheKey(sourceTableId, storage.serializeReplicaId(id));
+}
+
+/**
+ * Calculates a cache key for a stored ReplicaId. This is usually stored as a bytea/Buffer.
+ */
+export function encodedCacheKey(sourceTableId: string, storedKey: Buffer) {
+  return `${sourceTableId}.${storedKey.toString('base64')}`;
+}