480 lines
13 KiB
TypeScript
480 lines
13 KiB
TypeScript
import * as plugins from '../tsmdb.plugins.js';
|
|
import type { IStorageAdapter } from '../storage/IStorageAdapter.js';
|
|
import type {
|
|
Document,
|
|
IStoredDocument,
|
|
IIndexSpecification,
|
|
IIndexInfo,
|
|
ICreateIndexOptions,
|
|
} from '../types/interfaces.js';
|
|
import { TsmdbDuplicateKeyError, TsmdbIndexError } from '../errors/TsmdbErrors.js';
|
|
import { QueryEngine } from './QueryEngine.js';
|
|
|
|
/**
|
|
* Index data structure for fast lookups
|
|
*/
|
|
interface IIndexData {
|
|
name: string;
|
|
key: Record<string, 1 | -1 | string>;
|
|
unique: boolean;
|
|
sparse: boolean;
|
|
expireAfterSeconds?: number;
|
|
// Map from index key value to document _id(s)
|
|
entries: Map<string, Set<string>>;
|
|
}
|
|
|
|
/**
|
|
* Index engine for managing indexes and query optimization
|
|
*/
|
|
export class IndexEngine {
|
|
private dbName: string;
|
|
private collName: string;
|
|
private storage: IStorageAdapter;
|
|
private indexes: Map<string, IIndexData> = new Map();
|
|
private initialized = false;
|
|
|
|
constructor(dbName: string, collName: string, storage: IStorageAdapter) {
|
|
this.dbName = dbName;
|
|
this.collName = collName;
|
|
this.storage = storage;
|
|
}
|
|
|
|
/**
|
|
* Initialize indexes from storage
|
|
*/
|
|
async initialize(): Promise<void> {
|
|
if (this.initialized) return;
|
|
|
|
const storedIndexes = await this.storage.getIndexes(this.dbName, this.collName);
|
|
const documents = await this.storage.findAll(this.dbName, this.collName);
|
|
|
|
for (const indexSpec of storedIndexes) {
|
|
const indexData: IIndexData = {
|
|
name: indexSpec.name,
|
|
key: indexSpec.key,
|
|
unique: indexSpec.unique || false,
|
|
sparse: indexSpec.sparse || false,
|
|
expireAfterSeconds: indexSpec.expireAfterSeconds,
|
|
entries: new Map(),
|
|
};
|
|
|
|
// Build index entries
|
|
for (const doc of documents) {
|
|
const keyValue = this.extractKeyValue(doc, indexSpec.key);
|
|
if (keyValue !== null || !indexData.sparse) {
|
|
const keyStr = JSON.stringify(keyValue);
|
|
if (!indexData.entries.has(keyStr)) {
|
|
indexData.entries.set(keyStr, new Set());
|
|
}
|
|
indexData.entries.get(keyStr)!.add(doc._id.toHexString());
|
|
}
|
|
}
|
|
|
|
this.indexes.set(indexSpec.name, indexData);
|
|
}
|
|
|
|
this.initialized = true;
|
|
}
|
|
|
|
/**
|
|
* Create a new index
|
|
*/
|
|
async createIndex(
|
|
key: Record<string, 1 | -1 | 'text' | '2dsphere'>,
|
|
options?: ICreateIndexOptions
|
|
): Promise<string> {
|
|
await this.initialize();
|
|
|
|
// Generate index name if not provided
|
|
const name = options?.name || this.generateIndexName(key);
|
|
|
|
// Check if index already exists
|
|
if (this.indexes.has(name)) {
|
|
return name;
|
|
}
|
|
|
|
// Create index data structure
|
|
const indexData: IIndexData = {
|
|
name,
|
|
key: key as Record<string, 1 | -1 | string>,
|
|
unique: options?.unique || false,
|
|
sparse: options?.sparse || false,
|
|
expireAfterSeconds: options?.expireAfterSeconds,
|
|
entries: new Map(),
|
|
};
|
|
|
|
// Build index from existing documents
|
|
const documents = await this.storage.findAll(this.dbName, this.collName);
|
|
|
|
for (const doc of documents) {
|
|
const keyValue = this.extractKeyValue(doc, key);
|
|
|
|
if (keyValue === null && indexData.sparse) {
|
|
continue;
|
|
}
|
|
|
|
const keyStr = JSON.stringify(keyValue);
|
|
|
|
if (indexData.unique && indexData.entries.has(keyStr)) {
|
|
throw new TsmdbDuplicateKeyError(
|
|
`E11000 duplicate key error index: ${this.dbName}.${this.collName}.$${name}`,
|
|
key as Record<string, 1>,
|
|
keyValue
|
|
);
|
|
}
|
|
|
|
if (!indexData.entries.has(keyStr)) {
|
|
indexData.entries.set(keyStr, new Set());
|
|
}
|
|
indexData.entries.get(keyStr)!.add(doc._id.toHexString());
|
|
}
|
|
|
|
// Store index
|
|
this.indexes.set(name, indexData);
|
|
await this.storage.saveIndex(this.dbName, this.collName, name, {
|
|
key,
|
|
unique: options?.unique,
|
|
sparse: options?.sparse,
|
|
expireAfterSeconds: options?.expireAfterSeconds,
|
|
});
|
|
|
|
return name;
|
|
}
|
|
|
|
/**
|
|
* Drop an index
|
|
*/
|
|
async dropIndex(name: string): Promise<void> {
|
|
await this.initialize();
|
|
|
|
if (name === '_id_') {
|
|
throw new TsmdbIndexError('cannot drop _id index');
|
|
}
|
|
|
|
if (!this.indexes.has(name)) {
|
|
throw new TsmdbIndexError(`index not found: ${name}`);
|
|
}
|
|
|
|
this.indexes.delete(name);
|
|
await this.storage.dropIndex(this.dbName, this.collName, name);
|
|
}
|
|
|
|
/**
|
|
* Drop all indexes except _id
|
|
*/
|
|
async dropAllIndexes(): Promise<void> {
|
|
await this.initialize();
|
|
|
|
const names = Array.from(this.indexes.keys()).filter(n => n !== '_id_');
|
|
for (const name of names) {
|
|
this.indexes.delete(name);
|
|
await this.storage.dropIndex(this.dbName, this.collName, name);
|
|
}
|
|
}
|
|
|
|
/**
|
|
* List all indexes
|
|
*/
|
|
async listIndexes(): Promise<IIndexInfo[]> {
|
|
await this.initialize();
|
|
|
|
return Array.from(this.indexes.values()).map(idx => ({
|
|
v: 2,
|
|
key: idx.key,
|
|
name: idx.name,
|
|
unique: idx.unique || undefined,
|
|
sparse: idx.sparse || undefined,
|
|
expireAfterSeconds: idx.expireAfterSeconds,
|
|
}));
|
|
}
|
|
|
|
/**
|
|
* Check if an index exists
|
|
*/
|
|
async indexExists(name: string): Promise<boolean> {
|
|
await this.initialize();
|
|
return this.indexes.has(name);
|
|
}
|
|
|
|
/**
|
|
* Update index entries after document insert
|
|
*/
|
|
async onInsert(doc: IStoredDocument): Promise<void> {
|
|
await this.initialize();
|
|
|
|
for (const [name, indexData] of this.indexes) {
|
|
const keyValue = this.extractKeyValue(doc, indexData.key);
|
|
|
|
if (keyValue === null && indexData.sparse) {
|
|
continue;
|
|
}
|
|
|
|
const keyStr = JSON.stringify(keyValue);
|
|
|
|
// Check unique constraint
|
|
if (indexData.unique) {
|
|
const existing = indexData.entries.get(keyStr);
|
|
if (existing && existing.size > 0) {
|
|
throw new TsmdbDuplicateKeyError(
|
|
`E11000 duplicate key error collection: ${this.dbName}.${this.collName} index: ${name}`,
|
|
indexData.key as Record<string, 1>,
|
|
keyValue
|
|
);
|
|
}
|
|
}
|
|
|
|
if (!indexData.entries.has(keyStr)) {
|
|
indexData.entries.set(keyStr, new Set());
|
|
}
|
|
indexData.entries.get(keyStr)!.add(doc._id.toHexString());
|
|
}
|
|
}
|
|
|
|
/**
|
|
* Update index entries after document update
|
|
*/
|
|
async onUpdate(oldDoc: IStoredDocument, newDoc: IStoredDocument): Promise<void> {
|
|
await this.initialize();
|
|
|
|
for (const [name, indexData] of this.indexes) {
|
|
const oldKeyValue = this.extractKeyValue(oldDoc, indexData.key);
|
|
const newKeyValue = this.extractKeyValue(newDoc, indexData.key);
|
|
const oldKeyStr = JSON.stringify(oldKeyValue);
|
|
const newKeyStr = JSON.stringify(newKeyValue);
|
|
|
|
// Remove old entry if key changed
|
|
if (oldKeyStr !== newKeyStr) {
|
|
if (oldKeyValue !== null || !indexData.sparse) {
|
|
const oldSet = indexData.entries.get(oldKeyStr);
|
|
if (oldSet) {
|
|
oldSet.delete(oldDoc._id.toHexString());
|
|
if (oldSet.size === 0) {
|
|
indexData.entries.delete(oldKeyStr);
|
|
}
|
|
}
|
|
}
|
|
|
|
// Add new entry
|
|
if (newKeyValue !== null || !indexData.sparse) {
|
|
// Check unique constraint
|
|
if (indexData.unique) {
|
|
const existing = indexData.entries.get(newKeyStr);
|
|
if (existing && existing.size > 0) {
|
|
throw new TsmdbDuplicateKeyError(
|
|
`E11000 duplicate key error collection: ${this.dbName}.${this.collName} index: ${name}`,
|
|
indexData.key as Record<string, 1>,
|
|
newKeyValue
|
|
);
|
|
}
|
|
}
|
|
|
|
if (!indexData.entries.has(newKeyStr)) {
|
|
indexData.entries.set(newKeyStr, new Set());
|
|
}
|
|
indexData.entries.get(newKeyStr)!.add(newDoc._id.toHexString());
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
/**
|
|
* Update index entries after document delete
|
|
*/
|
|
async onDelete(doc: IStoredDocument): Promise<void> {
|
|
await this.initialize();
|
|
|
|
for (const indexData of this.indexes.values()) {
|
|
const keyValue = this.extractKeyValue(doc, indexData.key);
|
|
|
|
if (keyValue === null && indexData.sparse) {
|
|
continue;
|
|
}
|
|
|
|
const keyStr = JSON.stringify(keyValue);
|
|
const set = indexData.entries.get(keyStr);
|
|
if (set) {
|
|
set.delete(doc._id.toHexString());
|
|
if (set.size === 0) {
|
|
indexData.entries.delete(keyStr);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
/**
|
|
* Find the best index for a query
|
|
*/
|
|
selectIndex(filter: Document): { name: string; data: IIndexData } | null {
|
|
if (!filter || Object.keys(filter).length === 0) {
|
|
return null;
|
|
}
|
|
|
|
// Get filter fields
|
|
const filterFields = new Set(this.getFilterFields(filter));
|
|
|
|
// Score each index
|
|
let bestIndex: { name: string; data: IIndexData } | null = null;
|
|
let bestScore = 0;
|
|
|
|
for (const [name, indexData] of this.indexes) {
|
|
const indexFields = Object.keys(indexData.key);
|
|
let score = 0;
|
|
|
|
// Count how many index fields are in the filter
|
|
for (const field of indexFields) {
|
|
if (filterFields.has(field)) {
|
|
score++;
|
|
} else {
|
|
break; // Index fields must be contiguous
|
|
}
|
|
}
|
|
|
|
// Prefer unique indexes
|
|
if (indexData.unique && score > 0) {
|
|
score += 0.5;
|
|
}
|
|
|
|
if (score > bestScore) {
|
|
bestScore = score;
|
|
bestIndex = { name, data: indexData };
|
|
}
|
|
}
|
|
|
|
return bestIndex;
|
|
}
|
|
|
|
/**
|
|
* Use index to find candidate document IDs
|
|
*/
|
|
async findCandidateIds(filter: Document): Promise<Set<string> | null> {
|
|
await this.initialize();
|
|
|
|
const index = this.selectIndex(filter);
|
|
if (!index) return null;
|
|
|
|
// Try to use the index for equality matches
|
|
const indexFields = Object.keys(index.data.key);
|
|
const equalityValues: Record<string, any> = {};
|
|
|
|
for (const field of indexFields) {
|
|
const filterValue = this.getFilterValue(filter, field);
|
|
if (filterValue === undefined) break;
|
|
|
|
// Only use equality matches for index lookup
|
|
if (typeof filterValue === 'object' && filterValue !== null) {
|
|
if (filterValue.$eq !== undefined) {
|
|
equalityValues[field] = filterValue.$eq;
|
|
} else if (filterValue.$in !== undefined) {
|
|
// Handle $in with multiple lookups
|
|
const results = new Set<string>();
|
|
for (const val of filterValue.$in) {
|
|
equalityValues[field] = val;
|
|
const keyStr = JSON.stringify(this.buildKeyValue(equalityValues, index.data.key));
|
|
const ids = index.data.entries.get(keyStr);
|
|
if (ids) {
|
|
for (const id of ids) {
|
|
results.add(id);
|
|
}
|
|
}
|
|
}
|
|
return results;
|
|
} else {
|
|
break; // Non-equality operator, stop here
|
|
}
|
|
} else {
|
|
equalityValues[field] = filterValue;
|
|
}
|
|
}
|
|
|
|
if (Object.keys(equalityValues).length === 0) {
|
|
return null;
|
|
}
|
|
|
|
const keyStr = JSON.stringify(this.buildKeyValue(equalityValues, index.data.key));
|
|
return index.data.entries.get(keyStr) || new Set();
|
|
}
|
|
|
|
// ============================================================================
|
|
// Helper Methods
|
|
// ============================================================================
|
|
|
|
private generateIndexName(key: Record<string, any>): string {
|
|
return Object.entries(key)
|
|
.map(([field, dir]) => `${field}_${dir}`)
|
|
.join('_');
|
|
}
|
|
|
|
private extractKeyValue(doc: Document, key: Record<string, any>): any {
|
|
const values: any[] = [];
|
|
|
|
for (const field of Object.keys(key)) {
|
|
const value = QueryEngine.getNestedValue(doc, field);
|
|
values.push(value === undefined ? null : value);
|
|
}
|
|
|
|
// For single-field index, return the value directly
|
|
if (values.length === 1) {
|
|
return values[0];
|
|
}
|
|
|
|
return values;
|
|
}
|
|
|
|
private buildKeyValue(values: Record<string, any>, key: Record<string, any>): any {
|
|
const result: any[] = [];
|
|
|
|
for (const field of Object.keys(key)) {
|
|
result.push(values[field] !== undefined ? values[field] : null);
|
|
}
|
|
|
|
if (result.length === 1) {
|
|
return result[0];
|
|
}
|
|
|
|
return result;
|
|
}
|
|
|
|
private getFilterFields(filter: Document, prefix = ''): string[] {
|
|
const fields: string[] = [];
|
|
|
|
for (const [key, value] of Object.entries(filter)) {
|
|
if (key.startsWith('$')) {
|
|
// Logical operator
|
|
if (key === '$and' || key === '$or' || key === '$nor') {
|
|
for (const subFilter of value as Document[]) {
|
|
fields.push(...this.getFilterFields(subFilter, prefix));
|
|
}
|
|
}
|
|
} else {
|
|
const fullKey = prefix ? `${prefix}.${key}` : key;
|
|
fields.push(fullKey);
|
|
|
|
// Check for nested filters
|
|
if (typeof value === 'object' && value !== null && !Array.isArray(value)) {
|
|
const subKeys = Object.keys(value);
|
|
if (subKeys.length > 0 && !subKeys[0].startsWith('$')) {
|
|
fields.push(...this.getFilterFields(value, fullKey));
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
return fields;
|
|
}
|
|
|
|
private getFilterValue(filter: Document, field: string): any {
|
|
// Handle dot notation
|
|
const parts = field.split('.');
|
|
let current: any = filter;
|
|
|
|
for (const part of parts) {
|
|
if (current === null || current === undefined) {
|
|
return undefined;
|
|
}
|
|
current = current[part];
|
|
}
|
|
|
|
return current;
|
|
}
|
|
}
|