@elysiumjs/mnemosyne-drizzle 0.1.0

package/src/table.ts

@@ -0,0 +1,70 @@
+// 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 { PgColumnBuilderBase, PgTableWithColumns } from 'drizzle-orm/pg-core';
+
+import { pgTable as drizzlePgTable } from 'drizzle-orm/pg-core';
+
+// ============================================================================
+// Column Builder Registry
+// ============================================================================
+
+/**
+ * Maps tables to their original column builder definitions.
+ *
+ * When `pgTable()` is called, column builders are consumed and turned into
+ * 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>>();
+
+/**
+ * Retrieves the original column builders for a table.
+ * @param table The table created with enhanced pgTable.
+ * @returns The original column builder definitions, or `undefined`.
+ */
+export const getTableBuilders = (
+	table: PgTableWithColumns<any>
+): Record<string, PgColumnBuilderBase> | undefined => {
+	return tableRegistry.get(table);
+};
+
+// ============================================================================
+// Enhanced pgTable Function
+// ============================================================================
+
+/**
+ * Creates a PostgreSQL table with automatic column builder registration.
+ *
+ * This is an enhanced version of Drizzle's `pgTable` that automatically
+ * registers column builders for tenancy support. It has the exact same
+ * signature and behavior as the original.
+ *
+ * @param tableName The name of the database table.
+ * @param columns The column definitions.
+ * @param extraConfig Optional extra Drizzle table configuration.
+ * @returns A Drizzle PostgreSQL table with registered column builders.
+ */
+export const pgTable = <
+	TTableName extends string,
+	TColumns extends Record<string, PgColumnBuilderBase>
+>(
+	tableName: TTableName,
+	columns: TColumns,
+	extraConfig?: (table: any) => any
+) => {
+	const table = drizzlePgTable(tableName, columns, extraConfig);
+	tableRegistry.set(table, columns);
+	return table;
+};

package/src/tenancy.ts

