Deduplicate buckets

Author: simolus3

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

@@ -112,7 +112,7 @@ export class BucketChecksumState {
 
     /** Set of all buckets in this checkpoint. */
     const bucketDescriptionMap = new Map(
-      allBuckets.map((b) => [b.bucket, this.parameterState.overrideBucketDescription(b)])
+      allBuckets.map((b) => [b.bucket, this.parameterState.translateResolvedBucket(b)])
     );
 
     if (bucketDescriptionMap.size > this.context.maxBuckets) {
@@ -345,7 +345,7 @@ export interface CheckpointUpdate {
   /**
    * All buckets forming part of the checkpoint.
    */
-  buckets: BucketDescription[];
+  buckets: ResolvedBucket[];
 
   /**
    * If present, a set of buckets that have been updated since the last checkpoint.
@@ -367,7 +367,7 @@ export class BucketParameterState {
   private readonly explicitStreamSubscriptions: Record<string, util.RequestedStreamSubscription>;
   private readonly subscribedStreamNames: Set<string>;
   private readonly logger: Logger;
-  private cachedDynamicBuckets: BucketDescription[] | null = null;
+  private cachedDynamicBuckets: ResolvedBucket[] | null = null;
   private cachedDynamicBucketSet: Set<string> | null = null;
 
   private readonly lookups: Set<string>;
@@ -412,7 +412,10 @@ export class BucketParameterState {
       hasDefaultStreams: this.includeDefaultStreams,
       streams: streamsByName
     });
-    this.staticBuckets = new Map<string, BucketDescription>(this.querier.staticBuckets.map((b) => [b.bucket, b]));
+
+    this.staticBuckets = new Map<string, BucketDescription>(
+      mergeBuckets(this.querier.staticBuckets).map((b) => [b.bucket, b])
+    );
     this.lookups = new Set<string>(this.querier.parameterQueryLookups.map((l) => JSONBig.stringify(l.values)));
     this.subscribedStreamNames = new Set(Object.keys(streamsByName));
   }
@@ -422,7 +425,8 @@ export class BucketParameterState {
    * {@link util.ClientBucketDescription}.
    */
   translateResolvedBucket(description: ResolvedBucket): util.ClientBucketDescription {
-    // Assign
+    // If the client is overriding the priority of any stream that yields this bucket, sync the bucket with that
+    // priority.
     let priorityOverride: BucketPriority | null = null;
     for (const reason of description.inclusion_reasons) {
       if (reason != 'default') {
@@ -531,7 +535,7 @@ export class BucketParameterState {
       }
     }
 
-    let dynamicBuckets: BucketDescription[];
+    let dynamicBuckets: ResolvedBucket[];
     if (hasParameterChange || this.cachedDynamicBuckets == null || this.cachedDynamicBucketSet == null) {
       dynamicBuckets = await querier.queryDynamicBucketDescriptions({
         getParameterSets(lookups) {
@@ -612,3 +616,32 @@ function limitedBuckets(buckets: string[] | { bucket: string }[], limit: number)
   const limited = buckets.slice(0, limit);
   return `${JSON.stringify(limited)}...`;
 }
+
+/**
+ * Resolves duplicate buckets in the given array, merging the inclusion reasons for duplicate.
+ *
+ * It's possible for duplicates to occur when a stream has multiple subscriptions, consider e.g.
+ *
+ * ```
+ * sync_streams:
+ *  assets_by_category:
+ *    query: select * from assets where category in (request.parameters() -> 'categories')
+ * ```
+ *
+ * Here, a client might subscribe once with `{"categories": [1]}` and once with `{"categories": [1, 2]}`. Since each
+ * subscription is evaluated independently, this would lead to three buckets, with a duplicate `assets_by_category[1]`
+ * bucket.
+ */
+function mergeBuckets(buckets: ResolvedBucket[]): ResolvedBucket[] {
+  const byDefinition: Record<string, ResolvedBucket> = {};
+
+  for (const bucket of buckets) {
+    if (Object.hasOwn(byDefinition, bucket.definition)) {
+      byDefinition[bucket.definition].inclusion_reasons.push(...bucket.inclusion_reasons);
+    } else {
+      byDefinition[bucket.definition] = bucket;
+    }
+  }
+
+  return Object.values(byDefinition);
+}

packages/sync-rules/src/BucketParameterQuerier.ts

@@ -1,4 +1,4 @@
-import { BucketDescription } from './BucketDescription.js';
+import { BucketDescription, ResolvedBucket } from './BucketDescription.js';
 import { RequestParameters, SqliteJsonRow, SqliteJsonValue } from './types.js';
 import { normalizeParameterValue } from './utils.js';
 
@@ -14,7 +14,7 @@ export interface BucketParameterQuerier {
    *     select request.user_id() as user_id()
    *     select value as project_id from json_each(request.jwt() -> 'project_ids')
    */
-  readonly staticBuckets: BucketDescription[];
+  readonly staticBuckets: ResolvedBucket[];
 
   /**
    * True if there are dynamic buckets, meaning queryDynamicBucketDescriptions() should be used.
@@ -36,7 +36,7 @@ export interface BucketParameterQuerier {
    *
    *     select id as user_id from users where users.id = request.user_id()
    */
-  queryDynamicBucketDescriptions(source: ParameterLookupSource): Promise<BucketDescription[]>;
+  queryDynamicBucketDescriptions(source: ParameterLookupSource): Promise<ResolvedBucket[]>;
 }
 
 export interface ParameterLookupSource {
@@ -54,7 +54,7 @@ export function mergeBucketParameterQueriers(queriers: BucketParameterQuerier[])
     hasDynamicBuckets: parameterQueryLookups.length > 0,
     parameterQueryLookups: parameterQueryLookups,
     async queryDynamicBucketDescriptions(source: ParameterLookupSource) {
-      let results: BucketDescription[] = [];
+      let results: ResolvedBucket[] = [];
       for (let q of queriers) {
         if (q.hasDynamicBuckets) {
           results.push(...(await q.queryDynamicBucketDescriptions(source)));

packages/sync-rules/src/SqlBucketDescriptor.ts

@@ -1,10 +1,10 @@
-import { BucketDescription } from './BucketDescription.js';
+import { BucketDescription, BucketInclusionReason, ResolvedBucket } from './BucketDescription.js';
 import { BucketParameterQuerier, mergeBucketParameterQueriers } from './BucketParameterQuerier.js';
 import { IdSequence } from './IdSequence.js';
 import { SourceTableInterface } from './SourceTableInterface.js';
 import { SqlDataQuery } from './SqlDataQuery.js';
 import { SqlParameterQuery } from './SqlParameterQuery.js';
-import { SyncRulesOptions } from './SqlSyncRules.js';
+import { GetQuerierOptions, SyncRulesOptions } from './SqlSyncRules.js';
 import { StaticSqlParameterQuery } from './StaticSqlParameterQuery.js';
 import { StreamQuery } from './StreamQuery.js';
 import { TablePattern } from './TablePattern.js';
@@ -134,27 +134,48 @@ export class SqlBucketDescriptor {
     return results;
   }
 
-  getBucketParameterQuerier(parameters: RequestParameters): BucketParameterQuerier {
-    const staticBuckets = this.getStaticBucketDescriptions(parameters);
+  /**
+   * @deprecated Use `pushBucketParameterQueriers` instead and merge at the top-level.
+   */
+  getBucketParameterQuerier(options: GetQuerierOptions, parameters: RequestParameters): BucketParameterQuerier {
+    const queriers: BucketParameterQuerier[] = [];
+    this.pushBucketParameterQueriers(queriers, options, parameters);
+
+    return mergeBucketParameterQueriers(queriers);
+  }
+
+  pushBucketParameterQueriers(
+    result: BucketParameterQuerier[],
+    options: GetQuerierOptions,
+    parameters: RequestParameters
+  ) {
+    const reasons = [this.bucketInclusionReason(options)];
+    const staticBuckets = this.getStaticBucketDescriptions(parameters, reasons);
     const staticQuerier = {
       staticBuckets,
       hasDynamicBuckets: false,
       parameterQueryLookups: [],
       queryDynamicBucketDescriptions: async () => []
     } satisfies BucketParameterQuerier;
+    result.push(staticQuerier);
 
     if (this.parameterQueries.length == 0) {
-      return staticQuerier;
+      return;
     }
 
-    const dynamicQueriers = this.parameterQueries.map((query) => query.getBucketParameterQuerier(parameters));
-    return mergeBucketParameterQueriers([staticQuerier, ...dynamicQueriers]);
+    const dynamicQueriers = this.parameterQueries.map((query) => query.getBucketParameterQuerier(parameters, reasons));
+    result.push(...dynamicQueriers);
   }
 
-  getStaticBucketDescriptions(parameters: RequestParameters): BucketDescription[] {
-    let results: BucketDescription[] = [];
+  getStaticBucketDescriptions(parameters: RequestParameters, reasons: BucketInclusionReason[]): ResolvedBucket[] {
+    let results: ResolvedBucket[] = [];
     for (let query of this.globalParameterQueries) {
-      results.push(...query.getStaticBucketDescriptions(parameters));
+      for (const desc of query.getStaticBucketDescriptions(parameters)) {
+        results.push({
+          ...desc,
+          inclusion_reasons: reasons
+        });
+      }
     }
     return results;
   }
@@ -177,6 +198,14 @@ export class SqlBucketDescriptor {
     return result;
   }
 
+  private bucketInclusionReason(parameters: GetQuerierOptions): BucketInclusionReason {
+    if (this.type == SqlBucketDescriptorType.STREAM && !this.subscribedToByDefault) {
+      return { subscription: this.name };
+    } else {
+      return 'default';
+    }
+  }
+
   tableSyncsData(table: SourceTableInterface): boolean {
     for (let query of this.dataQueries) {
       if (query.applies(table)) {

packages/sync-rules/src/SqlParameterQuery.ts

@@ -1,5 +1,10 @@
 import { parse, SelectedColumn } from 'pgsql-ast-parser';
-import { BucketDescription, BucketPriority, DEFAULT_BUCKET_PRIORITY } from './BucketDescription.js';
+import {
+  BucketDescription,
+  BucketInclusionReason,
+  BucketPriority,
+  DEFAULT_BUCKET_PRIORITY
+} from './BucketDescription.js';
 import { BucketParameterQuerier, ParameterLookup, ParameterLookupSource } from './BucketParameterQuerier.js';
 import { SqlRuleError } from './errors.js';
 import { SourceTableInterface } from './SourceTableInterface.js';
@@ -451,7 +456,10 @@ export class SqlParameterQuery {
     }
   }
 
-  getBucketParameterQuerier(requestParameters: RequestParameters): BucketParameterQuerier {
+  getBucketParameterQuerier(
+    requestParameters: RequestParameters,
+    reasons: BucketInclusionReason[]
+  ): BucketParameterQuerier {
     const lookups = this.getLookups(requestParameters);
     if (lookups.length == 0) {
       // This typically happens when the query is pre-filtered using a where clause
@@ -470,7 +478,10 @@ export class SqlParameterQuery {
       parameterQueryLookups: lookups,
       queryDynamicBucketDescriptions: async (source: ParameterLookupSource) => {
         const bucketParameters = await source.getParameterSets(lookups);
-        return this.resolveBucketDescriptions(bucketParameters, requestParameters);
+        return this.resolveBucketDescriptions(bucketParameters, requestParameters).map((bucket) => ({
+          ...bucket,
+          inclusion_reasons: reasons
+        }));
       }
     };
   }

packages/sync-rules/src/SqlSyncRules.ts

@@ -411,16 +411,16 @@ export class SqlSyncRules implements SyncRules {
             hasExplicitDefaultSubscription = true;
           }
 
-          queriers.push(descriptor.getBucketParameterQuerier(subscriptionParams));
+          descriptor.pushBucketParameterQueriers(queriers, options, subscriptionParams);
         }
 
         // If the stream is subscribed to by default and there is no explicit subscription that would match the default
         // subscription, also include the default querier.
         if (descriptor.subscribedToByDefault && !hasExplicitDefaultSubscription) {
-          queriers.push(descriptor.getBucketParameterQuerier(params));
+          descriptor.pushBucketParameterQueriers(queriers, options, params);
         }
       } else {
-        queriers.push(descriptor.getBucketParameterQuerier(params));
+        descriptor.pushBucketParameterQueriers(queriers, options, params);
       }
     }