Include query streams in checkpoint message

Author: simolus3

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

@@ -14,6 +14,7 @@ import { JSONBig } from '@powersync/service-jsonbig';
 import { BucketParameterQuerier } from '@powersync/service-sync-rules/src/BucketParameterQuerier.js';
 import { SyncContext } from './SyncContext.js';
 import { getIntersection, hasIntersection } from './util.js';
+import { SqlBucketDescriptor, SqlBucketDescriptorType } from '@powersync/service-sync-rules/src/SqlBucketDescriptor.js';
 
 export interface BucketChecksumStateOptions {
   syncContext: SyncContext;
@@ -102,7 +103,9 @@ export class BucketChecksumState {
     const { buckets: allBuckets, updatedBuckets } = update;
 
     /** Set of all buckets in this checkpoint. */
-    const bucketDescriptionMap = new Map(allBuckets.map((b) => [b.bucket, b]));
+    const bucketDescriptionMap = new Map(
+      allBuckets.map((b) => [b.bucket, this.parameterState.overrideBucketDescription(b)])
+    );
 
     if (bucketDescriptionMap.size > this.context.maxBuckets) {
       throw new ServiceError(
@@ -223,14 +226,27 @@ export class BucketChecksumState {
         this.logger.info(message, { checkpoint: base.checkpoint, user_id: user_id, buckets: allBuckets.length });
       };
       bucketsToFetch = allBuckets;
+      this.parameterState.syncRules.bucketDescriptors;
+
+      const subscriptions: util.SubscribedStream[] = [];
+      for (const desc of this.parameterState.syncRules.bucketDescriptors) {
+        if (desc.type == SqlBucketDescriptorType.STREAM && this.parameterState.isSubscribedToStream(desc)) {
+          subscriptions.push({
+            name: desc.name,
+            is_default: desc.subscribedToByDefault
+          });
+        }
+      }
+
       checkpointLine = {
         checkpoint: {
           last_op_id: util.internalToExternalOpId(base.checkpoint),
           write_checkpoint: writeCheckpoint ? String(writeCheckpoint) : undefined,
           buckets: [...checksumMap.values()].map((e) => ({
             ...e,
             ...bucketDescriptionMap.get(e.bucket)!
-          }))
+          })),
+          included_subscriptions: subscriptions
         }
       } satisfies util.StreamingSyncCheckpoint;
     }
@@ -338,6 +354,7 @@ export class BucketParameterState {
   public readonly syncParams: RequestParameters;
   private readonly querier: BucketParameterQuerier;
   private readonly staticBuckets: Map<string, BucketDescription>;
+  private readonly includeDefaultStreams: boolean;
   private readonly explicitStreamSubscriptions: Record<string, util.StreamSubscription>;
   private readonly logger: Logger;
   private cachedDynamicBuckets: BucketDescription[] | null = null;
@@ -366,11 +383,12 @@ export class BucketParameterState {
         explicitStreamSubscriptions[subscription.stream] = subscription;
       }
     }
+    this.includeDefaultStreams = subscriptions?.include_defaults ?? true;
     this.explicitStreamSubscriptions = explicitStreamSubscriptions;
 
     this.querier = syncRules.getBucketParameterQuerier({
       globalParameters: this.syncParams,
-      hasDefaultSubscriptions: subscriptions?.include_defaults ?? true,
+      hasDefaultSubscriptions: this.includeDefaultStreams,
       resolveSubscription(name) {
         const subscription = explicitStreamSubscriptions[name];
         if (subscription) {
@@ -401,6 +419,10 @@ export class BucketParameterState {
     }
   }
 
+  isSubscribedToStream(desc: SqlBucketDescriptor): boolean {
+    return (desc.subscribedToByDefault && this.includeDefaultStreams) || desc.name in this.explicitStreamSubscriptions;
+  }
+
   async getCheckpointUpdate(checkpoint: storage.StorageCheckpointUpdate): Promise<CheckpointUpdate> {
     const querier = this.querier;
     let update: CheckpointUpdate;

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

@@ -146,10 +146,16 @@ export type StreamingSyncLine =
  */
 export type ProtocolOpId = string;
 
+export interface SubscribedStream {
+  name: string;
+  is_default: boolean;
+}
+
 export interface Checkpoint {
   last_op_id: ProtocolOpId;
   write_checkpoint?: ProtocolOpId;
   buckets: BucketChecksumWithDescription[];
+  included_subscriptions: SubscribedStream[];
 }
 
 export interface BucketState {

packages/sync-rules/src/SqlSyncRules.ts

@@ -1,9 +1,8 @@
 import { isScalar, LineCounter, parseDocument, Scalar, YAMLMap, YAMLSeq } from 'yaml';
-import { BucketPriority, isValidPriority } from './BucketDescription.js';
+import { isValidPriority } from './BucketDescription.js';
 import { BucketParameterQuerier, mergeBucketParameterQueriers } from './BucketParameterQuerier.js';
 import { SqlRuleError, SyncRulesErrors, YamlError } from './errors.js';
 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, SqlBucketDescriptorType } from './SqlBucketDescriptor.js';
@@ -375,16 +374,20 @@ export class SqlSyncRules implements SyncRules {
     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 (descriptor.type == SqlBucketDescriptorType.STREAM) {
+        const subscription = options.resolveSubscription(descriptor.name);
+
+        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);
+        }
 
-      if (subscription != null) {
-        params = params.withAddedParameters(subscription);
+        queriers.push(descriptor.getBucketParameterQuerier(params));
       }
 
       queriers.push(descriptor.getBucketParameterQuerier(params));