Files
db/src/index/manager.ts
2026-01-07 02:34:18 +01:00

401 lines
11 KiB
TypeScript

import { Query } from "mingo";
import type { Cursor } from "mingo/cursor";
import type { Criteria } from "mingo/types";
import type { AnyDocument, QueryCriteria, StringKeyOf } from "../types.ts";
import { PrimaryIndex, type PrimaryKey } from "./primary.ts";
import { SharedIndex } from "./shared.ts";
import { UniqueIndex } from "./unique.ts";
const OBJECT_PROTOTYPE = Object.getPrototypeOf({});
const OBJECT_TAG = "[object Object]";
const EMPTY_SET: ReadonlySet<PrimaryKey> = Object.freeze(new Set<PrimaryKey>());
export class IndexManager<TSchema extends AnyDocument> {
readonly primary: PrimaryIndex<TSchema>;
readonly unique: Map<StringKeyOf<TSchema>, UniqueIndex> = new Map<StringKeyOf<TSchema>, UniqueIndex>();
readonly shared: Map<StringKeyOf<TSchema>, SharedIndex> = new Map<StringKeyOf<TSchema>, SharedIndex>();
readonly specs: IndexSpec<TSchema>[];
constructor(specs: IndexSpec<TSchema>[]) {
const primary = specs.find((spec) => spec.kind === "primary");
if (primary === undefined) {
throw new Error("Primary index is required");
}
this.primary = new PrimaryIndex(primary.field);
for (const spec of specs) {
switch (spec.kind) {
case "unique": {
this.unique.set(spec.field, new UniqueIndex());
break;
}
case "shared": {
this.shared.set(spec.field, new SharedIndex());
break;
}
}
}
this.specs = specs;
}
#isPrimaryIndex(key: string): boolean {
for (const { field, kind } of this.specs) {
if (key === field && kind === "primary") {
return true;
}
}
return false;
}
#isUniqueIndex(key: string): boolean {
for (const { field, kind } of this.specs) {
if (key === field && kind === "unique") {
return true;
}
}
return false;
}
#isSharedIndex(key: string): boolean {
for (const { field, kind } of this.specs) {
if (key === field && kind === "shared") {
return true;
}
}
return false;
}
#getOptimalIndex(keys: string[]): string {
let best: string | undefined;
for (const key of keys) {
if (this.#isPrimaryIndex(key)) {
return key; // cannot beat primary
}
if (this.#isUniqueIndex(key)) {
best ??= key;
continue;
}
if (best === undefined && this.#isSharedIndex(key)) {
best = key;
}
}
return best ?? keys[0];
}
/**
* Atomic insert of the document into the index pools. If any part
* of the operation fails all changes are rolled back to their original
* states.
*
* @param document - Document to insert.
*/
insert(document: TSchema): void {
const pk = document[this.primary.key];
const insertedUniques: [StringKeyOf<TSchema>, any][] = [];
const insertedShared: [StringKeyOf<TSchema>, any][] = [];
try {
for (const [field, index] of this.unique) {
const value = document[field] as any;
if (value !== undefined) {
if (Array.isArray(value)) {
for (const innerValue of value) {
index.insert(innerValue, pk);
insertedUniques.push([field, innerValue]);
}
} else {
index.insert(value, pk);
insertedUniques.push([field, value]);
}
}
}
for (const [field, index] of this.shared) {
const value = document[field] as any;
if (value !== undefined) {
if (Array.isArray(value)) {
for (const innerValue of value) {
index.insert(innerValue, pk);
insertedShared.push([field, innerValue]);
}
} else {
index.insert(value, pk);
insertedShared.push([field, value]);
}
}
}
this.primary.insert(pk, document);
} catch (err) {
for (const [field, value] of insertedUniques) {
this.unique.get(field)?.delete(value);
}
for (const [field, value] of insertedShared) {
this.shared.get(field)?.delete(value, pk);
}
throw err;
}
}
getByCondition(condition: Criteria<TSchema> = {}): Cursor<TSchema> {
const indexes = resolveIndexesFromCondition(condition, this.specs);
if (indexes === undefined) {
return new Query(condition, {}).find<TSchema>(this.primary.documents);
}
const index = this.#getOptimalIndex(Object.keys(indexes));
const value = indexes[index];
if (Array.isArray(value)) {
const results: TSchema[] = [];
for (const innerValue of value) {
const records = this.getByIndex(index as any, innerValue);
results.push(...records);
}
const unique = new Map<any, TSchema>();
for (const doc of results) {
unique.set(doc[this.primary.key], doc);
}
return new Query(condition, {}).find<TSchema>(Array.from(unique.values()));
}
return new Query(condition, {}).find<TSchema>(this.getByIndex(index as any, value));
}
/**
* Get all primary keys found for given field => value pair.
*
* @param field - Field to lookup.
* @param value - Value to lookup.
*/
getPrimaryKeysByIndex(field: StringKeyOf<TSchema>, value: any): ReadonlySet<PrimaryKey> {
if (field === this.primary.key) {
if (this.primary.has(value)) {
return new Set([value]);
}
return EMPTY_SET;
}
if (this.unique.has(field)) {
const pk = this.unique.get(field)?.lookup(value);
if (pk === undefined) {
return EMPTY_SET;
}
return new Set([pk]);
}
return this.shared.get(field)?.lookup(value) ?? EMPTY_SET;
}
/**
* Get document by primary key.
*
* @param pk - Primary key to fetch document for.
*/
getByPrimary(pk: string): TSchema | undefined {
return this.primary.get(pk);
}
/**
* Get a document found for given field => value pair.
*
* @param field - Field to lookup.
* @param value - Value to lookup.
*/
getByUnique(field: StringKeyOf<TSchema>, value: any): TSchema | undefined {
const pk = this.unique.get(field)?.lookup(value);
if (pk !== undefined) {
return this.primary.get(pk);
}
}
/**
* Get all documents found for given field => value pair.
*
* @note This method may clean up stale index entries during reads.
*
* @param field - Field to lookup.
* @param value - Value to lookup.
*/
getByIndex(field: StringKeyOf<TSchema>, value: any): TSchema[] {
if (field === this.primary.key) {
const document = this.getByPrimary(value);
if (document === undefined) {
return [];
}
return [document];
}
if (this.unique.has(field)) {
const document = this.getByUnique(field, value);
if (document === undefined) {
this.unique.get(field)?.delete(value);
return [];
}
return [document];
}
const pks = this.shared.get(field)?.lookup(value);
if (pks === undefined) {
return [];
}
const documents: TSchema[] = [];
for (const pk of pks) {
const document = this.primary.get(pk);
if (document === undefined) {
this.shared.get(field)?.delete(value, pk);
} else {
documents.push(document);
}
}
return documents;
}
/**
* Update indexes for given document.
*
* @note If the document does not exist it will be inserted.
*
* @param document - Document to update against current index.
*/
update(document: TSchema): void {
const pk = document[this.primary.key];
const current = this.primary.get(pk);
if (current === undefined) {
this.insert(document);
return;
}
const revertedUniques: [StringKeyOf<TSchema>, any][] = [];
const revertedShared: [StringKeyOf<TSchema>, any][] = [];
try {
for (const [field, index] of this.unique) {
if (current[field] !== document[field]) {
index.delete(current[field]);
index.insert(document[field], pk);
revertedUniques.push([field, current[field]]);
}
}
for (const [field, index] of this.shared) {
if (current[field] !== document[field]) {
index.delete(current[field], pk);
index.insert(document[field], pk);
revertedShared.push([field, current[field]]);
}
}
this.primary.replace(pk, document);
} catch (err) {
for (const [field, value] of revertedUniques) {
this.unique.get(field)?.insert(value, pk);
this.unique.get(field)?.delete(document[field]);
}
for (const [field, value] of revertedShared) {
this.shared.get(field)?.insert(value, pk);
this.shared.get(field)?.delete(document[field], pk);
}
throw err;
}
}
/**
* Remove all indexes related to given document.
*
* @param document - Document to remove.
*/
remove(document: TSchema) {
const pk = document[this.primary.key];
const current = this.primary.get(pk);
if (current === undefined) {
return;
}
for (const [field, index] of this.unique) {
index.delete(current[field]);
}
for (const [field, index] of this.shared) {
index.delete(current[field], pk);
}
this.primary.delete(pk);
}
flush() {
this.primary.flush();
this.unique.clear();
this.shared.clear();
for (const spec of this.specs) {
switch (spec.kind) {
case "unique": {
this.unique.set(spec.field, new UniqueIndex());
break;
}
case "shared": {
this.shared.set(spec.field, new SharedIndex());
break;
}
}
}
}
}
/*
|--------------------------------------------------------------------------------
| Utils
|--------------------------------------------------------------------------------
*/
function resolveIndexesFromCondition<TSchema extends AnyDocument>(
condition: QueryCriteria<TSchema>,
indexes: IndexSpec<TSchema>[],
): Record<StringKeyOf<TSchema>, any> | undefined {
const indexNames = indexes.map(({ field }) => field);
const index: any = {};
for (const key in condition) {
if (indexNames.includes(key as any) === true) {
let val: any;
if (isObject(condition[key]) === true) {
if ((condition as any)[key].$in !== undefined) {
val = (condition as any)[key].$in;
}
} else {
val = condition[key];
}
if (val !== undefined) {
index[key] = val;
}
}
}
if (Object.keys(index).length > 0) {
return index;
}
}
function isObject(v: any): v is object {
if (!v) {
return false;
}
const proto = Object.getPrototypeOf(v);
return (proto === OBJECT_PROTOTYPE || proto === null) && OBJECT_TAG === Object.prototype.toString.call(v);
}
/*
|--------------------------------------------------------------------------------
| Types
|--------------------------------------------------------------------------------
*/
export type IndexSpec<TSchema extends AnyDocument> = {
field: StringKeyOf<TSchema>;
kind: IndexKind;
};
type IndexKind = "primary" | "unique" | "shared";