Fix postgres storage truncate; typed SourceTable id.

Author: rkistner

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

@@ -29,7 +29,7 @@ import { MongoIdSequence } from './MongoIdSequence.js';
 import { batchCreateCustomWriteCheckpoints } from './MongoWriteCheckpointAPI.js';
 import { cacheKey, OperationBatch, RecordOperation } from './OperationBatch.js';
 import { PersistedBatch } from './PersistedBatch.js';
-import { idPrefixFilter } from '../../utils/util.js';
+import { idPrefixFilter, mongoTableId } from '../../utils/util.js';
 
 /**
  * 15MB
@@ -211,7 +211,7 @@ export class MongoBucketBatch
       // the order of processing, which then becomes really tricky to manage.
       // This now takes 2+ queries, but doesn't have any issues with order of operations.
       const sizeLookups: SourceKey[] = batch.batch.map((r) => {
-        return { g: this.group_id, t: r.record.sourceTable.id, k: r.beforeId };
+        return { g: this.group_id, t: mongoTableId(r.record.sourceTable.id), k: r.beforeId };
       });
 
       sizes = new Map<string, number>();
@@ -252,7 +252,7 @@ export class MongoBucketBatch
         continue;
       }
       const lookups: SourceKey[] = b.map((r) => {
-        return { g: this.group_id, t: r.record.sourceTable.id, k: r.beforeId };
+        return { g: this.group_id, t: mongoTableId(r.record.sourceTable.id), k: r.beforeId };
       });
       let current_data_lookup = new Map<string, CurrentDataDocument>();
       // With skipExistingRows, we only need to know whether or not the row exists.
@@ -326,7 +326,7 @@ export class MongoBucketBatch
     let existing_lookups: bson.Binary[] = [];
     let new_lookups: bson.Binary[] = [];
 
-    const before_key: SourceKey = { g: this.group_id, t: record.sourceTable.id, k: beforeId };
+    const before_key: SourceKey = { g: this.group_id, t: mongoTableId(record.sourceTable.id), k: beforeId };
 
     if (this.skipExistingRows) {
       if (record.tag == SaveOperationTag.INSERT) {
@@ -538,7 +538,7 @@ export class MongoBucketBatch
     // 5. TOAST: Update current data and bucket list.
     if (afterId) {
       // Insert or update
-      const after_key: SourceKey = { g: this.group_id, t: sourceTable.id, k: afterId };
+      const after_key: SourceKey = { g: this.group_id, t: sourceTable.id as bson.ObjectId, k: afterId };
       batch.upsertCurrentData(after_key, {
         data: afterData,
         buckets: new_buckets,
@@ -957,7 +957,7 @@ export class MongoBucketBatch
 
     await this.withTransaction(async () => {
       for (let table of sourceTables) {
-        await this.db.source_tables.deleteOne({ _id: table.id });
+        await this.db.source_tables.deleteOne({ _id: mongoTableId(table.id) });
       }
     });
     return result;
@@ -992,7 +992,7 @@ export class MongoBucketBatch
     while (lastBatchCount == BATCH_LIMIT) {
       await this.withReplicationTransaction(`Truncate ${sourceTable.qualifiedName}`, async (session, opSeq) => {
         const current_data_filter: mongo.Filter<CurrentDataDocument> = {
-          _id: idPrefixFilter<SourceKey>({ g: this.group_id, t: sourceTable.id }, ['k']),
+          _id: idPrefixFilter<SourceKey>({ g: this.group_id, t: mongoTableId(sourceTable.id) }, ['k']),
           // Skip soft-deleted data
           pending_delete: { $exists: false }
         };
@@ -1052,7 +1052,7 @@ export class MongoBucketBatch
 
     await this.withTransaction(async () => {
       await this.db.source_tables.updateOne(
-        { _id: table.id },
+        { _id: mongoTableId(table.id) },
         {
           $set: {
             snapshot_status: {
@@ -1103,7 +1103,7 @@ export class MongoBucketBatch
 
   async markTableSnapshotDone(tables: storage.SourceTable[], no_checkpoint_before_lsn?: string) {
     const session = this.session;
-    const ids = tables.map((table) => table.id);
+    const ids = tables.map((table) => mongoTableId(table.id));
 
     await this.withTransaction(async () => {
       await this.db.source_tables.updateMany(

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

@@ -2,6 +2,7 @@ import { ToastableSqliteRow } from '@powersync/service-sync-rules';
 import * as bson from 'bson';
 
 import { storage } from '@powersync/service-core';
+import { mongoTableId } from '../storage-index.js';
 
 /**
  * Maximum number of operations in a batch.
@@ -86,8 +87,8 @@ export class RecordOperation {
     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.internalBeforeKey = cacheKey(mongoTableId(record.sourceTable.id), beforeId);
+    this.internalAfterKey = afterId ? cacheKey(mongoTableId(record.sourceTable.id), afterId) : null;
 
     this.estimatedSize = estimateRowSize(record.before) + estimateRowSize(record.after);
   }

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

@@ -16,7 +16,7 @@ import {
   CurrentDataDocument,
   SourceKey
 } from './models.js';
-import { replicaIdToSubkey } from '../../utils/util.js';
+import { mongoTableId, replicaIdToSubkey } from '../../utils/util.js';
 
 /**
  * Maximum size of operations we write in a single transaction.
@@ -132,7 +132,7 @@ export class PersistedBatch {
               o: op_id
             },
             op: 'PUT',
-            source_table: options.table.id,
+            source_table: mongoTableId(options.table.id),
             source_key: options.sourceKey,
             table: k.table,
             row_id: k.id,
@@ -159,7 +159,7 @@ export class PersistedBatch {
               o: op_id
             },
             op: 'REMOVE',
-            source_table: options.table.id,
+            source_table: mongoTableId(options.table.id),
             source_key: options.sourceKey,
             table: bd.table,
             row_id: bd.id,
@@ -208,7 +208,7 @@ export class PersistedBatch {
             _id: op_id,
             key: {
               g: this.group_id,
-              t: sourceTable.id,
+              t: mongoTableId(sourceTable.id),
               k: sourceKey
             },
             lookup: binLookup,
@@ -230,7 +230,7 @@ export class PersistedBatch {
             _id: op_id,
             key: {
               g: this.group_id,
-              t: sourceTable.id,
+              t: mongoTableId(sourceTable.id),
               k: sourceKey
             },
             lookup: lookup,

modules/module-mongodb-storage/src/utils/util.ts

@@ -92,17 +92,32 @@ export function mapOpEntry(row: BucketDataDocument): utils.OplogEntry {
   }
 }
 
-export function replicaIdToSubkey(table: bson.ObjectId, id: storage.ReplicaId): string {
+export function replicaIdToSubkey(table: storage.SourceTableId, id: storage.ReplicaId): string {
   if (storage.isUUID(id)) {
     // Special case for UUID for backwards-compatiblity
-    return `${table.toHexString()}/${id.toHexString()}`;
+    return `${tableIdString(table)}/${id.toHexString()}`;
   } else {
     // Hashed UUID from the table and id
     const repr = bson.serialize({ table, id });
     return uuid.v5(repr, utils.ID_NAMESPACE);
   }
 }
 
+export function mongoTableId(table: storage.SourceTableId): bson.ObjectId {
+  if (typeof table == 'string') {
+    throw new ServiceAssertionError(`Got string table id, expected ObjectId`);
+  }
+  return table;
+}
+
+function tableIdString(table: storage.SourceTableId) {
+  if (typeof table == 'string') {
+    return table;
+  } else {
+    return table.toHexString();
+  }
+}
+
 export function setSessionSnapshotTime(session: mongo.ClientSession, time: bson.Timestamp) {
   // This is a workaround for the lack of direct support for snapshot reads in the MongoDB driver.
   if (!session.snapshotEnabled) {

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

@@ -5,6 +5,7 @@
 
 import { storage, utils } from '@powersync/service-core';
 import { RequiredOperationBatchLimits } from '../../types/types.js';
+import { postgresTableId } from './PostgresPersistedBatch.js';
 
 /**
  * Batch of input operations.
@@ -89,13 +90,13 @@ export class RecordOperation {
 /**
  * In-memory cache key - must not be persisted.
  */
