Move input conversion to source stream

Author: simolus3

.changeset/dry-pots-leave.md

@@ -0,0 +1,9 @@
+---
+'@powersync/service-module-postgres': patch
+'@powersync/service-module-mongodb': patch
+'@powersync/service-core': patch
+'@powersync/service-module-mysql': patch
+'@powersync/service-sync-rules': patch
+---
+
+Correctly handle custom types in primary keys.

modules/module-mongodb-storage/src/storage/implementation/MongoBucketBatch.ts

@@ -319,8 +319,7 @@ export class MongoBucketBatch
     const record = operation.record;
     const beforeId = operation.beforeId;
     const afterId = operation.afterId;
-    let sourceAfter = record.after;
-    let after = sourceAfter && this.sync_rules.applyRowContext(sourceAfter);
+    let after = record.after;
     const sourceTable = record.sourceTable;
 
     let existing_buckets: CurrentBucket[] = [];

modules/module-mongodb/src/replication/ChangeStream.ts

@@ -481,7 +481,7 @@ export class ChangeStream {
       // Pre-fetch next batch, so that we can read and write concurrently
       nextChunkPromise = query.nextChunk();
       for (let document of docBatch) {
-        const record = constructAfterRecord(document);
+        const record = this.constructAfterRecord(document);
 
         // This auto-flushes when the batch reaches its size limit
         await batch.save({
@@ -619,6 +619,11 @@ export class ChangeStream {
     return result.table;
   }
 
+  private constructAfterRecord(document: mongo.Document): SqliteRow {
+    const inputRow = constructAfterRecord(document);
+    return this.sync_rules.applyRowContext<never>(inputRow);
+  }
+
   async writeChange(
     batch: storage.BucketStorageBatch,
     table: storage.SourceTable,
@@ -631,7 +636,7 @@ export class ChangeStream {
 
     this.metrics.getCounter(ReplicationMetric.ROWS_REPLICATED).add(1);
     if (change.operationType == 'insert') {
-      const baseRecord = constructAfterRecord(change.fullDocument);
+      const baseRecord = this.constructAfterRecord(change.fullDocument);
       return await batch.save({
         tag: SaveOperationTag.INSERT,
         sourceTable: table,
@@ -650,7 +655,7 @@ export class ChangeStream {
           beforeReplicaId: change.documentKey._id
         });
       }
-      const after = constructAfterRecord(change.fullDocument!);
+      const after = this.constructAfterRecord(change.fullDocument!);
       return await batch.save({
         tag: SaveOperationTag.UPDATE,
         sourceTable: table,

modules/module-mongodb/src/replication/MongoRelation.ts

@@ -1,16 +1,13 @@
 import { mongo } from '@powersync/lib-service-mongodb';
 import { storage } from '@powersync/service-core';
-import { JSONBig, JsonContainer } from '@powersync/service-jsonbig';
+import { JsonContainer } from '@powersync/service-jsonbig';
 import {
   CompatibilityContext,
   CustomArray,
   CustomObject,
   CustomSqliteValue,
-  DatabaseInputValue,
   SqliteInputRow,
   SqliteInputValue,
-  SqliteRow,
-  SqliteValue,
   DateTimeValue
 } from '@powersync/service-sync-rules';
 

modules/module-mysql/src/replication/BinLogStream.ts

@@ -329,7 +329,7 @@ export class BinLogStream {
         throw new ReplicationAssertionError(`No 'fields' event emitted`);
       }
 
-      const record = common.toSQLiteRow(row, columns!);
+      const record = this.toSQLiteRow(row, columns!);
       await batch.save({
         tag: storage.SaveOperationTag.INSERT,
         sourceTable: table,
@@ -596,14 +596,19 @@ export class BinLogStream {
     return null;
   }
 
+  private toSQLiteRow(row: Record<string, any>, columns: Map<string, ColumnDescriptor>): sync_rules.SqliteRow {
+    const inputRecord = common.toSQLiteRow(row, columns);
+    return this.syncRules.applyRowContext<never>(inputRecord);
+  }
+
   private async writeChange(
     batch: storage.BucketStorageBatch,
     payload: WriteChangePayload
   ): Promise<storage.FlushedResult | null> {
     switch (payload.type) {
       case storage.SaveOperationTag.INSERT:
         this.metrics.getCounter(ReplicationMetric.ROWS_REPLICATED).add(1);
-        const record = common.toSQLiteRow(payload.row, payload.columns);
+        const record = this.toSQLiteRow(payload.row, payload.columns);
         return await batch.save({
           tag: storage.SaveOperationTag.INSERT,
           sourceTable: payload.sourceTable,
@@ -617,9 +622,9 @@ export class BinLogStream {
         // The previous row may be null if the replica id columns are unchanged.
         // It's fine to treat that the same as an insert.
         const beforeUpdated = payload.previous_row
-          ? common.toSQLiteRow(payload.previous_row, payload.columns)
+          ? this.toSQLiteRow(payload.previous_row, payload.columns)
           : undefined;
-        const after = common.toSQLiteRow(payload.row, payload.columns);
+        const after = this.toSQLiteRow(payload.row, payload.columns);
 
         return await batch.save({
           tag: storage.SaveOperationTag.UPDATE,
@@ -634,7 +639,7 @@ export class BinLogStream {
 
       case storage.SaveOperationTag.DELETE:
         this.metrics.getCounter(ReplicationMetric.ROWS_REPLICATED).add(1);
-        const beforeDeleted = common.toSQLiteRow(payload.row, payload.columns);
+        const beforeDeleted = this.toSQLiteRow(payload.row, payload.columns);
 
         return await batch.save({
           tag: storage.SaveOperationTag.DELETE,

modules/module-postgres-storage/src/storage/batch/PostgresBucketBatch.ts

@@ -687,8 +687,7 @@ export class PostgresBucketBatch
     // We store bytea colums for source keys
     const beforeId = operation.beforeId;
     const afterId = operation.afterId;
-    let sourceAfter = record.after;
-    let after = sourceAfter && this.sync_rules.applyRowContext(sourceAfter);
+    let after = record.after;
     const sourceTable = record.sourceTable;
 
     let existingBuckets: CurrentBucket[] = [];

modules/module-postgres/src/replication/WalStream.ts

@@ -25,8 +25,11 @@ import {
   CompatibilityContext,
   DatabaseInputRow,
   SqliteInputRow,
+  SqliteInputValue,
+  SqliteRow,
   SqlSyncRules,
   TablePattern,
+  ToastableSqliteRow,
   toSyncRulesRow
 } from '@powersync/service-sync-rules';
 
@@ -635,7 +638,8 @@ WHERE  oid = $1::regclass`,
           hasRemainingData = true;
         }
 
-        for (const record of WalStream.getQueryData(rows)) {
+        for (const inputRecord of WalStream.getQueryData(rows)) {
+          const record = this.syncRulesRecord(inputRecord);
           // This auto-flushes when the batch reaches its size limit
           await batch.save({
             tag: storage.SaveOperationTag.INSERT,
@@ -787,6 +791,20 @@ WHERE  oid = $1::regclass`,
     return table;
   }
 
+  private syncRulesRecord(row: SqliteInputRow): SqliteRow;
+  private syncRulesRecord(row: SqliteInputRow | undefined): SqliteRow | undefined;
+
+  private syncRulesRecord(row: SqliteInputRow | undefined): SqliteRow | undefined {
+    if (row == null) {
+      return undefined;
+    }
+    return this.sync_rules.applyRowContext<never>(row);
+  }
+
+  private toastableSyncRulesRecord(row: ToastableSqliteRow<SqliteInputValue>): ToastableSqliteRow {
+    return this.sync_rules.applyRowContext(row);
+  }
+
   async writeChange(
     batch: storage.BucketStorageBatch,
     msg: pgwire.PgoutputMessage
@@ -803,7 +821,7 @@ WHERE  oid = $1::regclass`,
 
       if (msg.tag == 'insert') {
         this.metrics.getCounter(ReplicationMetric.ROWS_REPLICATED).add(1);
-        const baseRecord = this.connections.types.constructAfterRecord(msg);
+        const baseRecord = this.syncRulesRecord(this.connections.types.constructAfterRecord(msg));
         return await batch.save({
           tag: storage.SaveOperationTag.INSERT,
           sourceTable: table,
@@ -816,8 +834,8 @@ WHERE  oid = $1::regclass`,
         this.metrics.getCounter(ReplicationMetric.ROWS_REPLICATED).add(1);
         // "before" may be null if the replica id columns are unchanged
         // It's fine to treat that the same as an insert.
-        const before = this.connections.types.constructBeforeRecord(msg);
-        const after = this.connections.types.constructAfterRecord(msg);
+        const before = this.syncRulesRecord(this.connections.types.constructBeforeRecord(msg));
+        const after = this.toastableSyncRulesRecord(this.connections.types.constructAfterRecord(msg));
         return await batch.save({
           tag: storage.SaveOperationTag.UPDATE,
           sourceTable: table,
@@ -828,7 +846,7 @@ WHERE  oid = $1::regclass`,
         });
       } else if (msg.tag == 'delete') {
         this.metrics.getCounter(ReplicationMetric.ROWS_REPLICATED).add(1);
-        const before = this.connections.types.constructBeforeRecord(msg)!;
+        const before = this.syncRulesRecord(this.connections.types.constructBeforeRecord(msg)!);
 
         return await batch.save({
           tag: storage.SaveOperationTag.DELETE,

modules/module-postgres/test/src/wal_stream.test.ts

@@ -348,4 +348,28 @@ config:
     const data = await context.getBucketData('1#stream|0[]');
     expect(data).toMatchObject([putOp('test_data', { id: 't1', description: '{"foo":1,"bar":2}' })]);
   });
+
+  test('custom types in primary key', async () => {
+    await using context = await WalStreamTestContext.open(factory);
+
+    await context.updateSyncRules(`
+streams:
+  stream:
+    query: SELECT id, * FROM "test_data"
+
+config:
+  edition: 2
+`);
+
+    const { pool } = context;
+    await pool.query(`DROP TABLE IF EXISTS test_data`);
+    await pool.query(`CREATE DOMAIN test_id AS TEXT;`);
+    await pool.query(`CREATE TABLE test_data(id test_id primary key);`);
+
+    await context.initializeReplication();
+    await pool.query(`INSERT INTO test_data(id) VALUES ('t1')`);
+
+    const data = await context.getBucketData('1#stream|0[]');
+    expect(data).toMatchObject([putOp('test_data', { id: 't1' })]);
+  });
 }

packages/service-core/src/storage/BucketStorageBatch.ts

@@ -1,11 +1,5 @@
 import { ObserverClient } from '@powersync/lib-services-framework';
-import {
-  EvaluatedParameters,
-  EvaluatedRow,
-  SqliteInputRow,
-  SqliteRow,
-  ToastableSqliteRow
-} from '@powersync/service-sync-rules';
+import { EvaluatedParameters, EvaluatedRow, SqliteRow, ToastableSqliteRow } from '@powersync/service-sync-rules';
 import { BSON } from 'bson';
 import { ReplicationEventPayload } from './ReplicationEventPayload.js';
 import { SourceTable, TableSnapshotStatus } from './SourceTable.js';
@@ -138,7 +132,7 @@ export interface SaveInsert {
   sourceTable: SourceTable;
   before?: undefined;
   beforeReplicaId?: undefined;
-  after: SqliteInputRow;
+  after: SqliteRow;
   afterReplicaId: ReplicaId;
 }
 
@@ -149,7 +143,7 @@ export interface SaveUpdate {
   /**
    * This is only present when the id has changed, and will only contain replica identity columns.
    */
-  before?: SqliteInputRow;
+  before?: SqliteRow;
   beforeReplicaId?: ReplicaId;
 
   /**
@@ -164,7 +158,7 @@ export interface SaveUpdate {
 export interface SaveDelete {
   tag: SaveOperationTag.DELETE;
   sourceTable: SourceTable;
-  before?: SqliteInputRow;
+  before?: SqliteRow;
   beforeReplicaId: ReplicaId;
   after?: undefined;
   afterReplicaId?: undefined;

packages/service-core/src/util/utils.ts

@@ -182,7 +182,7 @@ export function uuidForRowBson(row: sync_rules.SqliteRow): bson.UUID {
   return new bson.UUID(uuid.v5(repr, ID_NAMESPACE, buffer));
 }
 
-export function hasToastedValues(row: sync_rules.ToastableSqliteRow) {
+export function hasToastedValues<V>(row: sync_rules.ToastableSqliteRow<V>) {
   for (let key in row) {
     if (typeof row[key] == 'undefined') {
       return true;
@@ -196,10 +196,10 @@ export function hasToastedValues(row: sync_rules.ToastableSqliteRow) {
  *
  * If we don't store data, we assume we always have a complete row.
  */
-export function isCompleteRow(
+export function isCompleteRow<V>(
   storeData: boolean,
-  row: sync_rules.ToastableSqliteRow
-): row is sync_rules.SqliteInputRow {
+  row: sync_rules.ToastableSqliteRow<V>
+): row is sync_rules.SqliteRow<V> {
   if (!storeData) {
     // Assume the row is complete - no need to check
     return true;