Track bucket parameter changes.

Author: rkistner

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

@@ -9,22 +9,30 @@ import {
 } from '@powersync/lib-services-framework';
 import {
   BroadcastIterable,
+  getLookupBucketDefinitionName,
   ReplicationCheckpoint,
   storage,
   utils,
   WatchWriteCheckpointOptions
 } from '@powersync/service-core';
 import { SqliteJsonRow, SqliteJsonValue, SqlSyncRules } from '@powersync/service-sync-rules';
 import * as bson from 'bson';
+import { wrapWithAbort } from 'ix/asynciterable/operators/withabort.js';
 import * as timers from 'timers/promises';
 import { MongoBucketStorage } from '../MongoBucketStorage.js';
 import { PowerSyncMongo } from './db.js';
-import { BucketDataDocument, BucketDataKey, SourceKey, SyncRuleCheckpointState, SyncRuleDocument } from './models.js';
+import {
+  BucketDataDocument,
+  BucketDataKey,
+  BucketParameterDocument,
+  SourceKey,
+  SyncRuleCheckpointState,
+  SyncRuleDocument
+} from './models.js';
 import { MongoBucketBatch } from './MongoBucketBatch.js';
 import { MongoCompactor } from './MongoCompactor.js';
 import { MongoWriteCheckpointAPI } from './MongoWriteCheckpointAPI.js';
 import { idPrefixFilter, mapOpEntry, readSingleBatch } from './util.js';
-import { wrapWithAbort } from 'ix/asynciterable/operators/withabort.js';
 
 export class MongoSyncBucketStorage
   extends DisposableObserver<storage.SyncRulesBucketStorageListener>
@@ -734,7 +742,10 @@ export class MongoSyncBucketStorage
         }
       } else if (update.ns.coll == this.db.bucket_data.collectionName) {
         const bucket = (update.documentKey._id as unknown as BucketDataKey).b;
-        yield { bucket: bucket };
+        yield { dataBucket: bucket };
+      } else if (update.ns.coll == this.db.bucket_parameters.collectionName && update.fullDocument != null) {
+        const bucketDefinition = getLookupBucketDefinitionName((update.fullDocument as BucketParameterDocument).lookup);
+        yield { parameterBucketDefinition: bucketDefinition };
       }
     }
   }
