convex/src/browser/sync/local_state.ts

Client for the Convex Cloud

import { convexToJson } from "../../values/index.js";
import {
  AddQuery,
  RemoveQuery,
  QueryId,
  QuerySetModification,
  QuerySetVersion,
  IdentityVersion,
  Authenticate,
  QueryJournal,
  Transition,
} from "./protocol.js";
import {
  canonicalizeUdfPath,
  QueryToken,
  serializePathAndArgs,
} from "./udf_path_utils.js";

type LocalQuery = {
  id: QueryId;
  canonicalizedUdfPath: string;
  args: any[];
  numSubscribers: number;
  journal?: QueryJournal;
};

export class LocalSyncState {
  private nextQueryId: QueryId;
  private querySetVersion: QuerySetVersion;
  private readonly querySet: Map<QueryToken, LocalQuery>;
  private readonly queryIdToToken: Map<QueryId, QueryToken>;
  private identityVersion: IdentityVersion;
  private auth?: { tokenType: "Admin" | "User"; value: string };

  constructor() {
    this.nextQueryId = 0;
    this.querySetVersion = 0;
    this.identityVersion = 0;
    this.querySet = new Map();
    this.queryIdToToken = new Map();
  }

  subscribe(
    udfPath: string,
    args: any[],
    journal?: QueryJournal
  ): {
    queryToken: QueryToken;
    modification: QuerySetModification | null;
    unsubscribe: () => QuerySetModification | null;
  } {
    const canonicalizedUdfPath = canonicalizeUdfPath(udfPath);
    const queryToken = serializePathAndArgs(canonicalizedUdfPath, args);

    const existingEntry = this.querySet.get(queryToken);

    if (existingEntry !== undefined) {
      existingEntry.numSubscribers += 1;
      return {
        queryToken,
        modification: null,
        unsubscribe: () => this.removeSubscriber(queryToken),
      };
    } else {
      const queryId = this.nextQueryId++;
      const query: LocalQuery = {
        id: queryId,
        canonicalizedUdfPath,
        args,
        numSubscribers: 1,
        journal,
      };
      this.querySet.set(queryToken, query);
      this.queryIdToToken.set(queryId, queryToken);

      const baseVersion = this.querySetVersion;
      const newVersion = ++this.querySetVersion;

      const add: AddQuery = {
        type: "Add",
        queryId,
        udfPath: canonicalizedUdfPath,
        args: args.map(convexToJson),
        journal,
      };
      const modification: QuerySetModification = {
        type: "ModifyQuerySet",
        baseVersion,
        newVersion,
        modifications: [add],
      };
      return {
        queryToken,
        modification,
        unsubscribe: () => this.removeSubscriber(queryToken),
      };
    }
  }

  saveQueryJournals(transition: Transition) {
    for (const modification of transition.modifications) {
      switch (modification.type) {
        case "QueryUpdated":
        case "QueryFailed": {
          const journal = modification.journal;
          if (journal !== undefined) {
            const queryToken = this.queryIdToToken.get(modification.queryId);
            // We may have already unsubscribed to this query by the time the server
            // sends us the journal. If so, just ignore it.
            if (queryToken !== undefined) {
              this.querySet.get(queryToken)!.journal = journal;
            }
          }

          break;
        }
        case "QueryRemoved": {
          break;
        }
        default: {
          // Enforce that the switch-case is exhaustive.
          const _: never = modification;
          throw new Error(`Invalid modification ${modification}`);
        }
      }
    }
  }

  queryId(udfPath: string, args: any[]): QueryId | null {
    const canonicalizedUdfPath = canonicalizeUdfPath(udfPath);
    const queryToken = serializePathAndArgs(canonicalizedUdfPath, args);
    const existingEntry = this.querySet.get(queryToken);
    if (existingEntry !== undefined) {
      return existingEntry.id;
    }
    return null;
  }

  setAuth(value: string): Authenticate {
    this.auth = {
      tokenType: "User",
      value: value,
    };
    const baseVersion = this.identityVersion++;
    return {
      type: "Authenticate",
      baseVersion: baseVersion,
      ...this.auth,
    };
  }

  setAdminAuth(value: string): Authenticate {
    this.auth = {
      tokenType: "Admin",
      value: value,
    };
    const baseVersion = this.identityVersion++;
    return {
      type: "Authenticate",
      baseVersion: baseVersion,
      ...this.auth,
    };
  }

  clearAuth(): Authenticate {
    this.auth = undefined;
    const baseVersion = this.identityVersion++;
    return {
      type: "Authenticate",
      tokenType: "None",
      baseVersion: baseVersion,
    };
  }

  hasAuth(): boolean {
    return !!this.auth;
  }

  isNewAuth(value: string): boolean {
    return this.auth?.value !== value;
  }

  queryPath(queryId: QueryId): string | null {
    const pathAndArgs = this.queryIdToToken.get(queryId);
    if (pathAndArgs) {
      return this.querySet.get(pathAndArgs)!.canonicalizedUdfPath;
    }
    return null;
  }

  queryArgs(queryId: QueryId): any[] | null {
    const pathAndArgs = this.queryIdToToken.get(queryId);
    if (pathAndArgs) {
      return this.querySet.get(pathAndArgs)!.args;
    }
    return null;
  }

  queryToken(queryId: QueryId): string | null {
    return this.queryIdToToken.get(queryId) ?? null;
  }

  queryJournal(queryToken: QueryToken): QueryJournal | undefined {
    return this.querySet.get(queryToken)?.journal;
  }

  restart(): [QuerySetModification, Authenticate?] {
    const modifications = [];
    for (const localQuery of this.querySet.values()) {
      const add: AddQuery = {
        type: "Add",
        queryId: localQuery.id,
        udfPath: localQuery.canonicalizedUdfPath,
        args: localQuery.args.map(convexToJson),
        journal: localQuery.journal,
      };
      modifications.push(add);
    }
    this.querySetVersion = 1;
    const querySet: QuerySetModification = {
      type: "ModifyQuerySet",
      baseVersion: 0,
      newVersion: 1,
      modifications,
    };
    // If there's no auth, no need to send an update as the server will also start with an unknown identity.
    if (!this.auth) {
      this.identityVersion = 0;
      return [querySet, undefined];
    }
    const authenticate: Authenticate = {
      type: "Authenticate",
      baseVersion: 0,
      ...this.auth,
    };
    this.identityVersion = 1;
    return [querySet, authenticate];
  }

  private removeSubscriber(
    queryToken: QueryToken
  ): QuerySetModification | null {
    const localQuery = this.querySet.get(queryToken)!;

    if (localQuery.numSubscribers > 1) {
      localQuery.numSubscribers -= 1;
      return null;
    } else {
      this.querySet.delete(queryToken);
      this.queryIdToToken.delete(localQuery.id);
      const baseVersion = this.querySetVersion;
      const newVersion = ++this.querySetVersion;
      const remove: RemoveQuery = {
        type: "Remove",
        queryId: localQuery.id,
      };
      return {
        type: "ModifyQuerySet",
        baseVersion,
        newVersion,
        modifications: [remove],
      };
    }
  }
}