join on fields between convex tables

join.ts

import { DataModel, Doc, Id, TableNames } from "./_generated/dataModel";

import { GenericDatabaseReader, GenericDatabaseWriter } from "convex/server";

type Database =
  | GenericDatabaseReader<DataModel>
  | GenericDatabaseWriter<DataModel>;

type DocWithId<
  Table extends TableNames,
  Relation extends TableNames,
> = Doc<Table> & { [key in `${Relation}Id`]: Id<Relation> };

type Joined<
  Table extends TableNames,
  Relation extends TableNames,
> = Doc<Table> & {
  [key in Relation]: Doc<Relation> | undefined;
};

function nonNullable<T>(a: T): a is NonNullable<T> {
  return a != null;
}

/**
 * Joins two tables on the given field, example:
 *
 *   // This has { _id: Id<"orgRole">, orgId: Id<"org">, ... }
 *   const roles = await ctx.db.query("orgRole").collect();
 *
 *   // This has { _id: Id<"orgRole">, org: Doc<"org"> | undefined, ... }
 *   const data = await join<"orgRole", "org">(ctx.db, roles, "orgId");
 */
export async function join<
  BaseTable extends TableNames = never,
  JoinTable extends TableNames = never,
>(
  db: Database,
  items: DocWithId<BaseTable, JoinTable>[],
  field: `${JoinTable}Id`,
): Promise<Joined<BaseTable, JoinTable>[]> {
  const ids = new Set(items.map((item) => item[field]));
  const docs = await Promise.all(Array.from(ids).map((id) => db.get(id)));
  const docMap = new Map(docs.filter(nonNullable).map((doc) => [doc._id, doc]));
  const targetField = field.slice(0, -2) as JoinTable;
  const result = items.map((item) => {
    return {
      ...item,
      [targetField]: docMap.get(item[field]),
    };
  });
  return result;
}