@@ -757,15 +768,24 @@ export class MongoSyncBucketStorage
     for await (const event of iter) {
       if (filter) {
         if (event.invalidate) {
-          shouldUpdate ||= filter({
-            invalidate: true
-          });
+          shouldUpdate =
+            filter({
+              invalidate: true
+            }) || shouldUpdate;
         }
 
-        if (event.bucket) {
-          shouldUpdate ||= filter({
-            bucket: event.bucket
-          });
+        if (event.dataBucket) {
+          shouldUpdate =
+            filter({
+              changedDataBucket: event.dataBucket
+            }) || shouldUpdate;
+        }
+
+        if (event.parameterBucketDefinition) {
+          shouldUpdate =
+            filter({
+              changedParameterBucketDefinition: event.parameterBucketDefinition
+            }) || shouldUpdate;
         }
       }
 
@@ -836,23 +856,57 @@ export class MongoSyncBucketStorage
 
   private getChangeStreamPipeline() {
     const syncRulesId = this.group_id;
+    const parsedSyncRules = this.getParsedSyncRules({
+      // Not applicable here
+      defaultSchema: 'default'
+    });
+    const staticBucketDefinitions = parsedSyncRules.getStaticBucketDefinitionList();
+    const staticBucketFilters = staticBucketDefinitions.map((name) => {
+      return {
+        // Check that the docuemnt starts with bucket name
+        'documentKey._id.b': {
+          $gte: name + '[',
+          $lt: name + '[\uFFFF'
+        }
+      };
+    });
+
+    let filters: mongo.Document[] = [
+      // sync_rules events
+      {
+        'ns.coll': this.db.sync_rules.collectionName,
+        'documentKey._id': syncRulesId,
+        operationType: { $in: ['insert', 'update', 'replace'] }
+      },
+      // bucket_data events
+      {
+        'ns.coll': this.db.bucket_data.collectionName,
+        'documentKey._id.g': syncRulesId,
+        operationType: 'insert',
+        $or: staticBucketFilters
+      },
+      // bucket_parameters events
+      {
+        'ns.coll': this.db.bucket_parameters.collectionName,
+        'fullDocument.key.g': syncRulesId,
+        operationType: 'insert'
+      }
+    ];
+
+    if (staticBucketFilters.length > 0) {
+      // bucket_data events, only if there are static bucket definitions
+      filters.push({
+        'ns.coll': this.db.bucket_data.collectionName,
+        'documentKey._id.g': syncRulesId,
+        operationType: 'insert',
+        $or: staticBucketFilters
+      });
+    }
+
     const pipeline: mongo.Document[] = [
       {
         $match: {
-          $or: [
-            // sync_rules events
-            {
-              'ns.coll': this.db.sync_rules.collectionName,
-              'documentKey._id': syncRulesId,
-              operationType: { $in: ['insert', 'update', 'replace'] }
-            },
-            // bucket_data events
-            {
-              'ns.coll': this.db.bucket_data.collectionName,
-              'documentKey._id.g': syncRulesId,
-              operationType: 'insert'
-            }
-          ]
+          $or: filters
         }
       },
       {
@@ -865,6 +919,10 @@ export class MongoSyncBucketStorage
           // For sync_rules, this contains the sync_rules id
           'documentKey._id': 1,
 
+          // For bucket_parameters
+          'fullDocument.key': 1,
+          'fullDocument.lookup': 1,
+
           // For sync_rules events
           'updateDescription.updatedFields.state': 1,
           'updateDescription.updatedFields.last_checkpoint': 1,
@@ -882,6 +940,7 @@ export class MongoSyncBucketStorage
 
 interface BucketCheckpointEvent {
   checkpoint?: ReplicationCheckpoint;
-  bucket?: string;
+  dataBucket?: string;
+  parameterBucketDefinition?: string;
   invalidate?: boolean;
 }

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

@@ -222,7 +222,8 @@ export interface WatchWriteCheckpointOptions {
 }
 
 export interface WatchFilterEvent {
-  bucket?: string;
+  changedDataBucket?: string;
+  changedParameterBucketDefinition?: string;
   invalidate?: boolean;
 }
 

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

@@ -29,6 +29,11 @@ export const serializeLookup = (lookup: SqliteJsonValue[]) => {
   return new bson.Binary(serializeLookupBuffer(lookup));
 };
 
+export const getLookupBucketDefinitionName = (lookup: bson.Binary) => {
+  const parsed = bson.deserialize(lookup.buffer, BSON_DESERIALIZE_OPTIONS).l as SqliteJsonValue[];
+  return parsed[0] as string;
+};
+
 /**
  * True if this is a bson.UUID.
  *

packages/service-core/src/sync/BucketChecksumState.ts

@@ -280,27 +280,44 @@ export class BucketParameterState {
       this.invalidated = true;
       this.pendingBuckets.clear();
       return true;
-    } else if (event.bucket != null) {
+    } else if (event.changedDataBucket != null) {
       const querier = this.querier;
       const staticBuckets = this.staticBuckets;
       if (querier.hasDynamicBuckets) {
-        // TODO: Check if the dynamic buckets may match this one
-        this.invalidated = true;
-        this.pendingBuckets.clear();
-        return true;
+        // Check for a maybe-match on dynamic buckets
+        // We could expand this to check our last set of buckets
+        for (let def of this.querier.dynamicBucketDefinitions) {
+          if (event.changedDataBucket.startsWith(def)) {
+            return this.invalidatePendingBuckets();
+          }
+        }
       }
 
-      if (staticBuckets.has(event.bucket)) {
-        this.pendingBuckets.add(event.bucket);
+      // For static buckets we can check for an exact match
+      if (staticBuckets.has(event.changedDataBucket)) {
+        this.pendingBuckets.add(event.changedDataBucket);
         return true;
       } else {
         return false;
       }
+    } else if (event.changedParameterBucketDefinition) {
+      if (this.querier.dynamicBucketDefinitions.has(event.changedParameterBucketDefinition)) {
+        // Potential change in dynamic bucket list - invalidate
+        return this.invalidatePendingBuckets();
+      }
+      // The parameter bucket definition is not relevant for this connection
+      return false;
     } else {
       return false;
     }
   };
 
+  private invalidatePendingBuckets() {
+    this.invalidated = true;
+    this.pendingBuckets.clear();
+    return true;
+  }
+
   async getCheckpointUpdate(checkpoint: util.OpId): Promise<CheckpointUpdate> {
     const querier = this.querier;
     let update: CheckpointUpdate;

packages/service-core/test/src/sync/BucketChecksumState.test.ts

@@ -252,7 +252,7 @@ bucket_definitions:
     // Invalidate the state for global[1] - will only re-check the single bucket.
     // This is essentially inconsistent state, but is the simplest way to test that
     // the filter is working.
-    state.checkpointFilter({ bucket: 'global[1]' });
+    state.checkpointFilter({ changedDataBucket: 'global[1]' });
 
     const line2 = (await state.buildNextCheckpointLine({
       base: { checkpoint: '2', lsn: '2' },
@@ -460,8 +460,7 @@ bucket_definitions:
     state.updateBucketPosition({ bucket: 'by_project[1]', nextAfter: '1', hasMore: false });
     state.updateBucketPosition({ bucket: 'by_project[2]', nextAfter: '1', hasMore: false });
 
-    // Update bucket storage state
-    state.checkpointFilter({ invalidate: true });
+    state.checkpointFilter({ changedParameterBucketDefinition: 'by_project' });
 
     storage.getParameterSets = async (checkpoint: OpId, lookups: SqliteJsonValue[][]): Promise<SqliteJsonRow[]> => {
       expect(checkpoint).toEqual('2');
@@ -494,7 +493,7 @@ class MockBucketChecksumStateStorage implements BucketChecksumStateStorage {
 
   updateTestChecksum(checksum: BucketChecksum): void {
     this.state.set(checksum.bucket, checksum);
-    this.filter?.({ bucket: checksum.bucket });
+    this.filter?.({ changedDataBucket: checksum.bucket });
   }
 
   invalidate() {

packages/sync-rules/src/BucketParameterQuerier.ts

@@ -18,10 +18,13 @@ export interface BucketParameterQuerier {
   /**
    * True if there are dynamic buckets, meaning queryDynamicBucketDescriptions() should be used.
    *
-   * If this is false, queryDynamicBucketDescriptions() will always return an empty array.
+   * If this is false, queryDynamicBucketDescriptions() will always return an empty array,
+   * and dynamicBucketDefinitions.size == 0.
    */
   readonly hasDynamicBuckets: boolean;
 
+  readonly dynamicBucketDefinitions: Set<string>;
+
   /**
    * These buckets depend on parameter storage, and needs to be retrieved dynamically for each checkpoint.
    *
@@ -44,9 +47,11 @@ export interface QueryBucketDescriptorOptions extends ParameterLookupSource {
 }
 
 export function mergeBucketParameterQueriers(queriers: BucketParameterQuerier[]): BucketParameterQuerier {
+  const dynamicBucketDefinitions = new Set<string>(...queriers.flatMap((q) => [...q.dynamicBucketDefinitions]));
   return {
     staticBuckets: queriers.flatMap((q) => q.staticBuckets),
-    hasDynamicBuckets: queriers.some((q) => q.hasDynamicBuckets),
+    hasDynamicBuckets: dynamicBucketDefinitions.size > 0,
+    dynamicBucketDefinitions,
     async queryDynamicBucketDescriptions(source: ParameterLookupSource) {
       let results: BucketDescription[] = [];
       for (let q of queriers) {

packages/sync-rules/src/SqlBucketDescriptor.ts

@@ -1,9 +1,5 @@
 import { BucketDescription } from './BucketDescription.js';
-import {
-  BucketParameterQuerier,
-  mergeBucketParameterQueriers,
-  QueryBucketDescriptorOptions
-} from './BucketParameterQuerier.js';
+import { BucketParameterQuerier, mergeBucketParameterQueriers } from './BucketParameterQuerier.js';
 import { IdSequence } from './IdSequence.js';
 import { SourceTableInterface } from './SourceTableInterface.js';
 import { SqlDataQuery } from './SqlDataQuery.js';
@@ -13,8 +9,8 @@ import { StaticSqlParameterQuery } from './StaticSqlParameterQuery.js';
 import { TablePattern } from './TablePattern.js';
 import { SqlRuleError } from './errors.js';
 import {
-  EvaluateRowOptions,
   EvaluatedParametersResult,
+  EvaluateRowOptions,
   EvaluationResult,
   QueryParseOptions,
   RequestParameters,
@@ -117,6 +113,7 @@ export class SqlBucketDescriptor {
     const staticQuerier = {
       staticBuckets,
       hasDynamicBuckets: false,
+      dynamicBucketDefinitions: new Set<string>(),
       queryDynamicBucketDescriptions: async () => []
     } satisfies BucketParameterQuerier;
 
@@ -128,6 +125,13 @@ export class SqlBucketDescriptor {
     return mergeBucketParameterQueriers([staticQuerier, ...dynamicQueriers]);
   }
 
+  /**
+   * Return bucket definition names for all static buckets.
+   */
+  hasStaticBuckets(): boolean {
+    return this.global_parameter_queries.length > 0;
+  }
+
   getStaticBucketDescriptions(parameters: RequestParameters): BucketDescription[] {
     let results: BucketDescription[] = [];
     for (let query of this.global_parameter_queries) {

packages/sync-rules/src/SqlParameterQuery.ts

@@ -1,11 +1,14 @@
 import { parse, SelectedColumn } from 'pgsql-ast-parser';
+import { BucketDescription, BucketPriority, defaultBucketPriority } from './BucketDescription.js';
+import { BucketParameterQuerier, ParameterLookupSource } from './BucketParameterQuerier.js';
 import { SqlRuleError } from './errors.js';
 import { SourceTableInterface } from './SourceTableInterface.js';
 import { SqlTools } from './sql_filters.js';
 import { checkUnsupportedFeatures, isClauseError, isParameterValueClause } from './sql_support.js';
 import { StaticSqlParameterQuery } from './StaticSqlParameterQuery.js';
 import { TablePattern } from './TablePattern.js';
 import { TableQuerySchema } from './TableQuerySchema.js';
+import { TableValuedFunctionSqlParameterQuery } from './TableValuedFunctionSqlParameterQuery.js';
 import {
   EvaluatedParameters,
   EvaluatedParametersResult,
@@ -21,13 +24,6 @@ import {
   SqliteRow
 } from './types.js';
 import { filterJsonRow, getBucketId, isJsonValue, isSelectStatement } from './utils.js';
-import { TableValuedFunctionSqlParameterQuery } from './TableValuedFunctionSqlParameterQuery.js';
-import { BucketDescription, BucketPriority, defaultBucketPriority } from './BucketDescription.js';
-import {
-  BucketParameterQuerier,
-  ParameterLookupSource,
-  QueryBucketDescriptorOptions
-} from './BucketParameterQuerier.js';
 
 /**
  * Represents a parameter query, such as:
@@ -193,6 +189,10 @@ export class SqlParameterQuery {
   priority?: BucketPriority;
 
   filter?: ParameterMatchClause;
+
+  /**
+   * Bucket definition name.
+   */
   descriptor_name?: string;
 
   /** _Input_ token / user parameters */
@@ -375,13 +375,15 @@ export class SqlParameterQuery {
       return {
         staticBuckets: [],
         hasDynamicBuckets: false,
+        dynamicBucketDefinitions: new Set<string>(),
         queryDynamicBucketDescriptions: async () => []
       };
     }
 
     return {
       staticBuckets: [],
       hasDynamicBuckets: true,
+      dynamicBucketDefinitions: new Set<string>(this.descriptor_name!),
       queryDynamicBucketDescriptions: async (source: ParameterLookupSource) => {
         const bucketParameters = await source.getParameterSets(lookups);
         return this.resolveBucketDescriptions(bucketParameters, requestParameters);