@@ -0,0 +1,249 @@
+// 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 { TenancyStrategy } from '@elysiumjs/mnemosyne';
+import type { PgColumnBuilderBase, PgSchema, PgTable, PgTransaction } from 'drizzle-orm/pg-core';
+import type { DrizzleConnection } from './database';
+
+import { getTenancyConfig } from '@elysiumjs/mnemosyne';
+import { sql } from 'drizzle-orm';
+import { getTableConfig, pgPolicy, pgSchema, pgTable } from 'drizzle-orm/pg-core';
+
+import { Database } from './database';
+import { getTableBuilders } from './table';
+
+// ============================================================================
+// Internal State
+// ============================================================================
+
+const tenantSchemas: Map<string, PgTable> = new Map();
+const schemaRegistry: Map<string, PgSchema> = new Map();
+
+// ============================================================================
+// Schema-based Tenancy Helpers
+// ============================================================================
+
+/**
+ * Returns the Drizzle `PgSchema` for the given tenant, creating it if needed.
+ * @author Axel Nana <axel.nana@workbud.com>
+ * @param tenant The tenant identifier.
+ * @returns The tenant's PostgreSQL schema.
+ */
+export const getTenantSchema = (tenant: string): PgSchema => {
+	return schemaRegistry.get(tenant) ?? registerTenantSchema(tenant);
+};
+
+/**
+ * Registers a new tenant schema in the registry.
+ * @author Axel Nana <axel.nana@workbud.com>
+ * @param tenant The tenant identifier.
+ * @returns The newly registered (or existing) `PgSchema`.
+ */
+export const registerTenantSchema = (tenant: string): PgSchema => {
+	if (schemaRegistry.has(tenant)) {
+		return schemaRegistry.get(tenant)!;
+	}
+	const tenantSchema = pgSchema(tenant);
+	schemaRegistry.set(tenant, tenantSchema);
+	return tenantSchema;
+};
+
+/**
+ * Wraps a table definition in a tenant-specific schema.
+ *
+ * For the `'public'` tenant, returns a standard `pgTable`.
+ * For other tenants, creates a schema-qualified table and caches it.
+ *
+ * @author Axel Nana <axel.nana@workbud.com>
+ * @param tenant The tenant identifier.
+ * @param tableName The table name.
+ * @param columns The column definitions.
+ * @returns A tenant-scoped `PgTable`.
+ */
+export const wrapTenantSchema = <
+	TTableName extends string,
+	TColumnsMap extends Record<string, PgColumnBuilderBase>
+>(
+	tenant: string,
+	tableName: TTableName,
+	columns: TColumnsMap
+): PgTable => {
+	if (tenant === 'public') {
+		return pgTable(tableName, columns);
+	}
+
+	const fullTableName = `${tenant}.${tableName}`;
+	if (tenantSchemas.has(fullTableName)) {
+		return tenantSchemas.get(fullTableName)!;
+	}
+
+	const tenantSchema = getTenantSchema(tenant);
+	const schemaTable = tenantSchema.table(tableName, columns);
+	tenantSchemas.set(fullTableName, schemaTable);
+
+	return schemaTable;
+};
+
+// ============================================================================
+// RLS-based Tenancy Helpers
+// ============================================================================
+
+/**
+ * Executes a callback within an RLS-scoped transaction for the given tenant.
+ * @author Axel Nana <axel.nana@workbud.com>
+ * @param tenant The tenant identifier (or `null` for 'public').
+ * @param callback The function to execute within the transaction.
+ * @param connectionName The database connection name.
+ * @returns The return value of the callback.
+ */
+export const withRLS = async <T>(
+	tenant: string | null,
+	callback: (tx: PgTransaction<any, any, any>) => Promise<T>,
+	connectionName: string = 'default'
+): Promise<T> => {
+	const db = Database.getConnection(connectionName);
+	const config = getTenancyConfig();
+	const varName = config.rls?.sessionVariable ?? 'app.current_tenant';
+	const value = tenant ?? 'public';
+
+	return (db as any).transaction(async (tx: PgTransaction<any, any, any>) => {
+		await tx.execute(sql`SELECT set_config(${varName}, ${value}, true)`);
+		return callback(tx);
+	});
+};
+
+/**
+ * Sets the current tenant on a persistent database connection (session-level).
+ * @author Axel Nana <axel.nana@workbud.com>
+ * @param tenant The tenant identifier (or `null` for 'public').
+ * @param connectionName The database connection name.
+ */
+export const setConnectionTenant = async (
+	tenant: string | null,
+	connectionName: string = 'default'
+): Promise<void> => {
+	const db = Database.getConnection(connectionName);
+	const config = getTenancyConfig();
+	const varName = config.rls?.sessionVariable ?? 'app.current_tenant';
+	const value = tenant ?? 'public';
+
+	await (db as any).execute(sql`SELECT set_config(${varName}, ${value}, false)`);
+};
+
+/**
+ * Reads the current tenant from the PostgreSQL session variable.
+ * @author Axel Nana <axel.nana@workbud.com>
+ * @param connectionName The database connection name.
+ * @returns The current tenant identifier, or `null`.
+ */
+export const getSessionTenant = async (
+	connectionName: string = 'default'
+): Promise<string | null> => {
+	const db = Database.getConnection(connectionName);
+	const config = getTenancyConfig();
+	const varName = config.rls?.sessionVariable ?? 'app.current_tenant';
+
+	const result = await (db as any).execute(sql`SELECT current_setting(${varName}, true) as tenant`);
+	return (result as any)?.rows?.[0]?.tenant ?? null;
+};
+
+/**
+ * Creates an RLS policy for tenant isolation on a table.
+ * @author Axel Nana <axel.nana@workbud.com>
+ * @param tableName The table name.
+ * @param columnName The tenant column name (defaults to config or 'tenant_id').
+ * @param policyName The policy name (defaults to `{tableName}_tenant_isolation`).
+ * @returns A Drizzle `PgPolicy` instance.
+ */
+export const createRLSPolicy = (tableName: string, columnName?: string, policyName?: string) => {
+	const config = getTenancyConfig();
+	const varName = config.rls?.sessionVariable ?? 'app.current_tenant';
+	const column = columnName ?? config.rls?.defaultColumn ?? 'tenant_id';
+	const name = policyName ?? `${tableName}_tenant_isolation`;
+
+	return pgPolicy(name, {
+		for: 'all',
+		to: 'public',
+		using: sql`${sql.identifier(column)} = current_setting(${varName}, true)::text`,
+		withCheck: sql`${sql.identifier(column)} = current_setting(${varName}, true)::text`
+	});
+};
+
+// ============================================================================
+// TenancyStrategy Implementations
+// ============================================================================
+
+/**
+ * Schema-based tenancy strategy for Drizzle.
+ *
+ * Each tenant gets its own PostgreSQL schema. Table resolution maps
+ * the base table to a schema-qualified equivalent.
+ *
+ * @author Axel Nana <axel.nana@workbud.com>
+ */
+export class DrizzleSchemaTenancy implements TenancyStrategy<PgTable, DrizzleConnection> {
+	readonly mode = 'schema';
+
+	resolveTable(table: PgTable, tenant: string): PgTable {
+		const builders = getTableBuilders(table as any);
+		if (!builders) {
+			throw new Error(
+				`No column builders registered for table. ` +
+					`Use pgTable() from '@elysiumjs/mnemosyne-drizzle' for tenant support.`
+			);
+		}
+		const config = getTableConfig(table as any);
+		return wrapTenantSchema(tenant, config.name, builders) as PgTable;
+	}
+
+	async withIsolation<T>(
+		_connection: DrizzleConnection,
+		_tenant: string,
+		callback: () => Promise<T>
+	): Promise<T> {
+		// Schema-based isolation doesn't need transaction wrapping
+		return callback();
+	}
+}
+
+/**
+ * Row-Level Security tenancy strategy for Drizzle.
+ *
+ * All tenants share the same schema. Isolation is achieved via
+ * PostgreSQL RLS policies and a session variable set in each transaction.
+ *
+ * @author Axel Nana <axel.nana@workbud.com>
+ */
+export class DrizzleRLSTenancy implements TenancyStrategy<PgTable, DrizzleConnection> {
+	readonly mode = 'rls';
+
+	resolveTable(table: PgTable, _tenant: string): PgTable {
+		// RLS doesn't modify the table - filtering is handled by policies
+		return table;
+	}
+
+	async withIsolation<T>(
+		connection: DrizzleConnection,
+		tenant: string,
+		callback: () => Promise<T>
+	): Promise<T> {
+		const config = getTenancyConfig();
+		const varName = config.rls?.sessionVariable ?? 'app.current_tenant';
+
+		return (connection as any).transaction(async (tx: PgTransaction<any, any, any>) => {
+			await tx.execute(sql`SELECT set_config(${varName}, ${tenant}, true)`);
+			return callback();
+		});
+	}
+}

