convex/src/server/impl/query_impl.ts

Client for the Convex Cloud

import { Value, JSONValue, jsonToConvex } from "../../values/index.js";
import { PaginationResult, PaginationOptions } from "../pagination.js";
import { performAsyncSyscall, performSyscall } from "./syscall.js";
import {
  filterBuilderImpl,
  serializeExpression,
} from "./filter_builder_impl.js";
import { Query, QueryInitializer } from "../query.js";
import { Expression, FilterBuilder } from "../filter_builder.js";
import { GenericTableInfo } from "../data_model.js";
import {
  IndexRangeBuilderImpl,
  SerializedRangeExpression,
} from "./index_range_builder_impl.js";

type QueryOperator = { filter: JSONValue } | { limit: number };
type Source =
  | { type: "FullTableScan"; tableName: string; order: "asc" | "desc" }
  | {
      type: "IndexRange";
      indexName: string;
      range: ReadonlyArray<SerializedRangeExpression>;
      order: "asc" | "desc";
    };
type SerializedQuery = {
  source: Source;
  operators: Array<QueryOperator>;
};

export class QueryInitializerImpl
  implements QueryInitializer<GenericTableInfo>
{
  private tableName: string;

  constructor(tableName: string) {
    this.tableName = tableName;
  }

  withIndex(
    indexName: string,
    indexRange: (q: IndexRangeBuilderImpl) => IndexRangeBuilderImpl
  ): QueryImpl {
    const rangeBuilder = indexRange(IndexRangeBuilderImpl.new());
    return new QueryImpl({
      source: {
        type: "IndexRange",
        indexName: this.tableName + "." + indexName,
        range: rangeBuilder.export(),
        order: "asc",
      },
      operators: [],
    });
  }

  fullTableScan(): QueryImpl {
    return new QueryImpl({
      source: {
        type: "FullTableScan",
        tableName: this.tableName,
        order: "asc",
      },
      operators: [],
    });
  }

  order(order: "asc" | "desc"): QueryImpl {
    return this.fullTableScan().order(order);
  }

  // This is internal API and should not be exposed to developers yet.
  async count(): Promise<number> {
    const syscallJSON = await performAsyncSyscall("count", {
      table: this.tableName,
    });
    const syscallResult = jsonToConvex(syscallJSON) as number;
    return syscallResult;
  }

  filter(
    predicate: (q: FilterBuilder<GenericTableInfo>) => Expression<boolean>
  ): QueryImpl {
    return this.fullTableScan().filter(predicate);
  }

  limit(n: number): QueryImpl {
    return this.fullTableScan().limit(n);
  }

  collect(): Promise<any[]> {
    return this.fullTableScan().collect();
  }

  take(n: number): Promise<Array<any>> {
    return this.fullTableScan().take(n);
  }

  paginate(options: PaginationOptions): Promise<PaginationResult<any>> {
    return this.fullTableScan().paginate(options);
  }

  first(): Promise<any> {
    return this.fullTableScan().first();
  }

  unique(): Promise<any> {
    return this.fullTableScan().unique();
  }

  [Symbol.asyncIterator](): AsyncIterableIterator<any> {
    return this.fullTableScan()[Symbol.asyncIterator]();
  }
}

/**
 * @param type Whether the query was consumed or closed.
 * @throws An error indicating the query has been closed.
 */
function throwClosedError(type: "closed" | "consumed"): never {
  throw new Error(
    type === "consumed"
      ? "This query is closed and can't emit any more values."
      : "This query has been chained with another operator and can't be reused."
  );
}

class QueryImpl implements Query<GenericTableInfo> {
  private state:
    | { type: "preparing"; query: SerializedQuery }
    | { type: "executing"; queryId: number }
    | { type: "closed" }
    | { type: "consumed" };

  constructor(query: SerializedQuery) {
    this.state = { type: "preparing", query };
  }

  private takeQuery(): SerializedQuery {
    if (this.state.type !== "preparing") {
      throw new Error(
        "A query can only be chained once and can't be chained after iteration begins."
      );
    }
    const query = this.state.query;
    this.state = { type: "closed" };
    return query;
  }

  private startQuery(): number {
    if (this.state.type === "executing") {
      throw new Error("Iteration can only begin on a query once.");
    }
    if (this.state.type === "closed" || this.state.type === "consumed") {
      throwClosedError(this.state.type);
    }
    const query = this.state.query;
    const { queryId } = performSyscall("queryStream", { query });
    this.state = { type: "executing", queryId };
    return queryId;
  }

  private closeQuery() {
    if (this.state.type === "executing") {
      const queryId = this.state.queryId;
      performSyscall("queryCleanup", { queryId });
    }
    this.state = { type: "consumed" };
  }

  order(order: "asc" | "desc"): QueryImpl {
    const query = this.takeQuery();
    query.source.order = order;
    return new QueryImpl(query);
  }

  filter(
    predicate: (q: FilterBuilder<GenericTableInfo>) => Expression<boolean>
  ): QueryImpl {
    const query = this.takeQuery();
    query.operators.push({
      filter: serializeExpression(predicate(filterBuilderImpl)),
    });
    return new QueryImpl(query);
  }

  limit(n: number): QueryImpl {
    const query = this.takeQuery();
    query.operators.push({ limit: n });
    return new QueryImpl(query);
  }

  [Symbol.asyncIterator](): AsyncIterableIterator<any> {
    this.startQuery();
    return this;
  }

  async next(): Promise<IteratorResult<any>> {
    if (this.state.type === "closed" || this.state.type === "consumed") {
      throwClosedError(this.state.type);
    }
    // Allow calling `.next()` when the query is in "preparing" state to implicitly start the
    // query. This allows the developer to call `.next()` on the query without having to use
    // a `for await` statement.
    const queryId =
      this.state.type === "preparing" ? this.startQuery() : this.state.queryId;
    const { value, done } = await performAsyncSyscall("queryStreamNext", {
      queryId,
    });
    if (done) {
      this.closeQuery();
    }
    const convexValue = jsonToConvex(value);
    return { value: convexValue, done };
  }

  return() {
    this.closeQuery();
    return Promise.resolve({ done: true, value: undefined });
  }

  async paginate(options: PaginationOptions): Promise<PaginationResult<any>> {
    const query = this.takeQuery();
    const pageSize = options.numItems;
    const cursor = options.cursor;
    const maximumRowsRead = options.maximumRowsRead ?? null;
    const { page, isDone, continueCursor } = await performAsyncSyscall(
      "queryPage",
      {
        query,
        cursor,
        pageSize,
        maximumRowsRead,
      }
    );
    return {
      page: page.map(jsonToConvex),
      isDone,
      continueCursor,
    };
  }

  async collect(): Promise<Array<any>> {
    const out: Value[] = [];
    for await (const item of this) {
      out.push(item);
    }
    return out;
  }

  async take(n: number): Promise<Array<any>> {
    return this.limit(n).collect();
  }

  async first(): Promise<any | null> {
    const first_array = await this.take(1);
    return first_array.length === 0 ? null : first_array[0];
  }

  async unique(): Promise<any | null> {
    const first_two_array = await this.take(2);
    if (first_two_array.length === 0) {
      return null;
    }
    if (first_two_array.length === 2) {
      throw new Error("unique() query returned more than one result");
    }
    return first_two_array[0];
  }
}