@elysiumjs/mnemosyne-drizzle 0.1.3 → 0.2.0

package/src/drizzle.module.ts

@@ -0,0 +1,97 @@
+// Copyright (c) 2026-present Workbud Technologies Inc. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+import type { TenancyConfig } from '@elysiumjs/mnemosyne';
+import type { DrizzleConnection, DrizzleConnectionProps } from './database';
+
+import { Module, needsDefaultAlias, runAll, Service } from '@elysiumjs/core';
+import {
+	AbstractDatabase,
+	MnemosyneModule,
+	TENANCY_CONFIG_SERVICE,
+	TENANCY_STRATEGY_SERVICE
+} from '@elysiumjs/mnemosyne';
+
+import { MigrationGenerateCommand } from './commands/migration_generate.command';
+import { MigrationRunCommand } from './commands/migration_run.command';
+import { Database } from './database';
+import { DrizzleRLSTenancy, DrizzleSchemaTenancy } from './tenancy';
+
+/**
+ * Configuration of the {@link DrizzleModule}.
+ * @author Axel Nana <axel.nana@workbud.com>
+ */
+export type DrizzleModuleConfig = {
+	/** The name of the connection that `default` resolves to. */
+	default: string;
+
+	/** The connections to declare, by name. */
+	connections: Record<string, DrizzleConnectionProps>;
+
+	/** The tenancy configuration. Defaults to schema-based tenancy. */
+	tenancy?: TenancyConfig;
+};
+
+/**
+ * Declares Drizzle connections as global services named `db.connection.{name}`, plus
+ * `db.connection.default` aliasing the configured default, and registers the tenancy config
+ * and strategy. Connections are created on first use. Also provides the migration commands.
+ * @author Axel Nana <axel.nana@workbud.com>
+ */
+@Module.register<DrizzleModuleConfig>({
+	imports: [MnemosyneModule],
+	commands: [MigrationGenerateCommand, MigrationRunCommand],
+	services: (config) => {
+		const alias = needsDefaultAlias('DrizzleModule', config);
+		const tenancy = config.tenancy ?? { mode: 'schema' };
+		return [
+			...Object.entries(config.connections).map(([name, props]) => ({
+				name: AbstractDatabase.getConnectionName(name),
+				global: true,
+				factory: () => Database.createConnection(props)
+			})),
+			...(alias
+				? [
+						{
+							name: AbstractDatabase.getConnectionName('default'),
+							global: true,
+							factory: (_config: unknown, get: <T>(name: string) => T) =>
+								get<DrizzleConnection>(AbstractDatabase.getConnectionName(config.default))
+						}
+					]
+				: []),
+			{ name: TENANCY_CONFIG_SERVICE, global: true, factory: () => tenancy },
+			{
+				name: TENANCY_STRATEGY_SERVICE,
+				global: true,
+				factory: () =>
+					tenancy.mode === 'rls' ? new DrizzleRLSTenancy() : new DrizzleSchemaTenancy()
+			}
+		];
+	}
+})
+export class DrizzleModule extends Module.configurable<DrizzleModuleConfig>({
+	defaults: { default: 'default' },
+	required: ['connections']
+}) {
+	/** Closes the clients of the connections that were created. */
+	public override async onDestroy(): Promise<void> {
+		const names = Object.keys(this.config.connections).map(AbstractDatabase.getConnectionName);
+		await runAll(
+			names
+				.filter((name) => Service.isConstructed(name))
+				.map((name) => () => Service.get<DrizzleConnection>(name)?.$client.close())
+		);
+	}
+}

package/src/index.ts

@@ -12,17 +12,15 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-import type { MnemosyneConfig } from '@elysiumjs/mnemosyne';
-import type { DrizzleConnectionProps } from './database';
-
-import { Application, Event } from '@elysiumjs/core';
-import { registerTenancyStrategy } from '@elysiumjs/mnemosyne';
-
-import { Database } from './database';
-import { DrizzleRLSTenancy, DrizzleSchemaTenancy } from './tenancy';
-
 // Re-export drizzle-orm/pg-core with enhanced pgTable
+export * from 'drizzle-orm';
 export * from 'drizzle-orm/pg-core';
+export type {
+	SelectedFields,
+	SelectedFieldsFlat,
+	SelectedFieldsOrdered,
+	TableConfig
+} from 'drizzle-orm/pg-core';
 export { pgTable, getTableBuilders } from './table';
 
 // Re-exports
@@ -34,6 +32,7 @@ export {
 	drizzleAdapter,
 	type DrizzleModelOptions
 } from './model';