package/tests/database.test.ts

@@ -0,0 +1,198 @@
+// 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 { afterAll, afterEach, beforeAll, beforeEach, describe, expect, it, jest, mock, spyOn } from 'bun:test';
+
+import { Application, Service } from '@elysiumjs/core';
+
+import { Database } from '../src/database';
+
+// Mock drizzle function
+mock.module('drizzle-orm/bun-sql', () => ({
+	drizzle: mock(() => ({ mockDrizzleInstance: true }))
+}));
+
+describe('DrizzleDatabase', () => {
+	// Mock console.error and process.exit to prevent actual console output and process exit
+	const originalConsoleError = console.error;
+	const originalProcessExit = process.exit;
+
+	beforeAll(() => {
+		// Mock Application.context so that getConnection('default') works
+		const mockContext = {
+			getStore: mock(() => new Map())
+		};
+		Object.defineProperty(Application, 'context', {
+			get: () => mockContext,
+			configurable: true
+		});
+	});
+
+	beforeEach(() => {
+		console.error = mock();
+		process.exit = mock() as any;
+	});
+
+	afterEach(() => {
+		console.error = originalConsoleError;
+		process.exit = originalProcessExit;
+		jest.clearAllMocks();
+	});
+
+	afterAll(() => {
+		mock.restore();
+	});
+
+	describe('registerConnection', () => {
+		it('should register a new connection', () => {
+			// Mock Service.exists to return false (connection doesn't exist)
+			const existsSpy = spyOn(Service, 'exists').mockReturnValueOnce(false);
+			const instanceSpy = spyOn(Service, 'instance');
+
+			// Call the function
+			const result = Database.registerConnection('test', { connection: 'sqlite:test.db' });
+
+			// Check if Service.exists was called with the correct connection name
+			expect(existsSpy).toHaveBeenCalledWith('db.connection.test');
+
+			// Check if Service.instance was called with the correct parameters
+			expect(instanceSpy).toHaveBeenCalledWith('db.connection.test', {
+				mockDrizzleInstance: true
+			});
+
+			// Check if the result is correct
+			expect(result as any).toEqual({ mockDrizzleInstance: true });
+		});
+
+		it('should exit if the connection already exists', () => {
+			// Mock Service.exists to return true (connection exists)
+			spyOn(Service, 'exists').mockReturnValueOnce(true);
+
+			// Call the function
+			Database.registerConnection('test', { connection: 'sqlite:test.db' });
+
+			// Check if console.error and process.exit were called
+			expect(console.error).toHaveBeenCalledWith(
+				expect.stringContaining('A connection with the name test has already been registered')
+			);
+			expect(process.exit).toHaveBeenCalledWith(1);
+		});
+	});
+
+	describe('getConnection', () => {
+		it('should retrieve an existing connection', () => {
+			// Mock Service.exists to return true (connection exists)
+			const existsSpy = spyOn(Service, 'exists').mockReturnValueOnce(true);
+
+			// Mock Service.get to return a mock connection
+			const mockConnection = { mockConnection: true };
+			const getSpy = spyOn(Service, 'get').mockReturnValue(mockConnection);
+
+			// Call the function
+			const result = Database.getConnection('test');
+
+			// Check if Service.exists was called with the correct connection name
+			expect(existsSpy).toHaveBeenCalledWith('db.connection.test');
+
+			// Check if Service.get was called with the correct connection name
+			expect(getSpy).toHaveBeenCalledWith('db.connection.test');
+
+			// Check if the result is correct
+			expect(result as any).toEqual(mockConnection);
+		});
+
+		it('should exit if the connection does not exist', () => {
+			// Mock Service.exists to return false (connection doesn't exist)
+			spyOn(Service, 'exists').mockReturnValueOnce(false);
+
+			// Call the function
+			Database.getConnection('test');
+
+			// Check if console.error and process.exit were called
+			expect(console.error).toHaveBeenCalledWith(
+				expect.stringContaining('No connection with name test found')
+			);
+			expect(process.exit).toHaveBeenCalledWith(1);
+		});
+	});
+
+	describe('connectionExists', () => {
+		it('should return true if the connection exists', () => {
+			// Mock Service.exists to return true
+			const existsSpy = spyOn(Service, 'exists').mockReturnValueOnce(true);
+
+			// Call the function
+			const result = Database.connectionExists('test');
+
+			// Check if Service.exists was called with the correct connection name
+			expect(existsSpy).toHaveBeenCalledWith('db.connection.test');
+
+			// Check if the result is correct
+			expect(result).toBe(true);
+		});
+
+		it('should return false if the connection does not exist', () => {
+			// Mock Service.exists to return false
+			const existsSpy = spyOn(Service, 'exists').mockReturnValueOnce(false);
+
+			// Call the function
+			const result = Database.connectionExists('test');
+
+			// Check if Service.exists was called with the correct connection name
+			expect(existsSpy).toHaveBeenCalledWith('db.connection.test');
+
+			// Check if the result is correct
+			expect(result).toBe(false);
+		});
+	});
+
+	describe('getDefaultConnection', () => {
+		it('should retrieve the default connection', () => {
+			// Mock Service to return the mock connection
+			const mockConnection = { mockConnection: true };
+			spyOn(Service, 'exists').mockReturnValueOnce(true);
+			spyOn(Service, 'get').mockReturnValue(mockConnection);
+
+			// Call the function
+			const result = Database.getDefaultConnection();
+
+			// Check if the result is correct
+			expect(result as any).toEqual(mockConnection);
+		});
+	});
+
+	describe('setDefaultConnection', () => {
+		it('should set the default connection', () => {
+			const instanceSpy = spyOn(Service, 'instance');
+			const removeSpy = spyOn(Service, 'remove');
+			spyOn(Service, 'exists').mockReturnValueOnce(true);
+
+			// Mock Service.get to return a mock connection
+			const mockConnection = { mockConnection: true };
+			spyOn(Service, 'get').mockReturnValue(mockConnection);
+
+			// Call the function
+			const result = Database.setDefaultConnection('test');
+
+			// Check if Service.remove was called with the correct connection name
+			expect(removeSpy).toHaveBeenCalledWith('db.connection.default');
+
+			// Check if Service.instance was called with the correct parameters
+			expect(instanceSpy).toHaveBeenCalledWith('db.connection.default', mockConnection);
+
+			// Check if the result is correct
+			expect(result as any).toEqual(mockConnection);
+		});
+	});
+});