Support customizing subscriptions

Author: simolus3

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

@@ -1,4 +1,4 @@
-import { BucketDescription, RequestParameters, SqlSyncRules } from '@powersync/service-sync-rules';
+import { BucketDescription, isValidPriority, RequestParameters, SqlSyncRules } from '@powersync/service-sync-rules';
 
 import * as storage from '../storage/storage-index.js';
 import * as util from '../util/util-index.js';
@@ -20,6 +20,7 @@ export interface BucketChecksumStateOptions {
   bucketStorage: BucketChecksumStateStorage;
   syncRules: SqlSyncRules;
   syncParams: RequestParameters;
+  syncRequest: util.StreamingSyncRequest;
   logger?: Logger;
   initialBucketPositions?: { name: string; after: util.InternalOpId }[];
 }
@@ -70,6 +71,7 @@ export class BucketChecksumState {
       options.bucketStorage,
       options.syncRules,
       options.syncParams,
+      options.syncRequest,
       this.logger
     );
     this.bucketDataPositions = new Map();
@@ -182,12 +184,12 @@ export class BucketChecksumState {
 
       const updatedBucketDescriptions = diff.updatedBuckets.map((e) => ({
         ...e,
-        priority: bucketDescriptionMap.get(e.bucket)!.priority
+        ...bucketDescriptionMap.get(e.bucket)!
       }));
       bucketsToFetch = [...generateBucketsToFetch].map((b) => {
         return {
-          bucket: b,
-          priority: bucketDescriptionMap.get(b)!.priority
+          ...bucketDescriptionMap.get(b)!,
+          bucket: b
         };
       });
 
@@ -227,7 +229,7 @@ export class BucketChecksumState {
           write_checkpoint: writeCheckpoint ? String(writeCheckpoint) : undefined,
           buckets: [...checksumMap.values()].map((e) => ({
             ...e,
-            priority: bucketDescriptionMap.get(e.bucket)!.priority
+            ...bucketDescriptionMap.get(e.bucket)!
           }))
         }
       } satisfies util.StreamingSyncCheckpoint;