-export function cacheKey(sourceTableId: string, id: storage.ReplicaId) {
+export function cacheKey(sourceTableId: storage.SourceTableId, 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')}`;
+export function encodedCacheKey(sourceTableId: storage.SourceTableId, storedKey: Buffer) {
+  return `${postgresTableId(sourceTableId)}.${storedKey.toString('base64')}`;
 }

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

@@ -19,7 +19,7 @@ import { NOTIFICATION_CHANNEL, sql } from '../../utils/db.js';
 import { pick } from '../../utils/ts-codec.js';
 import { batchCreateCustomWriteCheckpoints } from '../checkpoints/PostgresWriteCheckpointAPI.js';
 import { cacheKey, encodedCacheKey, OperationBatch, RecordOperation } from './OperationBatch.js';
-import { PostgresPersistedBatch } from './PostgresPersistedBatch.js';
+import { PostgresPersistedBatch, postgresTableId } from './PostgresPersistedBatch.js';
 import { bigint } from '../../types/codecs.js';
 
 export interface PostgresBucketBatchOptions {
@@ -196,6 +196,7 @@ export class PostgresBucketBatch
           WHERE
             group_id = ${{ type: 'int4', value: this.group_id }}
             AND source_table = ${{ type: 'varchar', value: sourceTable.id }}
+            AND pending_delete IS NULL
           LIMIT
             ${{ type: 'int4', value: BATCH_LIMIT }}
           FOR NO KEY UPDATE
@@ -220,7 +221,9 @@ export class PostgresBucketBatch
             persistedBatch.deleteCurrentData({
               // This is serialized since we got it from a DB query
               serialized_source_key: value.source_key,
-              source_table_id: sourceTable.id
+              source_table_id: postgresTableId(sourceTable.id),
+              // No need for soft delete, since this is not streaming replication
+              soft: false
             });
           }
         }
@@ -630,7 +633,7 @@ export class PostgresBucketBatch
       // exceeding memory limits.
       const sizeLookups = batch.batch.map((r) => {
         return {
-          source_table: r.record.sourceTable.id.toString(),
+          source_table: postgresTableId(r.record.sourceTable.id),
           /**
            * Encode to hex in order to pass a jsonb
            */
@@ -1012,7 +1015,7 @@ export class PostgresBucketBatch
         source_key: afterId,
         group_id: this.group_id,
         data: afterData!,
-        source_table: sourceTable.id,
+        source_table: postgresTableId(sourceTable.id),
         buckets: newBuckets,
         lookups: newLookups,
         pending_delete: null
@@ -1023,8 +1026,9 @@ export class PostgresBucketBatch
     if (afterId == null || !storage.replicaIdEquals(beforeId, afterId)) {
       // Either a delete (afterId == null), or replaced the old replication id
       persistedBatch.deleteCurrentData({
-        source_table_id: sourceTable.id,
-        source_key: beforeId!
+        source_table_id: postgresTableId(sourceTable.id),
+        source_key: beforeId!,
+        soft: true
       });
     }