convex-helpers
Version:
A collection of useful code to complement the official convex package.
1,559 lines (1,511 loc) • 60.6 kB
text/typescript
/* eslint-disable no-unexpected-multiline */
import type { Value } from "convex/values";
import { convexToJson, compareValues, jsonToConvex } from "convex/values";
import type {
DataModelFromSchemaDefinition,
DocumentByInfo,
DocumentByName,
GenericDataModel,
GenericDatabaseReader,
IndexNames,
IndexRange,
IndexRangeBuilder,
NamedIndex,
NamedTableInfo,
OrderedQuery,
PaginationOptions,
PaginationResult,
Query,
QueryInitializer,
SchemaDefinition,
SystemDataModel,
TableNamesInDataModel,
} from "convex/server";
export type IndexKey = (Value | undefined)[];
//
// Helper functions
//
function makeExclusive(boundType: "gt" | "lt" | "gte" | "lte") {
if (boundType === "gt" || boundType === "gte") {
return "gt";
}
return "lt";
}
type Bound = ["gt" | "lt" | "gte" | "lte" | "eq", string, Value];
/** Split a range query between two index keys into a series of range queries
* that should be executed in sequence. This is necessary because Convex only
* supports range queries of the form
* q.eq("f1", v).eq("f2", v).lt("f3", v).gt("f3", v).
* i.e. all fields must be equal except for the last field, which can have
* two inequalities.
*
* For example, the range from >[1, 2, 3] to <=[1, 3, 2] would be split into
* the following queries:
* 1. q.eq("f1", 1).eq("f2", 2).gt("f3", 3)
* 2. q.eq("f1", 1).gt("f2", 2).lt("f2", 3)
* 3. q.eq("f1", 1).eq("f2", 3).lte("f3", 2)
*/
function splitRange(
indexFields: string[],
// For descending queries, the resulting queries are reversed.
order: "asc" | "desc",
startBound: IndexKey,
endBound: IndexKey,
startBoundType: "gt" | "lt" | "gte" | "lte",
endBoundType: "gt" | "lt" | "gte" | "lte",
): Bound[][] {
// Three parts to the split:
// 1. reduce down from startBound to common prefix
// 2. range with common prefix
// 3. build back up from common prefix to endBound
const commonPrefix: Bound[] = [];
while (
startBound.length > 0 &&
endBound.length > 0 &&
compareValues(startBound[0]!, endBound[0]!) === 0
) {
const indexField = indexFields[0]!;
indexFields = indexFields.slice(1);
const eqBound = startBound[0]!;
startBound = startBound.slice(1);
endBound = endBound.slice(1);
commonPrefix.push(["eq", indexField, eqBound]);
}
const makeCompare = (
boundType: "gt" | "lt" | "gte" | "lte",
key: IndexKey,
) => {
const range = commonPrefix.slice();
let i = 0;
for (; i < key.length - 1; i++) {
range.push(["eq", indexFields[i]!, key[i]!]);
}
if (i < key.length) {
range.push([boundType, indexFields[i]!, key[i]!]);
}
return range;
};
// Stage 1.
const startRanges: Bound[][] = [];
while (startBound.length > 1) {
startRanges.push(makeCompare(startBoundType, startBound));
startBoundType = makeExclusive(startBoundType);
startBound = startBound.slice(0, -1);
}
// Stage 3.
const endRanges: Bound[][] = [];
while (endBound.length > 1) {
endRanges.push(makeCompare(endBoundType, endBound));
endBoundType = makeExclusive(endBoundType);
endBound = endBound.slice(0, -1);
}
endRanges.reverse();
// Stage 2.
let middleRange;
if (endBound.length === 0) {
middleRange = makeCompare(startBoundType, startBound);
} else if (startBound.length === 0) {
middleRange = makeCompare(endBoundType, endBound);
} else {
const startValue = startBound[0]!;
const endValue = endBound[0]!;
middleRange = commonPrefix.slice();
middleRange.push([startBoundType, indexFields[0]!, startValue]);
middleRange.push([endBoundType, indexFields[0]!, endValue]);
}
const ranges = [...startRanges, middleRange, ...endRanges];
if (order === "desc") {
ranges.reverse();
}
return ranges;
}
function rangeToQuery(range: Bound[]) {
return (q: any) => {
for (const [boundType, field, value] of range) {
q = q[boundType](field, value);
}
return q;
};
}
/**
* Get the ordered list of fields for a given table's index based on the schema.
*
* - For "by_creation_time", returns ["_creationTime", "_id"].
* - For "by_id", returns ["_id"].
* - Otherwise, looks up the named index in the schema and returns its fields
* followed by ["_creationTime", "_id"].
* e.g. for an index defined like `.index("abc", ["a", "b"])`,
* returns ["a", "b", "_creationTime", "_id"].
*/
export function getIndexFields<
Schema extends SchemaDefinition<any, boolean>,
T extends TableNamesInDataModel<DM<Schema>>,
>(
table: T,
index?: IndexNames<NamedTableInfo<DM<Schema>, T>>,
schema?: Schema,
): string[] {
const indexDescriptor = String(index ?? "by_creation_time");
if (indexDescriptor === "by_creation_time") {
return ["_creationTime", "_id"];
}
if (indexDescriptor === "by_id") {
return ["_id"];
}
if (!schema) {
throw new Error("schema is required to infer index fields");
}
const tableInfo = schema.tables[table];
const indexInfo = tableInfo.indexes.find(
(index: any) => index.indexDescriptor === indexDescriptor,
);
if (!indexInfo) {
throw new Error(`Index ${indexDescriptor} not found in table ${table}`);
}
const fields = indexInfo.fields.slice();
fields.push("_creationTime");
fields.push("_id");
return fields;
}
function getIndexKey<
DataModel extends GenericDataModel,
T extends TableNamesInDataModel<DataModel>,
>(doc: DocumentByName<DataModel, T>, indexFields: string[]): IndexKey {
const key: IndexKey = [];
for (const field of indexFields) {
let obj: any = doc;
for (const subfield of field.split(".")) {
obj = obj[subfield];
}
key.push(obj);
}
return key;
}
/**
* A "stream" is an async iterable of query results, ordered by an index on a table.
*
* Use it as you would use `ctx.db`.
* If using pagination in a reactive query, see the warnings on the `paginator`
* function. TL;DR: you need to pass in `endCursor` to prevent holes or overlaps
* between pages.
*
* Once you have a stream, you can use `mergeStreams` or `filterStream` to make
* more streams. Then use `queryStream` to convert it into an OrderedQuery,
* so you can call `.paginate()`, `.collect()`, etc.
*/
export function stream<Schema extends SchemaDefinition<any, boolean>>(
db: GenericDatabaseReader<DM<Schema>>,
schema: Schema,
): StreamDatabaseReader<Schema> {
return new StreamDatabaseReader(db, schema);
}
type GenericStreamItem = NonNullable<unknown>;
/**
* A "QueryStream" is an async iterable of query results, ordered by indexed fields.
*/
export abstract class QueryStream<T extends GenericStreamItem>
implements GenericOrderedQuery<T>
{
// Methods that subclasses must implement so OrderedQuery can be implemented.
abstract iterWithKeys(): AsyncIterable<[T | null, IndexKey]>;
abstract narrow(indexBounds: IndexBounds): QueryStream<T>;
// Methods so subclasses can make sure streams are combined correctly.
abstract getOrder(): "asc" | "desc";
abstract getIndexFields(): string[];
// Values that must match a prefix of the index key.
abstract getEqualityIndexFilter(): Value[];
/// Methods for creating new streams as modifications of the current stream.
/**
* Create a new stream with a TypeScript filter applied.
*
* This is similar to `db.query(tableName).filter(predicate)`, but it's more
* general because it can call arbitrary TypeScript code, including more
* database queries.
*
* All documents filtered out are still considered "read" from the database;
* they are just excluded from the output stream.
*
* In contrast to `filter` from convex-helpers/server/filter, this filterWith
* is applied *before* any pagination. That means if the filter excludes a lot
* of documents, the `.paginate()` method will read a lot of documents until
* it gets as many documents as it wants. If you run into issues with reading
* too much data, you can pass `maximumRowsRead` to `paginate()`.
*/
filterWith(predicate: (doc: T) => Promise<boolean>): QueryStream<T> {
const order = this.getOrder();
return new FlatMapStream(
this,
async (doc: T) => {
const filtered = (await predicate(doc)) ? doc : null;
return new SingletonStream(filtered, order, [], [], []);
},
[],
);
}
/**
* Create a new stream where each element is the result of applying the mapper
* function to the elements of the original stream.
*
* Similar to how [1, 2, 3].map(x => x * 2) => [2, 4, 6]
*/
map<U extends GenericStreamItem>(
mapper: (doc: T) => Promise<U | null>,
): QueryStream<U> {
const order = this.getOrder();
return new FlatMapStream(
this,
async (doc: T) => {
const mapped = await mapper(doc);
return new SingletonStream(mapped, order, [], [], []);
},
[],
);
}
/**
* Similar to flatMap on an array, but iterate over a stream, and the for each
* element, iterate over the stream created by the mapper function.
*
* Ordered by the original stream order, then the mapped stream. Similar to
* how ["a", "b"].flatMap(letter => [letter, letter]) => ["a", "a", "b", "b"]
*
* @param mapper A function that takes a document and returns a new stream.
* @param mappedIndexFields The index fields of the streams created by mapper.
* @returns A stream of documents returned by the mapper streams,
* grouped by the documents in the original stream.
*/
flatMap<U extends GenericStreamItem>(
mapper: (doc: T) => Promise<QueryStream<U>>,
mappedIndexFields: string[],
): QueryStream<U> {
normalizeIndexFields(mappedIndexFields);
return new FlatMapStream(this, mapper, mappedIndexFields);
}
/**
* Get the first item from the original stream for each distinct value of the
* selected index fields.
*
* e.g. if the stream has an equality filter on `a`, and index fields `[a, b, c]`,
* we can do `stream.distinct(["b"])` to get a stream of the first item for
* each distinct value of `b`.
* Similarly, you could do `stream.distinct(["a", "b"])` with the same result,
* or `stream.distinct(["a", "b", "c"])` to get the original stream.
*
* This stream efficiently skips past items with the same value for the selected
* distinct index fields.
*
* This can be used to perform a loose index scan.
*/
distinct(distinctIndexFields: string[]): QueryStream<T> {
return new DistinctStream(this, distinctIndexFields);
}
/// Implementation of OrderedQuery
filter(_predicate: any): never {
throw new Error(
"Cannot call .filter() directly on a query stream. Use .filterWith() for filtering or .collect() if you want to convert the stream to an array first.",
);
}
async paginate(
opts: PaginationOptions & {
endCursor?: string | null;
maximumRowsRead?: number;
},
): Promise<PaginationResult<T>> {
if (opts.numItems === 0) {
if (opts.cursor === null) {
throw new Error(
".paginate called with cursor of null and 0 for numItems. " +
"This is not supported, as null is not a valid continueCursor. " +
"Advice: avoid calling paginate entirely in these cases.",
);
}
return {
page: [],
isDone: false,
continueCursor: opts.cursor,
};
}
const order = this.getOrder();
let newStartKey = {
key: [] as IndexKey,
inclusive: true,
};
if (opts.cursor !== null) {
newStartKey = {
key: deserializeCursor(opts.cursor),
inclusive: false,
};
}
let newEndKey = {
key: [] as IndexKey,
inclusive: true,
};
const maxRowsToRead = opts.maximumRowsRead;
const softMaxRowsToRead = opts.numItems + 1;
let maxRows: number | undefined = opts.numItems;
if (opts.endCursor) {
newEndKey = {
key: deserializeCursor(opts.endCursor),
inclusive: true,
};
// If there's an endCursor, continue until we get there even if it's more
// than numItems.
maxRows = undefined;
}
const newLowerBound = order === "asc" ? newStartKey : newEndKey;
const newUpperBound = order === "asc" ? newEndKey : newStartKey;
const narrowStream = this.narrow({
lowerBound: newLowerBound.key,
lowerBoundInclusive: newLowerBound.inclusive,
upperBound: newUpperBound.key,
upperBoundInclusive: newUpperBound.inclusive,
});
const page: T[] = [];
const indexKeys: IndexKey[] = [];
let hasMore = opts.endCursor && opts.endCursor !== "[]";
let continueCursor = opts.endCursor ?? "[]";
for await (const [doc, indexKey] of narrowStream.iterWithKeys()) {
if (doc !== null) {
page.push(doc);
}
indexKeys.push(indexKey);
if (
(maxRows !== undefined && page.length >= maxRows) ||
(maxRowsToRead !== undefined && indexKeys.length >= maxRowsToRead)
) {
hasMore = true;
continueCursor = serializeCursor(indexKey);
break;
}
}
let pageStatus: "SplitRecommended" | "SplitRequired" | undefined =
undefined;
let splitCursor: IndexKey | undefined = undefined;
if (indexKeys.length === maxRowsToRead) {
pageStatus = "SplitRequired";
splitCursor = indexKeys[Math.floor((indexKeys.length - 1) / 2)];
} else if (indexKeys.length >= softMaxRowsToRead) {
pageStatus = "SplitRecommended";
splitCursor = indexKeys[Math.floor((indexKeys.length - 1) / 2)];
}
return {
page,
isDone: !hasMore,
continueCursor,
pageStatus,
splitCursor: splitCursor ? serializeCursor(splitCursor) : undefined,
};
}
async collect() {
return await this.take(Infinity);
}
async take(n: number) {
const results: T[] = [];
for await (const [doc, _] of this.iterWithKeys()) {
if (doc === null) {
continue;
}
results.push(doc);
if (results.length === n) {
break;
}
}
return results;
}
async unique() {
const docs = await this.take(2);
if (docs.length === 2) {
throw new Error("Query is not unique");
}
return docs[0] ?? null;
}
async first() {
const docs = await this.take(1);
return docs[0] ?? null;
}
[Symbol.asyncIterator]() {
const iterator = this.iterWithKeys()[Symbol.asyncIterator]();
return {
async next() {
const result = await iterator.next();
if (result.done) {
return { done: true as const, value: undefined };
}
return { done: false, value: result.value[0]! };
},
};
}
}
/**
* GenericOrderedQuery<DocumentByInfo<TableInfo>> is equivalent to OrderedQuery<TableInfo>
*/
export interface GenericOrderedQuery<T> extends AsyncIterable<T> {
/**
* Load a page of `n` results and obtain a {@link Cursor} for loading more.
*
* Note: If this is called from a reactive query function the number of
* results may not match `paginationOpts.numItems`!
*
* `paginationOpts.numItems` is only an initial value. After the first invocation,
* `paginate` will return all items in the original query range. This ensures
* that all pages will remain adjacent and non-overlapping.
*
* @param paginationOpts - A {@link PaginationOptions} object containing the number
* of items to load and the cursor to start at.
* @returns A {@link PaginationResult} containing the page of results and a
* cursor to continue paginating.
*/
paginate(paginationOpts: PaginationOptions): Promise<PaginationResult<T>>;
/**
* Execute the query and return all of the results as an array.
*
* Note: when processing a query with a lot of results, it's often better to use the `Query` as an
* `AsyncIterable` instead.
*
* @returns - An array of all of the query's results.
*/
collect(): Promise<Array<T>>;
/**
* Execute the query and return the first `n` results.
*
* @param n - The number of items to take.
* @returns - An array of the first `n` results of the query (or less if the
* query doesn't have `n` results).
*/
take(n: number): Promise<Array<T>>;
/**
* Execute the query and return the first result if there is one.
*
* @returns - The first value of the query or `null` if the query returned no results.
* */
first(): Promise<T | null>;
/**
* Execute the query and return the singular result if there is one.
*
* @returns - The single result returned from the query or null if none exists.
* @throws Will throw an error if the query returns more than one result.
*/
unique(): Promise<T | null>;
/**
* Not supported. Use `filterWith` instead.
*/
filter(predicate: any): this;
}
export class StreamDatabaseReader<Schema extends SchemaDefinition<any, boolean>>
implements GenericDatabaseReader<DM<Schema>>
{
// TODO: support system tables
public system: GenericDatabaseReader<SystemDataModel>["system"];
constructor(
public db: GenericDatabaseReader<DM<Schema>>,
public schema: Schema,
) {
this.system = db.system;
}
query<TableName extends TableNamesInDataModel<DM<Schema>>>(
tableName: TableName,
): StreamQueryInitializer<Schema, TableName> {
return new StreamQueryInitializer(this, tableName);
}
get(_id: any): any {
throw new Error("get() not supported for `paginator`");
}
normalizeId(_tableName: any, _id: any): any {
throw new Error("normalizeId() not supported for `paginator`.");
}
}
type DM<Schema extends SchemaDefinition<any, boolean>> =
DataModelFromSchemaDefinition<Schema>;
export type IndexBounds = {
lowerBound: IndexKey;
lowerBoundInclusive: boolean;
upperBound: IndexKey;
upperBoundInclusive: boolean;
};
export type QueryReflection<
Schema extends SchemaDefinition<any, boolean>,
T extends TableNamesInDataModel<DM<Schema>>,
IndexName extends IndexNames<NamedTableInfo<DM<Schema>, T>>,
> = {
db: GenericDatabaseReader<DataModelFromSchemaDefinition<Schema>>;
schema: Schema;
table: T;
index: IndexName;
indexFields: string[];
order: "asc" | "desc";
bounds: IndexBounds;
indexRange?: (
q: IndexRangeBuilder<
DocumentByInfo<NamedTableInfo<DM<Schema>, T>>,
NamedIndex<NamedTableInfo<DM<Schema>, T>, IndexName>
>,
) => IndexRange;
};
export abstract class StreamableQuery<
Schema extends SchemaDefinition<any, boolean>,
T extends TableNamesInDataModel<DM<Schema>>,
IndexName extends IndexNames<NamedTableInfo<DM<Schema>, T>>,
>
extends QueryStream<DocumentByInfo<NamedTableInfo<DM<Schema>, T>>>
// this "implements" is redundant, since QueryStream implies it, but it acts as a type-time assertion.
implements OrderedQuery<NamedTableInfo<DM<Schema>, T>>
{
abstract reflect(): QueryReflection<Schema, T, IndexName>;
}
export class StreamQueryInitializer<
Schema extends SchemaDefinition<any, boolean>,
T extends TableNamesInDataModel<DM<Schema>>,
>
extends StreamableQuery<Schema, T, "by_creation_time">
implements QueryInitializer<NamedTableInfo<DM<Schema>, T>>
{
constructor(
public parent: StreamDatabaseReader<Schema>,
public table: T,
) {
super();
}
fullTableScan(): StreamQuery<Schema, T, "by_creation_time"> {
return this.withIndex("by_creation_time");
}
withIndex<IndexName extends IndexNames<NamedTableInfo<DM<Schema>, T>>>(
indexName: IndexName,
indexRange?: (
q: IndexRangeBuilder<
DocumentByInfo<NamedTableInfo<DM<Schema>, T>>,
NamedIndex<NamedTableInfo<DM<Schema>, T>, IndexName>
>,
) => IndexRange,
): StreamQuery<Schema, T, IndexName> {
const indexFields = getIndexFields<Schema, T>(
this.table,
indexName,
this.parent.schema,
);
const q = new ReflectIndexRange(indexFields);
if (indexRange) {
indexRange(q as any);
}
return new StreamQuery(this, indexName, q, indexRange);
}
withSearchIndex(_indexName: any, _searchFilter: any): any {
throw new Error("Cannot paginate withSearchIndex");
}
inner() {
return this.fullTableScan();
}
order(
order: "asc" | "desc",
): OrderedStreamQuery<Schema, T, "by_creation_time"> {
return this.inner().order(order);
}
reflect() {
return this.inner().reflect();
}
iterWithKeys() {
return this.inner().iterWithKeys();
}
getOrder(): "asc" | "desc" {
return this.inner().getOrder();
}
getEqualityIndexFilter(): Value[] {
return this.inner().getEqualityIndexFilter();
}
getIndexFields(): string[] {
return this.inner().getIndexFields();
}
narrow(indexBounds: IndexBounds) {
return this.inner().narrow(indexBounds);
}
}
// Not to be confused with QueryStream or StreamableQuery.
export class StreamQuery<
Schema extends SchemaDefinition<any, boolean>,
T extends TableNamesInDataModel<DM<Schema>>,
IndexName extends IndexNames<NamedTableInfo<DM<Schema>, T>>,
>
extends StreamableQuery<Schema, T, IndexName>
implements Query<NamedTableInfo<DM<Schema>, T>>
{
constructor(
public parent: StreamQueryInitializer<Schema, T>,
public index: IndexName,
public q: ReflectIndexRange,
public indexRange:
| ((
q: IndexRangeBuilder<
DocumentByInfo<NamedTableInfo<DM<Schema>, T>>,
NamedIndex<NamedTableInfo<DM<Schema>, T>, IndexName>
>,
) => IndexRange)
| undefined,
) {
super();
}
order(order: "asc" | "desc") {
return new OrderedStreamQuery(this, order);
}
inner() {
return this.order("asc");
}
reflect() {
return this.inner().reflect();
}
iterWithKeys() {
return this.inner().iterWithKeys();
}
getOrder(): "asc" | "desc" {
return this.inner().getOrder();
}
getEqualityIndexFilter(): Value[] {
return this.inner().getEqualityIndexFilter();
}
getIndexFields(): string[] {
return this.inner().getIndexFields();
}
narrow(indexBounds: IndexBounds) {
return this.inner().narrow(indexBounds);
}
}
export class OrderedStreamQuery<
Schema extends SchemaDefinition<any, boolean>,
T extends TableNamesInDataModel<DM<Schema>>,
IndexName extends IndexNames<NamedTableInfo<DM<Schema>, T>>,
>
extends StreamableQuery<Schema, T, IndexName>
implements OrderedQuery<NamedTableInfo<DM<Schema>, T>>
{
constructor(
public parent: StreamQuery<Schema, T, IndexName>,
public order: "asc" | "desc",
) {
super();
}
reflect() {
return {
db: this.parent.parent.parent.db,
schema: this.parent.parent.parent.schema,
table: this.parent.parent.table,
index: this.parent.index,
indexFields: this.parent.q.indexFields,
order: this.order,
bounds: {
lowerBound: this.parent.q.lowerBoundIndexKey ?? [],
lowerBoundInclusive: this.parent.q.lowerBoundInclusive,
upperBound: this.parent.q.upperBoundIndexKey ?? [],
upperBoundInclusive: this.parent.q.upperBoundInclusive,
},
indexRange: this.parent.indexRange,
};
}
/**
* inner() is as if you had used ctx.db to construct the query.
*/
inner(): OrderedQuery<NamedTableInfo<DM<Schema>, T>> {
const { db, table, index, order, indexRange } = this.reflect();
return db.query(table).withIndex(index, indexRange).order(order);
}
iterWithKeys(): AsyncIterable<
[DocumentByName<DM<Schema>, T> | null, IndexKey]
> {
const { indexFields } = this.reflect();
const iterable = this.inner();
return {
[Symbol.asyncIterator]() {
const iterator = iterable[Symbol.asyncIterator]();
return {
async next() {
const result = await iterator.next();
if (result.done) {
return { done: true, value: undefined };
}
return {
done: false,
value: [result.value, getIndexKey(result.value, indexFields)],
};
},
};
},
};
}
getOrder(): "asc" | "desc" {
return this.order;
}
getEqualityIndexFilter(): Value[] {
return this.parent.q.equalityIndexFilter;
}
getIndexFields(): string[] {
return this.parent.q.indexFields;
}
narrow(indexBounds: IndexBounds) {
const { db, table, index, order, bounds, schema } = this.reflect();
let maxLowerBound = bounds.lowerBound;
let maxLowerBoundInclusive = bounds.lowerBoundInclusive;
if (
compareKeys(
{
value: indexBounds.lowerBound,
kind: indexBounds.lowerBoundInclusive ? "predecessor" : "successor",
},
{
value: bounds.lowerBound,
kind: bounds.lowerBoundInclusive ? "predecessor" : "successor",
},
) > 0
) {
maxLowerBound = indexBounds.lowerBound;
maxLowerBoundInclusive = indexBounds.lowerBoundInclusive;
}
let minUpperBound = bounds.upperBound;
let minUpperBoundInclusive = bounds.upperBoundInclusive;
if (
compareKeys(
{
value: indexBounds.upperBound,
kind: indexBounds.upperBoundInclusive ? "successor" : "predecessor",
},
{
value: bounds.upperBound,
kind: bounds.upperBoundInclusive ? "successor" : "predecessor",
},
) < 0
) {
minUpperBound = indexBounds.upperBound;
minUpperBoundInclusive = indexBounds.upperBoundInclusive;
}
return streamIndexRange(
db,
schema,
table,
index,
{
lowerBound: maxLowerBound,
lowerBoundInclusive: maxLowerBoundInclusive,
upperBound: minUpperBound,
upperBoundInclusive: minUpperBoundInclusive,
},
order,
);
}
}
/**
* Create a stream of documents using the given index and bounds.
*/
export function streamIndexRange<
Schema extends SchemaDefinition<any, boolean>,
T extends TableNamesInDataModel<DM<Schema>>,
IndexName extends IndexNames<NamedTableInfo<DM<Schema>, T>>,
>(
db: GenericDatabaseReader<DM<Schema>>,
schema: Schema,
table: T,
index: IndexName,
bounds: IndexBounds,
order: "asc" | "desc",
): QueryStream<DocumentByName<DM<Schema>, T>> {
const indexFields = getIndexFields(table, index, schema);
const splitBounds = splitRange(
indexFields,
order,
bounds.lowerBound,
bounds.upperBound,
bounds.lowerBoundInclusive ? "gte" : "gt",
bounds.upperBoundInclusive ? "lte" : "lt",
);
const subQueries = splitBounds.map((splitBound) =>
stream(db, schema)
.query(table)
.withIndex(index, rangeToQuery(splitBound))
.order(order),
);
return new ConcatStreams(...subQueries);
}
class ReflectIndexRange {
#hasSuffix = false;
public lowerBoundIndexKey: IndexKey | undefined = undefined;
public lowerBoundInclusive: boolean = true;
public upperBoundIndexKey: IndexKey | undefined = undefined;
public upperBoundInclusive: boolean = true;
public equalityIndexFilter: Value[] = [];
constructor(public indexFields: string[]) {}
eq(field: string, value: Value) {
if (!this.#canLowerBound(field) || !this.#canUpperBound(field)) {
throw new Error(`Cannot use eq on field '${field}'`);
}
this.lowerBoundIndexKey = this.lowerBoundIndexKey ?? [];
this.lowerBoundIndexKey.push(value);
this.upperBoundIndexKey = this.upperBoundIndexKey ?? [];
this.upperBoundIndexKey.push(value);
this.equalityIndexFilter.push(value);
return this;
}
lt(field: string, value: Value) {
if (!this.#canUpperBound(field)) {
throw new Error(`Cannot use lt on field '${field}'`);
}
this.upperBoundIndexKey = this.upperBoundIndexKey ?? [];
this.upperBoundIndexKey.push(value);
this.upperBoundInclusive = false;
this.#hasSuffix = true;
return this;
}
lte(field: string, value: Value) {
if (!this.#canUpperBound(field)) {
throw new Error(`Cannot use lte on field '${field}'`);
}
this.upperBoundIndexKey = this.upperBoundIndexKey ?? [];
this.upperBoundIndexKey.push(value);
this.#hasSuffix = true;
return this;
}
gt(field: string, value: Value) {
if (!this.#canLowerBound(field)) {
throw new Error(`Cannot use gt on field '${field}'`);
}
this.lowerBoundIndexKey = this.lowerBoundIndexKey ?? [];
this.lowerBoundIndexKey.push(value);
this.lowerBoundInclusive = false;
this.#hasSuffix = true;
return this;
}
gte(field: string, value: Value) {
if (!this.#canLowerBound(field)) {
throw new Error(`Cannot use gte on field '${field}'`);
}
this.lowerBoundIndexKey = this.lowerBoundIndexKey ?? [];
this.lowerBoundIndexKey.push(value);
this.#hasSuffix = true;
return this;
}
#canLowerBound(field: string) {
const currentLowerBoundLength = this.lowerBoundIndexKey?.length ?? 0;
const currentUpperBoundLength = this.upperBoundIndexKey?.length ?? 0;
if (currentLowerBoundLength > currentUpperBoundLength) {
// Already have a lower bound.
return false;
}
if (
currentLowerBoundLength === currentUpperBoundLength &&
this.#hasSuffix
) {
// Already have a lower bound and an upper bound.
return false;
}
return (
currentLowerBoundLength < this.indexFields.length &&
this.indexFields[currentLowerBoundLength] === field
);
}
#canUpperBound(field: string) {
const currentLowerBoundLength = this.lowerBoundIndexKey?.length ?? 0;
const currentUpperBoundLength = this.upperBoundIndexKey?.length ?? 0;
if (currentUpperBoundLength > currentLowerBoundLength) {
// Already have an upper bound.
return false;
}
if (
currentLowerBoundLength === currentUpperBoundLength &&
this.#hasSuffix
) {
// Already have a lower bound and an upper bound.
return false;
}
return (
currentUpperBoundLength < this.indexFields.length &&
this.indexFields[currentUpperBoundLength] === field
);
}
}
/**
* Merge multiple streams, provided in any order, into a single stream.
*
* The streams will be merged into a stream of documents ordered by the index keys,
* i.e. by "author" (then by the implicit "_creationTime").
*
* e.g. ```ts
* mergedStream([
* stream(db, schema).query("messages").withIndex("by_author", q => q.eq("author", "user3")),
* stream(db, schema).query("messages").withIndex("by_author", q => q.eq("author", "user1")),
* stream(db, schema).query("messages").withIndex("by_author", q => q.eq("author", "user2")),
* ], ["author"])
* ```
*
* returns a stream of messages for user1, then user2, then user3.
*
* You can also use `orderByIndexFields` to change the indexed fields before merging, which changes the order of the merged stream.
* This only works if the streams are already ordered by `orderByIndexFields`,
* which happens if each does a .eq(field, value) on all index fields before `orderByIndexFields`.
*
* e.g. if the "by_author" index is defined as being ordered by ["author", "_creationTime"],
* and each query does an equality lookup on "author", each individual query before merging is in fact ordered by "_creationTime".
*
* e.g. ```ts
* mergedStream([
* stream(db, schema).query("messages").withIndex("by_author", q => q.eq("author", "user3")),
* stream(db, schema).query("messages").withIndex("by_author", q => q.eq("author", "user1")),
* stream(db, schema).query("messages").withIndex("by_author", q => q.eq("author", "user2")),
* ], ["_creationTime"])
* ```
*
* This returns a stream of messages from all three users, sorted by creation time.
*/
export function mergedStream<T extends GenericStreamItem>(
streams: QueryStream<T>[],
orderByIndexFields: string[],
): QueryStream<T> {
return new MergedStream(streams, orderByIndexFields);
}
export class MergedStream<T extends GenericStreamItem> extends QueryStream<T> {
#order: "asc" | "desc";
#streams: QueryStream<T>[];
#equalityIndexFilter: Value[];
#indexFields: string[];
constructor(streams: QueryStream<T>[], orderByIndexFields: string[]) {
super();
if (streams.length === 0) {
throw new Error("Cannot union empty array of streams");
}
this.#order = allSame(
streams.map((stream) => stream.getOrder()),
"Cannot merge streams with different orders",
);
this.#streams = streams.map(
(stream) => new OrderByStream(stream, orderByIndexFields),
);
this.#indexFields = allSame(
this.#streams.map((stream) => stream.getIndexFields()),
"Cannot merge streams with different index fields. Consider using .orderBy()",
);
// Calculate common prefix of equality index filters.
this.#equalityIndexFilter = commonPrefix(
this.#streams.map((stream) => stream.getEqualityIndexFilter()),
);
}
iterWithKeys() {
const iterables = this.#streams.map((stream) => stream.iterWithKeys());
const comparisonInversion = this.#order === "asc" ? 1 : -1;
return {
[Symbol.asyncIterator]() {
const iterators = iterables.map((iterable) =>
iterable[Symbol.asyncIterator](),
);
const results = Array.from(
{ length: iterators.length },
(): IteratorResult<[T | null, IndexKey] | undefined> => ({
done: false,
value: undefined,
}),
);
return {
async next() {
// Fill results from iterators with no value yet.
await Promise.all(
iterators.map(async (iterator, i) => {
if (!results[i]!.done && !results[i]!.value) {
const result = await iterator.next();
results[i] = result;
}
}),
);
// Find index for the value with the lowest index key.
let minIndexKeyAndIndex: [IndexKey, number] | undefined = undefined;
for (let i = 0; i < results.length; i++) {
const result = results[i]!;
if (result.done || !result.value) {
continue;
}
const [_, resultIndexKey] = result.value;
if (minIndexKeyAndIndex === undefined) {
minIndexKeyAndIndex = [resultIndexKey, i];
continue;
}
const [prevMin, _prevMinIndex] = minIndexKeyAndIndex;
if (
compareKeys(
{ value: resultIndexKey, kind: "exact" },
{ value: prevMin, kind: "exact" },
) *
comparisonInversion <
0
) {
minIndexKeyAndIndex = [resultIndexKey, i];
}
}
if (minIndexKeyAndIndex === undefined) {
return { done: true, value: undefined };
}
const [_, minIndex] = minIndexKeyAndIndex;
const result = results[minIndex]!.value;
// indicate that we've used this result
results[minIndex]!.value = undefined;
return { done: false, value: result };
},
};
},
};
}
getOrder(): "asc" | "desc" {
return this.#order;
}
getEqualityIndexFilter(): Value[] {
return this.#equalityIndexFilter;
}
getIndexFields(): string[] {
return this.#indexFields;
}
narrow(indexBounds: IndexBounds) {
return new MergedStream(
this.#streams.map((stream) => stream.narrow(indexBounds)),
this.#indexFields,
);
}
}
function allSame<T extends Value>(values: T[], errorMessage: string): T {
const first = values[0]!;
for (const value of values) {
if (compareValues(value, first)) {
throw new Error(errorMessage);
}
}
return first;
}
function commonPrefix(values: Value[][]) {
let commonPrefix = values[0]!;
for (const value of values) {
for (let i = 0; i < commonPrefix.length; i++) {
if (i >= value.length || compareValues(commonPrefix[i], value[i])) {
commonPrefix = commonPrefix.slice(0, i);
break;
}
}
}
return commonPrefix;
}
/**
* Concatenate multiple streams into a single stream.
* This assumes that the streams correspond to disjoint index ranges,
* and are provided in the same order as the index ranges.
*
* e.g. ```ts
* new ConcatStreams(
* stream(db, schema).query("messages").withIndex("by_author", q => q.eq("author", "user1")),
* stream(db, schema).query("messages").withIndex("by_author", q => q.eq("author", "user2")),
* )
* ```
*
* is valid, but if the stream arguments were reversed, or the queries were
* `.order("desc")`, it would be invalid.
*
* It's not recommended to use `ConcatStreams` directly, since it has the same
* behavior as `MergedStream`, but with fewer runtime checks.
*/
class ConcatStreams<T extends GenericStreamItem> extends QueryStream<T> {
#order: "asc" | "desc";
#streams: QueryStream<T>[];
#equalityIndexFilter: Value[];
#indexFields: string[];
constructor(...streams: QueryStream<T>[]) {
super();
this.#streams = streams;
if (streams.length === 0) {
throw new Error("Cannot concat empty array of streams");
}
this.#order = allSame(
streams.map((stream) => stream.getOrder()),
"Cannot concat streams with different orders. Consider using .orderBy()",
);
this.#indexFields = allSame(
streams.map((stream) => stream.getIndexFields()),
"Cannot concat streams with different index fields. Consider using .orderBy()",
);
this.#equalityIndexFilter = commonPrefix(
streams.map((stream) => stream.getEqualityIndexFilter()),
);
}
iterWithKeys(): AsyncIterable<[T | null, IndexKey]> {
const iterables = this.#streams.map((stream) => stream.iterWithKeys());
const comparisonInversion = this.#order === "asc" ? 1 : -1;
let previousIndexKey: IndexKey | undefined = undefined;
return {
[Symbol.asyncIterator]() {
const iterators = iterables.map((iterable) =>
iterable[Symbol.asyncIterator](),
);
return {
async next() {
while (iterators.length > 0) {
const result = await iterators[0]!.next();
if (result.done) {
iterators.shift();
} else {
const [_, indexKey] = result.value;
if (
previousIndexKey !== undefined &&
compareKeys(
{
value: previousIndexKey,
kind: "exact",
},
{
value: indexKey,
kind: "exact",
},
) *
comparisonInversion >
0
) {
throw new Error(
`ConcatStreams in wrong order: ${JSON.stringify(previousIndexKey)}, ${JSON.stringify(indexKey)}`,
);
}
previousIndexKey = indexKey;
return result;
}
}
return { done: true, value: undefined };
},
};
},
};
}
getOrder(): "asc" | "desc" {
return this.#order;
}
getEqualityIndexFilter(): Value[] {
return this.#equalityIndexFilter;
}
getIndexFields(): string[] {
return this.#indexFields;
}
narrow(indexBounds: IndexBounds) {
return new ConcatStreams(
...this.#streams.map((stream) => stream.narrow(indexBounds)),
);
}
}
class FlatMapStreamIterator<
T extends GenericStreamItem,
U extends GenericStreamItem,
> implements AsyncIterator<[U | null, IndexKey]>
{
#outerStream: QueryStream<T>;
#outerIterator: AsyncIterator<[T | null, IndexKey]>;
#currentOuterItem: {
t: T | null;
indexKey: IndexKey;
innerIterator: AsyncIterator<[U | null, IndexKey]>;
count: number;
} | null = null;
#mapper: (doc: T) => Promise<QueryStream<U>>;
#mappedIndexFields: string[];
constructor(
outerStream: QueryStream<T>,
mapper: (doc: T) => Promise<QueryStream<U>>,
mappedIndexFields: string[],
) {
this.#outerIterator = outerStream.iterWithKeys()[Symbol.asyncIterator]();
this.#outerStream = outerStream;
this.#mapper = mapper;
this.#mappedIndexFields = mappedIndexFields;
}
singletonSkipInnerStream(): QueryStream<U> {
// If the outer stream is a filtered value, yield a singleton
// filtered value from the inner stream, with index key of nulls.
const indexKey = this.#mappedIndexFields.map(() => null);
return new SingletonStream<U>(
null,
this.#outerStream.getOrder(),
this.#mappedIndexFields,
indexKey,
indexKey,
);
}
async setCurrentOuterItem(item: [T | null, IndexKey]) {
const [t, indexKey] = item;
let innerStream: QueryStream<U>;
if (t === null) {
innerStream = this.singletonSkipInnerStream();
} else {
innerStream = await this.#mapper(t);
if (
!equalIndexFields(innerStream.getIndexFields(), this.#mappedIndexFields)
) {
throw new Error(
`FlatMapStream: inner stream has different index fields than expected: ${JSON.stringify(innerStream.getIndexFields())} vs ${JSON.stringify(this.#mappedIndexFields)}`,
);
}
if (innerStream.getOrder() !== this.#outerStream.getOrder()) {
throw new Error(
`FlatMapStream: inner stream has different order than outer stream: ${innerStream.getOrder()} vs ${this.#outerStream.getOrder()}`,
);
}
}
this.#currentOuterItem = {
t,
indexKey,
innerIterator: innerStream.iterWithKeys()[Symbol.asyncIterator](),
count: 0,
};
}
async next(): Promise<IteratorResult<[U | null, IndexKey]>> {
if (this.#currentOuterItem === null) {
const result = await this.#outerIterator.next();
if (result.done) {
return { done: true, value: undefined };
}
await this.setCurrentOuterItem(result.value);
return await this.next();
}
const result = await this.#currentOuterItem.innerIterator.next();
if (result.done) {
if (this.#currentOuterItem.count > 0) {
this.#currentOuterItem = null;
} else {
// The inner stream was completely empty, so we should inject a null
// (which will be skipped by everything except the maximumRowsRead count)
// to account for the cost of the outer stream.
this.#currentOuterItem.innerIterator = this.singletonSkipInnerStream()
.iterWithKeys()
[Symbol.asyncIterator]();
}
return await this.next();
}
const [u, indexKey] = result.value;
this.#currentOuterItem.count++;
const fullIndexKey = [...this.#currentOuterItem.indexKey, ...indexKey];
return { done: false, value: [u, fullIndexKey] };
}
}
class FlatMapStream<
T extends GenericStreamItem,
U extends GenericStreamItem,
> extends QueryStream<U> {
#stream: QueryStream<T>;
#mapper: (doc: T) => Promise<QueryStream<U>>;
#mappedIndexFields: string[];
constructor(
stream: QueryStream<T>,
mapper: (doc: T) => Promise<QueryStream<U>>,
mappedIndexFields: string[],
) {
super();
this.#stream = stream;
this.#mapper = mapper;
this.#mappedIndexFields = mappedIndexFields;
}
iterWithKeys(): AsyncIterable<[U | null, IndexKey]> {
const outerStream = this.#stream;
const mapper = this.#mapper;
const mappedIndexFields = this.#mappedIndexFields;
return {
[Symbol.asyncIterator]() {
return new FlatMapStreamIterator(
outerStream,
mapper,
mappedIndexFields,
);
},
};
}
getOrder(): "asc" | "desc" {
return this.#stream.getOrder();
}
getEqualityIndexFilter(): Value[] {
return this.#stream.getEqualityIndexFilter();
}
getIndexFields(): string[] {
return [...this.#stream.getIndexFields(), ...this.#mappedIndexFields];
}
narrow(indexBounds: IndexBounds) {
const outerLength = this.#stream.getIndexFields().length;
const outerLowerBound = indexBounds.lowerBound.slice(0, outerLength);
const outerUpperBound = indexBounds.upperBound.slice(0, outerLength);
const innerLowerBound = indexBounds.lowerBound.slice(outerLength);
const innerUpperBound = indexBounds.upperBound.slice(outerLength);
const outerIndexBounds = {
lowerBound: outerLowerBound,
lowerBoundInclusive:
innerLowerBound.length === 0 ? indexBounds.lowerBoundInclusive : true,
upperBound: outerUpperBound,
upperBoundInclusive:
innerUpperBound.length === 0 ? indexBounds.upperBoundInclusive : true,
};
const innerIndexBounds = {
lowerBound: innerLowerBound,
lowerBoundInclusive:
innerLowerBound.length === 0 ? true : indexBounds.lowerBoundInclusive,
upperBound: innerUpperBound,
upperBoundInclusive:
innerUpperBound.length === 0 ? true : indexBounds.upperBoundInclusive,
};
return new FlatMapStream(
this.#stream.narrow(outerIndexBounds),
async (t) => {
const innerStream = await this.#mapper(t);
return innerStream.narrow(innerIndexBounds);
},
this.#mappedIndexFields,
);
}
}
export class SingletonStream<
T extends GenericStreamItem,
> extends QueryStream<T> {
#value: T | null;
#order: "asc" | "desc";
#indexFields: string[];
#indexKey: IndexKey;
#equalityIndexFilter: Value[];
constructor(
value: T | null,
order: "asc" | "desc" = "asc",
indexFields: string[],
indexKey: IndexKey,
equalityIndexFilter: Value[],
) {
super();
this.#value = value;
this.#order = order;
this.#indexFields = indexFields;
this.#indexKey = indexKey;
this.#equalityIndexFilter = equalityIndexFilter;
if (indexKey.length !== indexFields.length) {
throw new Error(
`indexKey must have the same length as indexFields: ${JSON.stringify(
indexKey,
)} vs ${JSON.stringify(indexFields)}`,
);
}
}
iterWithKeys(): AsyncIterable<[T | null, IndexKey]> {
const value = this.#value;
const indexKey = this.#indexKey;
return {
[Symbol.asyncIterator]() {
let sent = false;
return {
async next() {
if (sent) {
return { done: true, value: undefined };
}
sent = true;
return { done: false, value: [value, indexKey] };
},
};
},
};
}
getOrder(): "asc" | "desc" {
return this.#order;
}
getIndexFields(): string[] {
return this.#indexFields;
}
getEqualityIndexFilter(): Value[] {
return this.#equalityIndexFilter;
}
narrow(indexBounds: IndexBounds): QueryStream<T> {
const compareLowerBound = compareKeys(
{
value: indexBounds.lowerBound,
kind: indexBounds.lowerBoundInclusive ? "exact" : "successor",
},
{
value: this.#indexKey,
kind: "exact",
},
);
const compareUpperBound = compareKeys(
{
value: this.#indexKey,
kind: "exact",
},
{
value: indexBounds.upperBound,
kind: indexBounds.upperBoundInclusive ? "exact" : "predecessor",
},
);
// If lowerBound <= this.indexKey <= upperBound, return this.value
if (compareLowerBound <= 0 && compareUpperBound <= 0) {
return new SingletonStream(
this.#value,
this.#order,
this.#indexFields,
this.#indexKey,
this.#equalityIndexFilter,
);
}
return new EmptyStream(this.#order, this.#indexFields);
}
}
/**
* This is a completely empty stream that yields no values, and in particular
* does not count towards maximumRowsRead.
* Compare to SingletonStream(null, ...), which yields no values but does count
* towards maximumRowsRead.
*/
export class EmptyStream<T extends GenericStreamItem> extends QueryStream<T> {
#order: "asc" | "desc";
#indexFields: string[];
constructor(order: "asc" | "desc", indexFields: string[]) {
super();
this.#order = order;
this.#indexFields = indexFields;
}
iterWithKeys(): AsyncIterable<[T | null, IndexKey]> {
return {
[Symbol.asyncIterator]() {
return {
async next() {
return { done: true, value: undefined };
},
};
},
};
}
getOrder(): "asc" | "desc" {
return this.#order;
}
getIndexFields(): string[] {
return this.#indexFields;
}
getEqualityIndexFilter(): Value[] {
return [];
}
narrow(_indexBounds: IndexBounds) {
return this;
}
}
function normalizeIndexFields(indexFields: string[]) {
// Append _creationTime and _id to the index fields if they're not already there
if (!indexFields.includes("_creationTime")) {
// With one exception: if indexFields is ["_id"], we don't need to add _creationTime
if (indexFields.length !== 1 || indexFields[0] !== "_id") {
indexFields.push("_creationTime");
}
}
if (!indexFields.includes("_id")) {
indexFields.push("_id");
}
}
// Given a stream ordered by `indexFields`, where the first `equalityIndexLength`
// fields are bounded by equality filters, return a generator of the possible
// index fields used for ordering.
function* getOrderingIndexFields<T extends GenericStreamItem>(
stream: QueryStream<T>,
): Generator<string[]> {
const streamEqualityIndexLength = stream.getEqualityIndexFilter().length;
const streamIndexFields = stream.getIndexFields();
for (let i = 0; i <= streamEqualityIndexLength; i++) {
yield streamIndexFields.slice(i);
}
}
class OrderByStream<T extends GenericStreamItem> extends QueryStream<T> {
#staticFilter: Value[];
#stream: QueryStream<T>;
#indexFields: string[];
constructor(stream: QueryStream<T>, indexFields: string[]) {
super();
this.#stream = stream;
this.#indexFields = indexFields;
normalizeIndexFields(this.#indexFields);
// indexFields must be a suffix of the stream's in