+export { DrizzleModule, type DrizzleModuleConfig } from './drizzle.module';
 export { DrizzleRepository } from './repository';
 export {
 	DrizzleSchemaTenancy,
@@ -46,24 +45,3 @@ export {
 	getSessionTenant,
 	createRLSPolicy
 } from './tenancy';
-
-// Auto-initialize on app launch
-Event.once('elysium:app:launched', () => {
-	const app = Application.instance;
-	const config = app.getConfig<MnemosyneConfig<DrizzleConnectionProps>>('elysium:mnemosyne' as any);
-
-	if (config?.database) {
-		for (const [name, props] of Object.entries(config.database.connections)) {
-			Database.registerConnection(name, props);
-		}
-		if (Database.connectionExists(config.database.default)) {
-			Database.setDefaultConnection(config.database.default);
-		}
-	}
-
-	if (config?.tenancy) {
-		const strategy =
-			config.tenancy.mode === 'rls' ? new DrizzleRLSTenancy() : new DrizzleSchemaTenancy();
-		registerTenancyStrategy(strategy);
-	}
-});

package/src/migrations.ts

@@ -0,0 +1,100 @@
+// Copyright (c) 2025-present Workbud Technologies Inc. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+import { basename, dirname } from 'node:path';
+
+/** A migration recorded as applied in `elysium_migrations`. */
+export type AppliedMigration = {
+	name: string;
+	appliedAt: Date;
+};
+
+/** The migrations to run, and the applied records to rename to their new names. */
+export type MigrationPlan = {
+	pending: string[];
+	renames: { from: string; to: string }[];
+};
+
+const FOLDER_MIGRATION_FILE = 'migration.sql';
+
+/**
+ * Gets the name under which a migration file is tracked. Drizzle Kit 1.0 stores each migration
+ * as `<folder>/migration.sql`, in which case the folder name identifies the migration.
+ * @author Axel Nana <axel.nana@workbud.com>
+ * @param path The path of the migration file.
+ */
+export const getMigrationName = (path: string): string =>
+	basename(path) === FOLDER_MIGRATION_FILE ? basename(dirname(path)) : basename(path);
+
+/** Drops the `.sql` extension and the leading index or timestamp: `0001_users.sql` → `users`. */
+const descriptiveName = (name: string): string =>
+	name
+		.replace(/\.sql$/, '')
+		.split('_')
+		.slice(1)
+		.join('_');
+
+/** Reads the UTC `YYYYMMDDHHmmss` timestamp that starts a Drizzle Kit 1.0 folder name. */
+const folderTime = (name: string): number | undefined => {
+	const match = /^(\d{4})(\d{2})(\d{2})(\d{2})(\d{2})(\d{2})/.exec(name);
+	return match
+		? Date.UTC(+match[1]!, +match[2]! - 1, +match[3]!, +match[4]!, +match[5]!, +match[6]!)
+		: undefined;
+};
+
+/**
+ * Decides which migration files to run.
+ *
+ * `drizzle-kit up` moves each flat `<prefix>_<name>.sql` file into a
+ * `<timestamp>_<name>/migration.sql` folder, where the timestamp is when the migration was
+ * generated. An applied record whose flat file is gone is matched, in order, to a folder with
+ * the same `<name>` generated before the record was applied; that folder is not run again and
+ * the record is renamed. Folders generated later are new migrations and run.
+ * @author Axel Nana <axel.nana@workbud.com>
+ * @param files The paths of the migration files.
+ * @param applied The applied migrations.
+ */
+export const planMigrations = (files: string[], applied: AppliedMigration[]): MigrationPlan => {
+	const appliedNames = new Set(applied.map(({ name }) => name));
+	const fileNames = new Set(files.map(getMigrationName));
+
+	const moved = new Map<string, AppliedMigration[]>();
+	for (const record of [...applied].sort((a, b) => a.name.localeCompare(b.name))) {
+		if (!record.name.endsWith('.sql') || fileNames.has(record.name)) continue;
+		const key = descriptiveName(record.name);
+		moved.set(key, [...(moved.get(key) ?? []), record]);
+	}
+
+	const plan: MigrationPlan = { pending: [], renames: [] };
+	for (const file of [...files].sort()) {
+		const name = getMigrationName(file);
+		if (appliedNames.has(name)) continue;
+
+		const from = basename(file) === FOLDER_MIGRATION_FILE ? takeMoved(moved, name) : undefined;
+		if (from) plan.renames.push({ from, to: name });
+		else plan.pending.push(file);
+	}
+	return plan;
+};
+
+/** Removes and returns the first moved record that the folder `name` can be the new home of. */
+const takeMoved = (moved: Map<string, AppliedMigration[]>, name: string): string | undefined => {
+	const records = moved.get(descriptiveName(name));
+	const generatedAt = folderTime(name);
+	const index =
+		records?.findIndex(
+			({ appliedAt }) => generatedAt === undefined || generatedAt <= appliedAt.getTime()
+		) ?? -1;
+	return index < 0 ? undefined : records!.splice(index, 1)[0]!.name;
+};

package/src/model.ts

@@ -13,7 +13,7 @@
 // limitations under the License.
 
 import type { ColumnMetadata, ModelAdapter } from '@elysiumjs/mnemosyne';
-import type { PgColumn, PgTableWithColumns } from 'drizzle-orm/pg-core';
+import type { PgColumn, PgTableWithColumns, TableConfig } from 'drizzle-orm/pg-core';
 
 import { AbstractModel, createSchemaFromModel } from '@elysiumjs/mnemosyne';
 import { getTableConfig } from 'drizzle-orm/pg-core';
@@ -32,23 +32,30 @@ import { wrapTenantSchema } from './tenancy';
  */
 const mapDrizzleDataType = (col: PgColumn): ColumnMetadata['dataType'] => {
 	if (col.columnType === 'PgUUID') return 'uuid';
-	switch (col.dataType) {
-		case 'string':
-			return 'string';
+	if ((col as { dimensions?: number }).dimensions) return 'array';
+
+	// Drizzle 1.0 column data types are `<kind>` or `<kind> <constraint>` (e.g. `object json`).
+	const [kind, constraint] = col.dataType.split(' ');
+	switch (kind) {
 		case 'number':
 			return 'number';
 		case 'boolean':
 			return 'boolean';
 		case 'array':
 			return 'array';
-		case 'json':
-			return 'json';
-		case 'date':
-			return 'date';
 		case 'bigint':
 			return 'bigint';
-		case 'buffer':
-			return 'buffer';
+		case 'object':
+			switch (constraint) {
+				case 'json':
+					return 'json';
+				case 'date':
+					return 'date';
+				case 'buffer':
+					return 'buffer';
+				default:
+					return 'string';
+			}
 		default:
 			return 'string';
 	}
@@ -145,7 +152,7 @@ export interface DrizzleModelOptions {
  * @param options Model configuration options.
  * @returns A model class with Drizzle-backed static metadata.
  */
-export const DrizzleModel = <TTable extends PgTableWithColumns<any>>(
+export const DrizzleModel = <TTable extends PgTableWithColumns<TableConfig>>(
 	table: TTable,
 	options?: DrizzleModelOptions
 ) => {

package/src/table.ts

@@ -12,7 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-import type { PgColumnBuilderBase, PgTableWithColumns } from 'drizzle-orm/pg-core';
+import type { AnyPgColumnBuilder, PgTableWithColumns } from 'drizzle-orm/pg-core';
 
 import { pgTable as drizzlePgTable } from 'drizzle-orm/pg-core';
 
@@ -27,7 +27,7 @@ import { pgTable as drizzlePgTable } from 'drizzle-orm/pg-core';
  * built `PgColumn` instances. This registry preserves the original builder
  * objects so that tenant tables can be re-created in different schemas.
  */
-const tableRegistry = new WeakMap<PgTableWithColumns<any>, Record<string, PgColumnBuilderBase>>();
+const tableRegistry = new WeakMap<PgTableWithColumns<any>, Record<string, AnyPgColumnBuilder>>();
 
 /**
  * Retrieves the original column builders for a table.
@@ -36,7 +36,7 @@ const tableRegistry = new WeakMap<PgTableWithColumns<any>, Record<string, PgColu
  */
 export const getTableBuilders = (
 	table: PgTableWithColumns<any>
-): Record<string, PgColumnBuilderBase> | undefined => {
+): Record<string, AnyPgColumnBuilder> | undefined => {
 	return tableRegistry.get(table);
 };
 
@@ -58,7 +58,7 @@ export const getTableBuilders = (
  */
 export const pgTable = <
 	TTableName extends string,
-	TColumns extends Record<string, PgColumnBuilderBase>
+	TColumns extends Record<string, AnyPgColumnBuilder>
 >(
 	tableName: TTableName,
 	columns: TColumns,

package/src/tenancy.ts

@@ -13,7 +13,7 @@
 // limitations under the License.
 
 import type { TenancyStrategy } from '@elysiumjs/mnemosyne';
-import type { PgColumnBuilderBase, PgSchema, PgTable, PgTransaction } from 'drizzle-orm/pg-core';
+import type { AnyPgColumnBuilder, PgAsyncTransaction, PgSchema, PgTable } from 'drizzle-orm/pg-core';
 import type { DrizzleConnection } from './database';
 
 import { getTenancyConfig } from '@elysiumjs/mnemosyne';
@@ -73,7 +73,7 @@ export const registerTenantSchema = (tenant: string): PgSchema => {
  */
 export const wrapTenantSchema = <
 	TTableName extends string,
-	TColumnsMap extends Record<string, PgColumnBuilderBase>
+	TColumnsMap extends Record<string, AnyPgColumnBuilder>
 >(
 	tenant: string,
 	tableName: TTableName,
@@ -109,7 +109,7 @@ export const wrapTenantSchema = <
  */
 export const withRLS = async <T>(
 	tenant: string | null,
-	callback: (tx: PgTransaction<any, any, any>) => Promise<T>,
+	callback: (tx: PgAsyncTransaction<any, any>) => Promise<T>,
 	connectionName: string = 'default'
 ): Promise<T> => {
 	const db = Database.getConnection(connectionName);
@@ -117,7 +117,7 @@ export const withRLS = async <T>(
 	const varName = config.rls?.sessionVariable ?? 'app.current_tenant';
 	const value = tenant ?? 'public';
 
-	return (db as any).transaction(async (tx: PgTransaction<any, any, any>) => {
+	return (db as any).transaction(async (tx: PgAsyncTransaction<any, any>) => {
 		await tx.execute(sql`SELECT set_config(${varName}, ${value}, true)`);
 		return callback(tx);
 	});
@@ -241,7 +241,7 @@ export class DrizzleRLSTenancy implements TenancyStrategy<PgTable, DrizzleConnec
 		const config = getTenancyConfig();
 		const varName = config.rls?.sessionVariable ?? 'app.current_tenant';
 
-		return (connection as any).transaction(async (tx: PgTransaction<any, any, any>) => {
+		return (connection as any).transaction(async (tx: PgAsyncTransaction<any, any>) => {
 			await tx.execute(sql`SELECT set_config(${varName}, ${tenant}, true)`);
 			return callback();
 		});

package/tests/drizzle.module.test.ts

@@ -0,0 +1,201 @@
+// Copyright (c) 2026-present Workbud Technologies Inc. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+import { Application, Event, Service } from '@elysiumjs/core';
+import { getTenancyConfig, getTenancyStrategy } from '@elysiumjs/mnemosyne';
+import { afterEach, beforeEach, describe, expect, it, mock, spyOn } from 'bun:test';
+
+import * as pkg from '../src';
+import { Database, DrizzleDatabaseCache, DrizzleModule, DrizzleRLSTenancy, DrizzleSchemaTenancy } from '../src';
+
+const URL = 'postgres://u:p@localhost:5432/db';
+
+const stubClient = () => ({
+	options: {},
+	connect: mock(async () => {}),
+	begin: mock(async (_callback: unknown) => {}),
+	close: mock(async () => {})
+});
+
+describe('DrizzleModule', () => {
+	let app: Awaited<ReturnType<typeof Application.testing>> | undefined;
+
+	afterEach(async () => {
+		await app?.shutdown();
+		app = undefined;
+	});
+
+	it('registers nothing when the package is imported', async () => {
+		const once = spyOn(Event, 'once');
+		const on = spyOn(Event, 'on');
+		await import('../src/index.ts?fresh-import');
+		expect(once).not.toHaveBeenCalled();
+		expect(on).not.toHaveBeenCalled();
+		once.mockRestore();
+		on.mockRestore();
+		expect(Service.exists('db.connection.default')).toBe(false);
+		expect(Service.exists('mnemosyne.tenancy.config')).toBe(false);
+		expect(typeof pkg.DrizzleModule.with).toBe('function');
+	});
+
+	it('provides connections, the default alias and the schema strategy by default', async () => {
+		app = await Application.testing({
+			modules: [DrizzleModule.with({ default: 'main', connections: { main: { connection: URL } } })],
+			boot: 'eager'
+		});
+
+		expect(Database.connectionExists('main')).toBe(true);
+		expect(Database.connectionExists('default')).toBe(true);
+		const main = Database.getConnection('main');
+		expect(main).toBeDefined();
+		expect(Database.getConnection()).toBe(main);
+		expect(getTenancyConfig()).toEqual({ mode: 'schema' });
+		expect(getTenancyStrategy()).toBeInstanceOf(DrizzleSchemaTenancy);
+	});
+
+	it('uses the RLS strategy when tenancy mode is rls', async () => {
+		app = await Application.testing({
+			modules: [
+				DrizzleModule.with({
+					connections: { default: { connection: URL } },
+					tenancy: { mode: 'rls' }
+				})
+			],
+			boot: 'eager'
+		});
+
+		expect(getTenancyConfig()).toEqual({ mode: 'rls' });
+		expect(getTenancyStrategy()).toBeInstanceOf(DrizzleRLSTenancy);
+	});
+
+	it('rejects a default that is not a declared connection', async () => {
+		await expect(
+			Application.testing({
+				modules: [DrizzleModule.with({ default: 'nope', connections: { main: { connection: URL } } })]
+			})
+		).rejects.toThrow(/nope/);
+	});
+
+	it('builds a connection with a cache on a lazy boot', async () => {
+		app = await Application.testing({
+			modules: [
+				DrizzleModule.with({
+					connections: { default: { connection: URL, cache: { storage: 'memory' } } }
+				})
+			],
+			boot: 'lazy'
+		});
+
+		expect(await app.get('db.connection.default')).toBeDefined();
+	});
+
+	it('resolves its cache storage on first use', async () => {
+		app = await Application.testing({
+			modules: [DrizzleModule.with({ connections: { default: { connection: URL } } })],
+			boot: 'eager'
+		});
+		const cache = new DrizzleDatabaseCache('all', 1, 'memory');
+		await cache.put('k', [1], ['t'], false);
+		expect(await cache.get('k')).toEqual([1]);
+	});
+
+	it('closes only the connections that were created', async () => {
+		const used = stubClient();
+		const unused = stubClient();
+		app = await Application.testing({
+			modules: [
+				DrizzleModule.with({
+					default: 'a',
+					connections: { a: { client: used as never }, b: { client: unused as never } }
+				})
+			],
+			boot: 'lazy'
+		});
+		await app.get('db.connection.default');
+
+		await app.shutdown();
+		app = undefined;
+
+		expect(used.close).toHaveBeenCalledTimes(1);
+		expect(unused.close).not.toHaveBeenCalled();
+	});
+});
+
+describe('migration commands', () => {
+	const originalExit = process.exit;
+	const originalStderr = process.stderr.write;
+	const originalStdout = process.stdout.write;
+	let errors = '';
+
+	beforeEach(() => {
+		errors = '';
+		process.exit = mock() as never;
+		process.stderr.write = mock((message: string) => {
+			errors += message;
+			return true;
+		}) as never;
+		process.stdout.write = mock(() => true) as never;
+	});
+
+	let started: Application | undefined;
+
+	afterEach(async () => {
+		await started?.kernel.shutdown();
+		started = undefined;
+		Service.clear();
+		process.exit = originalExit;
+		process.stderr.write = originalStderr;
+		process.stdout.write = originalStdout;
+		for (const signal of ['SIGINT', 'SIGTERM'] as const) process.removeAllListeners(signal);
+	});
+
+	const makeApp = () => {
+		const client = stubClient();
+		@Application.register({
+			modules: [
+				DrizzleModule.with({ connections: { default: { client: client as never } } })
+			]
+		})
+		class TestApp extends Application {}
+		const app = new TestApp();
+		started = app;
+		return { app, client };
+	};
+
+	it('migration:run finds db.connection.default', async () => {
+		const { app, client } = makeApp();
+		await app.run(['migration:run', '--tenant', 'public']);
+
+		expect(errors).not.toContain('does not exist');
+		expect(client.connect).toHaveBeenCalled();
+		expect(client.begin).toHaveBeenCalled();
+		expect(Database.connectionExists('default')).toBe(true);
+	});
+
+	it('lists the migration commands under NODE_ENV=production', async () => {
+		const previous = Bun.env.NODE_ENV;
+		Bun.env.NODE_ENV = 'production';
+		try {
+			const { app } = makeApp();
+			await app.run(['list']);
+			const names = (app.kernel.graph.collect('commands') as { item: { command: string } }[]).map(
+				({ item }) => item.command
+			);
+			expect(names).toContain('migration:run');
+			expect(names).toContain('migration:generate');
+		} finally {
+			Bun.env.NODE_ENV = previous;
+		}
+	});
+});