@@ -336,6 +338,7 @@ export class BucketParameterState {
   public readonly syncParams: RequestParameters;
   private readonly querier: BucketParameterQuerier;
   private readonly staticBuckets: Map<string, BucketDescription>;
+  private readonly explicitStreamSubscriptions: Record<string, util.StreamSubscription>;
   private readonly logger: Logger;
   private cachedDynamicBuckets: BucketDescription[] | null = null;
   private cachedDynamicBucketSet: Set<string> | null = null;
@@ -347,6 +350,7 @@ export class BucketParameterState {
     bucketStorage: BucketChecksumStateStorage,
     syncRules: SqlSyncRules,
     syncParams: RequestParameters,
+    request: util.StreamingSyncRequest,
     logger: Logger
   ) {
     this.context = context;
@@ -355,11 +359,48 @@ export class BucketParameterState {
     this.syncParams = syncParams;
     this.logger = logger;
 
-    this.querier = syncRules.getBucketParameterQuerier(this.syncParams);
+    const explicitStreamSubscriptions: Record<string, util.StreamSubscription> = {};
+    const subscriptions = request.subscriptions;
+    if (subscriptions) {
+      for (const subscription of subscriptions.subscriptions) {
+        explicitStreamSubscriptions[subscription.stream] = subscription;
+      }
+    }
+    this.explicitStreamSubscriptions = explicitStreamSubscriptions;
+
+    this.querier = syncRules.getBucketParameterQuerier({
+      globalParameters: this.syncParams,
+      hasDefaultSubscriptions: subscriptions?.include_defaults ?? true,
+      resolveSubscription(name) {
+        const subscription = explicitStreamSubscriptions[name];
+        if (subscription) {
+          return subscription.parameters ?? {};
+        } else {
+          return null;
+        }
+      }
+    });
     this.staticBuckets = new Map<string, BucketDescription>(this.querier.staticBuckets.map((b) => [b.bucket, b]));
     this.lookups = new Set<string>(this.querier.parameterQueryLookups.map((l) => JSONBig.stringify(l.values)));
   }
 
+  /**
+   * Overrides the `description` based on subscriptions from the client.
+   *
+   * In partiuclar, this can override the priority assigned to a bucket.
+   */
+  overrideBucketDescription(description: BucketDescription): BucketDescription {
+    const changedPriority = this.explicitStreamSubscriptions[description.definition]?.override_priority;
+    if (changedPriority != null && isValidPriority(changedPriority)) {
+      return {
+        ...description,
+        priority: changedPriority
+      };
+    } else {
+      return description;
+    }
+  }
+
   async getCheckpointUpdate(checkpoint: storage.StorageCheckpointUpdate): Promise<CheckpointUpdate> {
     const querier = this.querier;
     let update: CheckpointUpdate;

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

@@ -100,6 +100,7 @@ async function* streamResponseInner(
     bucketStorage,
     syncRules,
     syncParams,
+    syncRequest: params,
     initialBucketPositions: params.buckets?.map((bucket) => ({
       name: bucket.name,
       after: BigInt(bucket.after)

packages/service-core/src/util/protocol-types.ts

@@ -13,9 +13,51 @@ export const BucketRequest = t.object({
 
 export type BucketRequest = t.Decoded<typeof BucketRequest>;
 
+/**
+ * An explicit subscription to a defined sync stream made by the client.
+ */
+export const StreamSubscription = t.object({
+  /**
+   * The defined name of the stream as it appears in sync rules.
+   */
+  stream: t.string,
+  /**
+   * An optional dictionary of parameters to pass to this specific stream.
+   */
+  parameters: t.record(t.any).optional(),
+  /**
+   * Set when the client wishes to re-assign a different priority to this subscription.
+   *
+   * Streams and sync rules can also assign a default priority, but clients are allowed to override those. This can be
+   * useful when the priority for partial syncs depends on e.g. the current page opened in a client.
+   */
+  override_priority: t.number.optional()
+});
+
+export type StreamSubscription = t.Decoded<typeof StreamSubscription>;
+
+/**
+ * An overview of all subscriptions as part of a streaming sync request.
+ */
+export const StreamSubscriptions = t.object({
+  /**
+   * Whether to sync default streams.
+   *
+   * When disabled,only
+   */
+  include_defaults: t.boolean.optional(),
+
+  /**
+   * An array of streams the client has subscribed to.
+   */
+  subscriptions: t.array(StreamSubscription)
+});
+
+export type StreamSubscriptions = t.Decoded<typeof StreamSubscriptions>;
+
 export const StreamingSyncRequest = t.object({
   /**
-   * Existing bucket states.
+   * Existing client-side bucket states.
    */
   buckets: t.array(BucketRequest).optional(),
 
@@ -47,7 +89,12 @@ export const StreamingSyncRequest = t.object({
   /**
    * Unique client id.
    */
-  client_id: t.string.optional()
+  client_id: t.string.optional(),
+
+  /**
+   * If the client is aware of stream subscriptions, an array of streams the client is subscribing to.
+   */
+  subscriptions: StreamSubscriptions.optional()
 });
 
 export type StreamingSyncRequest = t.Decoded<typeof StreamingSyncRequest>;

packages/sync-rules/src/BucketDescription.ts

@@ -20,6 +20,10 @@ export const isValidPriority = (i: number): i is BucketPriority => {
 };
 
 export interface BucketDescription {
+  /**
+   * The name of the sync rule or stream definition from which the bucket is derived.
+   */
+  definition: string;
   /**
    * The id of the bucket, which is derived from the name of the bucket's definition
    * in the sync rules as well as the values returned by the parameter queries.

packages/sync-rules/src/SqlBucketDescriptor.ts

@@ -29,12 +29,23 @@ export interface QueryParseResult {
   errors: SqlRuleError[];
 }
 
+export enum SqlBucketDescriptorType {
+  SYNC_RULE,
+  STREAM
+}
+
 export class SqlBucketDescriptor {
   name: string;
   bucketParameters?: string[];
+  type: SqlBucketDescriptorType;
+  subscribedToByDefault: boolean;
 
-  constructor(name: string) {
+  constructor(name: string, type: SqlBucketDescriptorType) {
     this.name = name;
+    this.type = type;
+
+    // Sync-rule style buckets are subscribed to by default, streams are opt-in unless their definition says otherwise.
+    this.subscribedToByDefault = type == SqlBucketDescriptorType.SYNC_RULE;
   }
 
   /**
@@ -93,6 +104,7 @@ export class SqlBucketDescriptor {
       }
     }
     this.dataQueries.push(query.data);
+    this.subscribedToByDefault = options.default ?? false;
 
     return {
       parsed: true,

packages/sync-rules/src/SqlParameterQuery.ts

@@ -367,6 +367,7 @@ export class SqlParameterQuery {
         }
 
         return {
+          definition: this.descriptorName,
           bucket: getBucketId(this.descriptorName, this.bucketParameters, result),
           priority: this.priority
         };

packages/sync-rules/src/SqlSyncRules.ts

@@ -6,7 +6,7 @@ import { SqlEventDescriptor } from './events/SqlEventDescriptor.js';
 import { IdSequence } from './IdSequence.js';
 import { validateSyncRulesSchema } from './json_schema.js';
 import { SourceTableInterface } from './SourceTableInterface.js';
-import { QueryParseResult, SqlBucketDescriptor } from './SqlBucketDescriptor.js';
+import { QueryParseResult, SqlBucketDescriptor, SqlBucketDescriptorType } from './SqlBucketDescriptor.js';
 import { TablePattern } from './TablePattern.js';
 import {
   EvaluatedParameters,
@@ -40,6 +40,22 @@ export interface SyncRulesOptions {
   throwOnError?: boolean;
 }
 
+export interface GetQuerierOptions {
+  globalParameters: RequestParameters;
+  /**
+   * Whether the client is subscribing to default subscriptions (the default).
+   */
+  hasDefaultSubscriptions: boolean;
+  /**
+   * For streams, this is invoked to check whether the client has requested a subscription to
+   * the stream.
+   *
+   * @param name The name of the stream as it appears in the sync rule definitions.
+   * @returns If a subscription is active, the stream parameters for that particular stream. Otherwise null.
+   */
+  resolveSubscription: (name: string) => Record<string, any> | null;
+}
+
 export class SqlSyncRules implements SyncRules {
   bucketDescriptors: SqlBucketDescriptor[] = [];
   eventDescriptors: SqlEventDescriptor[] = [];
@@ -144,7 +160,7 @@ export class SqlSyncRules implements SyncRules {
       const parameters = value.get('parameters', true) as unknown;
       const dataQueries = value.get('data', true) as unknown;
 
-      const descriptor = new SqlBucketDescriptor(key);
+      const descriptor = new SqlBucketDescriptor(key, SqlBucketDescriptorType.SYNC_RULE);
 
       if (parameters instanceof Scalar) {
         rules.withScalar(parameters, (q) => {
@@ -179,7 +195,7 @@ export class SqlSyncRules implements SyncRules {
         continue;
       }
 
-      const descriptor = new SqlBucketDescriptor(key);
+      const descriptor = new SqlBucketDescriptor(key, SqlBucketDescriptorType.STREAM);
 
       const accept_potentially_dangerous_queries =
         value.get('accept_potentially_dangerous_queries', true)?.value == true;
@@ -355,8 +371,25 @@ export class SqlSyncRules implements SyncRules {
     return { results, errors };
   }
 
-  getBucketParameterQuerier(parameters: RequestParameters): BucketParameterQuerier {
-    const queriers = this.bucketDescriptors.map((query) => query.getBucketParameterQuerier(parameters));
+  getBucketParameterQuerier(options: GetQuerierOptions): BucketParameterQuerier {
+    const queriers: BucketParameterQuerier[] = [];
+    for (const descriptor of this.bucketDescriptors) {
+      let params = options.globalParameters;
+      const subscription =
+        descriptor.type == SqlBucketDescriptorType.STREAM ? options.resolveSubscription(descriptor.name) : null;
+
+      if (!descriptor.subscribedToByDefault && subscription == null) {
+        // The client is not subscribing to this stream, so don't query buckets related to it.
+        continue;
+      }
+
+      if (subscription != null) {
+        params = params.withAddedParameters(subscription);
+      }
+
+      queriers.push(descriptor.getBucketParameterQuerier(params));
+    }
+
     return mergeBucketParameterQueriers(queriers);
   }
 

packages/sync-rules/src/StaticSqlParameterQuery.ts

@@ -177,6 +177,7 @@ export class StaticSqlParameterQuery {
 
     return [
       {
+        definition: this.descriptorName,
         bucket: getBucketId(this.descriptorName, this.bucketParameters, result),
         priority: this.priority
       }

packages/sync-rules/src/TableValuedFunctionSqlParameterQuery.ts

@@ -249,6 +249,7 @@ export class TableValuedFunctionSqlParameterQuery {
     }
 
     return {
+      definition: this.descriptorName,
       bucket: getBucketId(this.descriptorName, this.bucketParameters, result),
       priority: this.priority
     };

packages/sync-rules/src/types.ts

@@ -135,6 +135,14 @@ export class RequestParameters implements ParameterValueSet {
     }
     throw new Error(`Unknown table: ${table}`);
   }
+
+  withAddedParameters(params: Record<string, any>): RequestParameters {
+    const clone = structuredClone(this);
+    clone.rawUserParameters = JSONBig.stringify(params);
+    clone.userParameters = toSyncRulesParameters(params);
+
+    return clone;
+  }
 }
 
 /**