diff --git a/.changeset/twelve-keys-hunt.md b/.changeset/twelve-keys-hunt.md new file mode 100644 index 0000000..7181cb0 --- /dev/null +++ b/.changeset/twelve-keys-hunt.md @@ -0,0 +1,5 @@ +--- +"monarch-orm": minor +--- + +Add searchIndexes support. Use `schema.searchIndexes()` for defining Atlas Search Indexes diff --git a/src/database.ts b/src/database.ts index ae62279..92085dd 100644 --- a/src/database.ts +++ b/src/database.ts @@ -4,6 +4,7 @@ import { Collection } from "./collection/collection"; import { applyIndexes } from "./schema/indexes"; import type { AnySchema, Schemas } from "./schema/schema"; import { Schema } from "./schema/schema"; +import { applySearchIndexes } from "./schema/search-indexes"; import { getValidator, type SchemaValidation, type Validator } from "./schema/validation"; import type { DbCollections } from "./type-helpers"; import { createAsyncLimiter, createAsyncResolver, type AsyncResolver } from "./utils/misc"; @@ -92,16 +93,22 @@ export class Database> { /** * Creates collections with indexes and document validation if provided. * - * @param options - Init options + * @param options - Init options. Pass `true` to run all steps, or an object to selectively enable steps. + * @param collections - Select collections. Omit to initialize all collections, or an object to selectively enable collections. */ - public async initialize(options?: InitOptions): Promise { + public async initialize( + options?: InitOptions | true, + collections?: InitCollections, + ): Promise { const promises: Promise[] = []; - const collections = Object.values(this.collections).map((c: Collection): CollectionInit => { - const resolver = createAsyncResolver(); - promises.push(resolver.promise); - return { schema: c.schema, defaultValidation: this.options?.validation, resolver }; - }); - initializeCollections(this.db, collections, options); + const collectionInits = (Object.values(this.collections) as Collection[]) + .filter((c) => !collections || collections[c.schema.name] === true) + .map((c): CollectionInit => { + const resolver = createAsyncResolver(); + promises.push(resolver.promise); + return { schema: c.schema, defaultValidation: this.options?.validation, resolver }; + }); + initializeCollections(this.db, collectionInits, options); return Promise.all(promises).then(() => undefined); } @@ -145,19 +152,33 @@ export function createDatabase>( return new Database(db, schemas, options); } -type InitOptions = { - indexes?: boolean; - validation?: boolean; - collections?: Partial>; +/** + * Initialization options. When provided, only fields explicitly set to `true` will run. + * When omitted, all initialization steps run. + */ +export type InitOptions = { + /** Create or update schema indexes. */ + indexes?: true; + /** Create or update search indexes. */ + searchIndexes?: true; + /** Apply document validation rules. */ + validation?: true; }; +/** + * Limit initialization to specific collections. + * When provided, only collections with a `true` value will be initialized. + */ +type InitCollections = { [K in T]?: true }; + type CollectionInit = { schema: AnySchema; defaultValidation?: SchemaValidation; resolver: AsyncResolver; }; -function initializeCollections(db: Db, collections: CollectionInit[], options?: InitOptions) { +function initializeCollections(db: Db, collections: CollectionInit[], options?: InitOptions | true) { + const opts = options === true ? undefined : options; const run = createAsyncLimiter(10); const existingPromise = db .listCollections({}, { nameOnly: true }) @@ -165,13 +186,6 @@ function initializeCollections(db: Db, collections: CollectionInit[], options?: .then((colls) => new Set(colls.map((c) => c.name))); for (const c of collections) { - // Skip disabled collections - const enabled = options?.collections ? options.collections[c.schema.name] === true : true; - if (!enabled) { - c.resolver.resolve(); - continue; - } - run(async () => { const existing = await existingPromise; const exists = existing.has(c.schema.name); @@ -180,7 +194,7 @@ function initializeCollections(db: Db, collections: CollectionInit[], options?: // Get schema validation let validation: (SchemaValidation & { validator: Validator }) | undefined; const validationOptions = schemaOptions.validation ?? c.defaultValidation; - if ((options?.validation ?? true) && validationOptions) { + if ((opts === undefined || opts.validation) && validationOptions) { validation = { ...validationOptions, validator: getValidator(c.schema) }; } @@ -194,9 +208,14 @@ function initializeCollections(db: Db, collections: CollectionInit[], options?: } // Create schema indexes - if ((options?.indexes ?? true) && schemaOptions.indexes) { + if ((opts === undefined || opts.indexes) && schemaOptions.indexes) { await applyIndexes(coll, schemaOptions.indexes); } + + // Create schema search indexes + if ((opts === undefined || opts.searchIndexes) && schemaOptions.searchIndexes) { + await applySearchIndexes(coll, schemaOptions.searchIndexes); + } }) .then(c.resolver.resolve) .catch(c.resolver.reject); diff --git a/src/index.ts b/src/index.ts index a30efce..70c7a91 100644 --- a/src/index.ts +++ b/src/index.ts @@ -37,7 +37,7 @@ export type { ReplaceOneQueryOptions } from "./collection/query/replace-one"; export type { UpdateManyQueryOptions } from "./collection/query/update-many"; export type { UpdateOneQueryOptions } from "./collection/query/update-one"; export type { PipelineStage } from "./collection/types/pipeline-stage"; -export { createClient, createDatabase, Database, type DatabaseOptions } from "./database"; +export { createClient, createDatabase, Database, type DatabaseOptions, type InitOptions } from "./database"; export { MonarchError, MonarchParseError } from "./errors"; export { mergeRelations, @@ -46,7 +46,14 @@ export { type RelationsFn, type SchemasRelations, } from "./relations/relations"; +export type { CreateIndexesOptions, SchemaIndex } from "./schema/indexes"; export { createSchema, defineSchemas, mergeSchemas, Schema, Schemas } from "./schema/schema"; +export type { + SchemaSearchIndex, + SchemaSearchIndexDefinition, + SearchIndexDefinition, + VectorSearchIndexDefinition, +} from "./schema/search-indexes"; export type { Condition, CreateIndexKey, diff --git a/src/schema/indexes.ts b/src/schema/indexes.ts index 5b52565..e8b08f5 100644 --- a/src/schema/indexes.ts +++ b/src/schema/indexes.ts @@ -8,6 +8,8 @@ import { MonarchError } from "../errors"; import type { AnyMonarchType } from "../types/type"; import type { CreateIndexKey } from "./type-helpers"; +export type { CreateIndexesOptions }; + export type CreateIndex> = ( key: CreateIndexKey, options?: CreateIndexesOptions, @@ -28,15 +30,15 @@ export type SchemaIndexes> = (options: [k: string]: SchemaIndex; }; -export function makeIndexes>(indexesFn: SchemaIndexes) { - return indexesFn({ +export function makeIndexes>(fn: SchemaIndexes) { + return fn({ createIndex: (key, options) => ({ key, options }), unique: (key) => ({ key: { [key]: 1 } as CreateIndexKey, options: { unique: true } }), }); } -export async function applyIndexes(coll: MongoCollection, indexesFn: SchemaIndexes) { - const indexes = Object.entries(makeIndexes(indexesFn)); +export async function applyIndexes(coll: MongoCollection, fn: SchemaIndexes) { + const indexes = Object.entries(makeIndexes(fn)); const desiredIndexes = new Map(indexes.map(([_, idx]) => [JSON.stringify(idx.key), idx.options ?? {}])); const existingIndexes = await coll.indexes(); const indexesToDrop: string[] = []; diff --git a/src/schema/schema.ts b/src/schema/schema.ts index f1aaaba..37b9d7c 100644 --- a/src/schema/schema.ts +++ b/src/schema/schema.ts @@ -7,6 +7,7 @@ import { MonarchObjectId, objectId } from "../types/objectId"; import { MonarchNullable, MonarchOptional, MonarchType, type AnyMonarchType } from "../types/type"; import type { MergeAll, MergeN1All, Pretty, RequiredObject } from "../utils/type-helpers"; import type { SchemaIndexes } from "./indexes"; +import type { SchemaSearchIndexes } from "./search-indexes"; import type { InferSchemaData, InferSchemaInput, @@ -52,6 +53,7 @@ export class Schema< private options: { omit?: SchemaOmit; indexes?: SchemaIndexes; + searchIndexes?: SchemaSearchIndexes; validation?: SchemaValidation; virtuals?: SchemaVirtuals; renames?: TRenames; @@ -104,7 +106,6 @@ export class Schema< * This method allows you to specify indexes that should be created for the schema. * * @param indexes - A function that defines the indexes to be created. - * * @returns The current schema instance for method chaining. * * @example @@ -121,6 +122,33 @@ export class Schema< return this; } + /** + * Defines the search indexes for the schema. + * + * Search indexes are only supported on MongoDB Atlas clusters and are applied during initialization. + * + * @param searchIndexes - A function that defines the search indexes to be created. + * @returns The current schema instance for method chaining. + * + * @example + * const articleSchema = createSchema("articles", { + * title: string(), + * body: string(), + * embedding: array(number()), + * }).searchIndexes(({ searchIndex, vectorSearchIndex }) => ({ + * fullText: searchIndex("articles_search", { + * mappings: { dynamic: false, fields: { title: { type: "string" }, body: { type: "string" } } }, + * }), + * semantic: vectorSearchIndex("articles_vector", { + * fields: [{ type: "vector", path: "embedding", numDimensions: 1536, similarity: "cosine" }], + * }), + * })); + */ + public searchIndexes(searchIndexes: SchemaSearchIndexes) { + this.options.searchIndexes = searchIndexes; + return this; + } + /** * Sets MongoDB document validation for this schema. * diff --git a/src/schema/search-indexes.ts b/src/schema/search-indexes.ts new file mode 100644 index 0000000..776ec37 --- /dev/null +++ b/src/schema/search-indexes.ts @@ -0,0 +1,94 @@ +import type { Document, Collection as MongoCollection } from "mongodb"; +import { MonarchError } from "../errors"; +import type { AnyMonarchType } from "../types/type"; +import type { CreateIndexKey } from "./type-helpers"; + +export type SearchIndexDefinition> = { + mappings?: { + dynamic?: boolean; + fields?: { [K in keyof CreateIndexKey]?: Document }; + }; +}; +export type VectorSearchIndexDefinition> = { + fields?: Array<{ + type: "vector" | "filter"; + path: keyof CreateIndexKey; + numDimensions?: number; + similarity?: "euclidean" | "cosine" | "dotProduct"; + [key: string]: unknown; + }>; +}; +export type SchemaSearchIndexDefinition> = + | SearchIndexDefinition + | VectorSearchIndexDefinition; + +export type SchemaSearchIndex> = { + name: string; + type?: "search" | "vectorSearch"; + definition: SchemaSearchIndexDefinition; +}; + +export type SchemaSearchIndexes> = (options: { + searchIndex: (name: string, definition: SearchIndexDefinition) => SchemaSearchIndex; + vectorSearchIndex: (name: string, definition: VectorSearchIndexDefinition) => SchemaSearchIndex; +}) => Record>; + +export function makeSearchIndexes>(fn: SchemaSearchIndexes) { + return fn({ + searchIndex: (name, definition) => ({ name, type: "search", definition }), + vectorSearchIndex: (name, definition) => ({ name, type: "vectorSearch", definition }), + }); +} + +type ExistingSearchIndex = { + id: string; + name: string; + status: string; + queryable: boolean; + latestDefinition: Document; +}; + +export async function applySearchIndexes(coll: MongoCollection, fn: SchemaSearchIndexes) { + const desired = Object.values(makeSearchIndexes(fn)); + const desiredByName = new Map(desired.map((idx) => [idx.name, idx])); + + let existing: ExistingSearchIndex[]; + try { + existing = (await coll.listSearchIndexes().toArray()) as ExistingSearchIndex[]; + } catch { + return; + } + + const existingByName = new Map(existing.map((idx) => [idx.name, idx])); + + // Drop stale indexes + await Promise.all( + Array.from(existingByName.keys()) + .filter((name) => !desiredByName.has(name)) + .map((name) => + coll.dropSearchIndex(name).catch((error) => { + throw new MonarchError(`failed to drop search index '${name}': ${error}`, error); + }), + ), + ); + + // Create or update indexes + await Promise.all( + desired.map(async (idx) => { + const existing = existingByName.get(idx.name); + if (!existing) { + await coll.createSearchIndex({ name: idx.name, type: idx.type, definition: idx.definition }).catch((error) => { + throw new MonarchError(`failed to create search index '${idx.name}': ${error}`, error); + }); + return; + } + + const defChanged = JSON.stringify(existing.latestDefinition) !== JSON.stringify(idx.definition); + if (defChanged) { + await coll.updateSearchIndex(idx.name, idx.definition).catch((error) => { + throw new MonarchError(`failed to update search index '${idx.name}': ${error}`, error); + }); + } + }), + ); +} diff --git a/tests/database.test.ts b/tests/database.test.ts index 501f4db..3b5a611 100644 --- a/tests/database.test.ts +++ b/tests/database.test.ts @@ -73,7 +73,7 @@ describe("Database options", async () => { ).rejects.toThrow("Document failed validation"); }); - it("initialize validation false skips applying validators", async () => { + it("initialize without validation option skips applying validators", async () => { const schema = createSchema("users", { name: string(), nickname: string(), @@ -86,13 +86,13 @@ describe("Database options", async () => { validationAction: "error", }, }); - await db.initialize({ validation: false }); + await db.initialize({}); const rawCollection = client.db().collection("users"); await expect(rawCollection.insertOne({})).resolves.toMatchObject({ acknowledged: true }); }); - it("initialize indexes false skips schema index creation", async () => { + it("initialize without indexes option skips schema index creation", async () => { const schema = createSchema("users", { username: string(), }).indexes(({ unique }) => ({ @@ -100,7 +100,7 @@ describe("Database options", async () => { })); const db = createDatabase(client.db(), defineSchemas({ users: schema }), { initialize: false }); - await db.initialize({ indexes: false }); + await db.initialize({}); const rawCollection = client.db().collection("users"); await rawCollection.insertOne({ username: "same-user" }); @@ -123,7 +123,7 @@ describe("Database options", async () => { validationAction: "error", }, }); - await db.initialize({ collections: { users: true } }); + await db.initialize(true, { users: true }); const existing = await client.db().listCollections({}, { nameOnly: true }).toArray(); const existingNames = new Set(existing.map((collection) => collection.name)); @@ -164,7 +164,7 @@ describe("Database options", async () => { }); const firstDb = createDatabase(client.db(), defineSchemas({ users: schema }), { initialize: false }); - await firstDb.initialize({ validation: false }); + await firstDb.initialize({}); const rawCollection = client.db().collection("users"); await expect(rawCollection.insertOne({})).resolves.toMatchObject({ acknowledged: true });