@apisr/drizzle-model 2.0.2 → 2.0.4

package/src/model/builder.ts

@@ -5,6 +5,78 @@ import type { ModelDialect } from "./dialect.ts";
 import type { Model } from "./model.ts";
 import type { ModelOptions } from "./options.ts";
 
+type AnyRecord = Record<string, unknown>;
+
+/**
+ * Wraps a {@link ModelRuntime} instance into a plain target object that:
+ *
+ * 1. Exposes all runtime prototype methods, intercepting any that return a
+ *    `ModelRuntime` so the result is re-wrapped (preserving custom methods).
+ * 2. Defines the `$model`, `$modelName`, and `$format` identity getters.
+ * 3. Attaches user-defined custom methods from `options.methods` via
+ *    `runtime.attachMethods`, binding each to the new `target`.
+ *
+ * This must be called every time a new `ModelRuntime` is produced (e.g. by
+ * `.where()`, `.extend()`, `.db()`) so that custom methods are never lost.
+ */
+function buildTarget(runtime: ModelRuntime): AnyRecord {
+	const target: AnyRecord = {};
+
+	for (const key of Object.getOwnPropertyNames(
+		Object.getPrototypeOf(runtime)
+	)) {
+		if (key === "constructor") {
+			continue;
+		}
+
+		const value = (runtime as unknown as AnyRecord)[key];
+
+		if (typeof value === "function") {
+			target[key] = (...args: unknown[]) => {
+				const result = (value as (...args: unknown[]) => unknown).apply(
+					runtime,
+					args
+				);
+
+				// Lifecycle methods (where, extend, db) return a new ModelRuntime.
+				// Re-wrap it so custom methods are re-attached on the fresh target.
+				if (result instanceof ModelRuntime) {
+					return buildTarget(result);
+				}
+
+				return result;
+			};
+		}
+	}
+
+	Object.defineProperty(target, "$model", {
+		get: () => runtime.$model,
+		enumerable: true,
+	});
+	Object.defineProperty(target, "$modelName", {
+		get: () => runtime.$modelName,
+		enumerable: true,
+	});
+	Object.defineProperty(target, "$format", {
+		get: () => runtime.$format,
+		enumerable: true,
+	});
+	Object.defineProperty(target, "$where", {
+		get: () => runtime.$where,
+		enumerable: true,
+	});
+	Object.defineProperty(target, "$tableName", {
+		get: () => runtime.$tableName,
+		enumerable: true,
+	});
+
+	// Attach user-defined methods, binding each to this target so that `this`
+	// inside a custom method refers to the model-like object (with all methods).
+	runtime.attachMethods(target);
+
+	return target;
+}
+
 export function modelBuilder<
 	TFullSchema extends Record<string, unknown> = Record<string, never>,
 	TRelations extends AnyRelations = EmptyRelations,
@@ -44,36 +116,7 @@ export function modelBuilder<
 			options: (options ?? {}) as any,
 		});
 
-		const target: Record<string, unknown> = {};
-
-		for (const key of Object.getOwnPropertyNames(
-			Object.getPrototypeOf(runtime)
-		)) {
-			if (key === "constructor") {
-				continue;
-			}
-			const value = (runtime as unknown as Record<string, unknown>)[key];
-			if (typeof value === "function") {
-				target[key] = value.bind(runtime);
-			}
-		}
-
-		Object.defineProperty(target, "$model", {
-			get: () => runtime.$model,
-			enumerable: true,
-		});
-		Object.defineProperty(target, "$modelName", {
-			get: () => runtime.$modelName,
-			enumerable: true,
-		});
-		Object.defineProperty(target, "$format", {
-			get: () => runtime.$format,
-			enumerable: true,
-		});
-
-		runtime.attachMethods(target);
-
-		return target as unknown as Model<
+		return buildTarget(runtime) as unknown as Model<
 			ModelConfig<TRelations, TRelations[TTableName], TDialect, TOptions>
 		>;
 	};

package/tests/base/custom-methods.test.ts

@@ -0,0 +1,335 @@
+import { describe, expect, test } from "bun:test";
+import { model } from "tests/base";
+import { db } from "tests/db";
+import { esc } from "@/model";
+
+// ─── Helpers ──────────────────────────────────────────────────────────────────
+
+function uid(): string {
+	return `${Date.now()}-${Math.random().toString(36).slice(2)}`;
+}
+
+// ─── Models ───────────────────────────────────────────────────────────────────
+
+/**
+ * A user model that defines several custom methods.
+ * These must remain accessible after every lifecycle call
+ * (.where, .extend, .db).
+ */
+const userModel = model("user", {});
+
+const serviceQueries = userModel.extend({
+	methods: {
+		findByName(name: string) {
+			return serviceQueries.where({ name: esc(name) }).findFirst();
+		},
+		findAllVerified() {
+			return serviceQueries.where({ isVerified: esc(true) }).findMany();
+		},
+		findByMinAge(minAge: number) {
+			return serviceQueries.where({ age: { gte: esc(minAge) } }).findMany();
+		},
+	},
+});
+
+const postsModel = model("userPosts", {});
+
+const postQueries = postsModel.extend({
+	methods: {
+		findByTitle(title: string) {
+			return postQueries.where({ title: esc(title) }).findFirst();
+		},
+		findFeatured() {
+			return postQueries.where({ featured: esc(true) }).findMany();
+		},
+	},
+});
+
+// ─── Tests ────────────────────────────────────────────────────────────────────
+
+describe("custom methods", () => {
+	// -------------------------------------------------------------------------
+	// Presence — methods must be callable on the initial model object
+	// -------------------------------------------------------------------------
+
+	describe("presence on initial model", () => {
+		test("all declared methods are functions on the model", () => {
+			expect(typeof serviceQueries.findByName).toBe("function");
+			expect(typeof serviceQueries.findAllVerified).toBe("function");
+			expect(typeof serviceQueries.findByMinAge).toBe("function");
+		});
+
+		test("multiple models keep their own method sets independent", () => {
+			expect(typeof serviceQueries.findByName).toBe("function");
+			expect(typeof postQueries.findByTitle).toBe("function");
+			// cross-model — methods must NOT bleed across models
+			expect(
+				(serviceQueries as { findByTitle?: unknown }).findByTitle
+			).toBeUndefined();
+			expect(
+				(postQueries as { findByName?: unknown }).findByName
+			).toBeUndefined();
+		});
+	});
+
+	// -------------------------------------------------------------------------
+	// Survival through .where() — THE PRIMARY BUG
+	// Custom methods must survive every `.where()` call.
+	// -------------------------------------------------------------------------
+
+	describe("survival after .where()", () => {
+		test("custom methods are present on the result of .where()", () => {
+			const scoped = serviceQueries.where({ age: { gte: esc(0) } });
+			expect(typeof scoped.findByName).toBe("function");
+			expect(typeof scoped.findAllVerified).toBe("function");
+			expect(typeof scoped.findByMinAge).toBe("function");
+		});
+
+		test("custom methods are present after chained .where() calls", () => {
+			const scoped = serviceQueries
+				.where({ age: { gte: esc(0) } })
+				.where({ isVerified: esc(false) });
+			expect(typeof scoped.findByName).toBe("function");
+		});
+
+		test("base methods are still present after .where()", () => {
+			const scoped = serviceQueries.where({ age: { gte: esc(0) } });
+			expect(typeof scoped.findFirst).toBe("function");
+			expect(typeof scoped.findMany).toBe("function");
+			expect(typeof scoped.insert).toBe("function");
+			expect(typeof scoped.update).toBe("function");
+			expect(typeof scoped.delete).toBe("function");
+		});
+	});
+
+	// -------------------------------------------------------------------------
+	// Survival through .extend()
+	// -------------------------------------------------------------------------
+
+	describe("survival after .extend()", () => {
+		test("original custom methods survive .extend()", () => {
+			const extended = serviceQueries.extend({});
+			expect(typeof extended.findByName).toBe("function");
+			expect(typeof extended.findAllVerified).toBe("function");
+		});
+
+		test("new methods added via .extend() are accessible", () => {
+			const extended = serviceQueries.extend({
+				methods: {
+					findAdults() {
+						return extended.where({ age: { gte: esc(18) } }).findMany();
+					},
+				},
+			});
+			expect(typeof extended.findAdults).toBe("function");
+		});
+
+		test("parent custom methods survive on an extended model", () => {
+			const extended = serviceQueries.extend({
+				methods: {
+					findAdults() {
+						return extended.where({ age: { gte: esc(18) } }).findMany();
+					},
+				},
+			});
+			// Parent methods must still be present on the child
+			expect(typeof extended.findByName).toBe("function");
+			expect(typeof extended.findAllVerified).toBe("function");
+			// New method must be present too
+			expect(typeof extended.findAdults).toBe("function");
+		});
+
+		test("child method overrides parent method of the same name", () => {
+			const SENTINEL = Symbol("override-sentinel");
+			const extended = serviceQueries.extend({
+				methods: {
+					findByName(_name: string) {
+						return SENTINEL;
+					},
+				},
+			});
+			// The extended model's findByName must shadow the parent's
+			expect(extended.findByName("ignored")).toBe(SENTINEL);
+		});
+
+		test("custom methods survive .where() on an extended model", () => {
+			const extended = serviceQueries.extend({
+				methods: {
+					findAdults() {
+						return extended.where({ age: { gte: esc(18) } }).findMany();
+					},
+				},
+			});
+			const scoped = extended.where({ isVerified: esc(true) });
+			expect(typeof scoped.findByName).toBe("function");
+			expect(typeof scoped.findAdults).toBe("function");
+		});
+	});
+
+	// -------------------------------------------------------------------------
+	// Survival through .db()
+	// -------------------------------------------------------------------------
+
+	describe("survival after .db()", () => {
+		test("custom methods are present after .db() rebind", () => {
+			const rebound = serviceQueries.db(db);
+			expect(typeof rebound.findByName).toBe("function");
+			expect(typeof rebound.findAllVerified).toBe("function");
+			expect(typeof rebound.findByMinAge).toBe("function");
+		});
+
+		test("base methods are still present after .db() rebind", () => {
+			const rebound = serviceQueries.db(db);
+			expect(typeof rebound.findFirst).toBe("function");
+			expect(typeof rebound.findMany).toBe("function");
+		});
+
+		test("custom methods survive .where() after .db() rebind", () => {
+			const rebound = serviceQueries.db(db).where({ age: { gte: esc(0) } });
+			expect(typeof rebound.findByName).toBe("function");
+		});
+	});
+
+	// -------------------------------------------------------------------------
+	// Functional correctness — custom methods actually return the right data
+	// -------------------------------------------------------------------------
+
+	describe("functional correctness", () => {
+		test("findByName returns a single user with matching name", async () => {
+			const uniqueName = `TestUser-${uid()}`;
+			const uniqueEmail = `${uid()}@custom-methods.test`;
+
+			await serviceQueries.insert({
+				name: uniqueName,
+				email: uniqueEmail,
+				age: 25,
+			});
+
+			const result = await serviceQueries.findByName(uniqueName);
+
+			expect(result).toBeDefined();
+			expect(result?.name).toBe(uniqueName);
+			expect(typeof result?.id).toBe("number");
+		});
+
+		test("findAllVerified returns only verified users", async () => {
+			const result = await serviceQueries.findAllVerified();
+
+			expect(Array.isArray(result)).toBe(true);
+			for (const user of result) {
+				expect(user.isVerified).toBe(true);
+			}
+		});
+
+		test("findByMinAge returns users at or above the age threshold", async () => {
+			const result = await serviceQueries.findByMinAge(18);
+
+			expect(Array.isArray(result)).toBe(true);
+			for (const user of result) {
+				expect(user.age).toBeGreaterThanOrEqual(18);
+			}
+		});
+
+		test("custom method result is consistent with equivalent base query", async () => {
+			const uniqueName = `Consistent-${uid()}`;
+			const uniqueEmail = `${uid()}@custom-methods.test`;
+
+			await serviceQueries.insert({
+				name: uniqueName,
+				email: uniqueEmail,
+				age: 30,
+			});
+
+			const viaCustom = await serviceQueries.findByName(uniqueName);
+			const viaBase = await serviceQueries
+				.where({ name: esc(uniqueName) })
+				.findFirst();
+
+			expect(viaCustom).toEqual(viaBase);
+		});
+
+		test("posts custom method findByTitle returns correct post", async () => {
+			const uniqueTitle = `PostTitle-${uid()}`;
+
+			await postQueries.insert({
+				title: uniqueTitle,
+				featured: false,
+				userId: 1,
+			});
+
+			const result = await postQueries.findByTitle(uniqueTitle);
+
+			expect(result).toBeDefined();
+			expect(result?.title).toBe(uniqueTitle);
+		});
+
+		test("posts custom method findFeatured returns only featured posts", async () => {
+			await postQueries.insert({
+				title: `Featured-${uid()}`,
+				featured: true,
+				userId: 1,
+			});
+
+			const result = await postQueries.findFeatured();
+
+			expect(Array.isArray(result)).toBe(true);
+			expect(result.length).toBeGreaterThan(0);
+			for (const post of result) {
+				expect(post.featured).toBe(true);
+			}
+		});
+
+		test("custom method is still functional after .where() (the bug scenario)", async () => {
+			const uniqueName = `BugCheck-${uid()}`;
+			const uniqueEmail = `${uid()}@custom-methods.test`;
+
+			await serviceQueries.insert({
+				name: uniqueName,
+				email: uniqueEmail,
+				age: 99,
+			});
+
+			const scopedModel = serviceQueries.where({ age: { gte: esc(0) } });
+			const result = await scopedModel.findByName(uniqueName);
+
+			expect(result).toBeDefined();
+			expect(result?.name).toBe(uniqueName);
+		});
+
+		test("custom method is still functional after .extend()", async () => {
+			const uniqueName = `ExtendBug-${uid()}`;
+			const uniqueEmail = `${uid()}@custom-methods.test`;
+
+			await serviceQueries.insert({
+				name: uniqueName,
+				email: uniqueEmail,
+				age: 22,
+			});
+
+			const extended = serviceQueries.extend({
+				format: (u) => ({ ...u, _extended: true }),
+			});
+			const result = await extended.findByName(uniqueName);
+
+			expect(result).toBeDefined();
+			expect(result?.name).toBe(uniqueName);
+		});
+
+		test("custom method is still functional after .db()", async () => {
+			const uniqueName = `DbBug-${uid()}`;
+			const uniqueEmail = `${uid()}@custom-methods.test`;
+
+			await serviceQueries.insert({
+				name: uniqueName,
+				email: uniqueEmail,
+				age: 33,
+			});
+
+			const rebound = serviceQueries.db(db);
+			const result = await rebound.findByName(uniqueName);
+
+			expect(result).toBeDefined();
+			expect(result?.name).toBe(uniqueName);
+		});
+	});
+});

package/tests/snippets/x-2.ts

@@ -13,6 +13,7 @@ const model = modelBuilder({
 // create model
 const userModel = model("user", {});
 
+// #1 syntax
 await userModel
 	.where({
 		name: {
@@ -21,6 +22,7 @@ await userModel
 	})
 	.findFirst();
 
+// #2 syntax
 await userModel
 	.where({
 		name: esc.like("A%"),

package/dist/model/core/joins.mjs

@@ -1,184 +0,0 @@
-import { and, eq } from "drizzle-orm";
-
-//#region src/model/core/joins.ts
-function isDrizzleColumn(value) {
-	return !!value && typeof value === "object" && typeof value.getSQL === "function";
-}
-function getPrimaryKeyField(table) {
-	for (const [k, v] of Object.entries(table)) {
-		if (!isDrizzleColumn(v)) continue;
-		if (v.primary === true) return k;
-		if (v.config?.primaryKey === true) return k;
-	}
-	if ("id" in table) return "id";
-	return Object.keys(table).find((k) => isDrizzleColumn(table[k])) ?? "id";
-}
-function isAllNullRow(obj) {
-	if (!obj || typeof obj !== "object") return true;
-	for (const v of Object.values(obj)) if (v !== null && v !== void 0) return false;
-	return true;
-}
-async function aliasTable(table, aliasName, dialect) {
-	if (dialect === "PostgreSQL") {
-		const mod = await import("drizzle-orm/pg-core");
-		if (typeof mod.alias === "function") return mod.alias(table, aliasName);
-	}
-	if (dialect === "MySQL") {
-		const mod = await import("drizzle-orm/mysql-core");
-		if (typeof mod.alias === "function") return mod.alias(table, aliasName);
-	}
-	if (dialect === "SQLite") {
-		const mod = await import("drizzle-orm/sqlite-core");
-		if (typeof mod.alias === "function") return mod.alias(table, aliasName);
-	}
-	return table;
-}
-function buildJoinOn(node) {
-	const parts = node.sourceColumns.map((src, i) => {
-		const tgt = node.targetColumns[i];
-		const tgtKey = Object.entries(node.targetTable).find(([, v]) => v === tgt)?.[0];
-		return eq(tgtKey ? node.targetAliasTable[tgtKey] : tgt, src);
-	});
-	return parts.length === 1 ? parts[0] : and(...parts);
-}
-function buildSelectMapForTable(table) {
-	const out = {};
-	for (const [k, v] of Object.entries(table)) if (isDrizzleColumn(v)) out[k] = v;
-	return out;
-}
-async function executeWithJoins(args) {
-	const { db, schema, relations, baseTableName, baseTable, dialect, whereSql, withValue, limitOne } = args;
-	const usedAliasKeys = /* @__PURE__ */ new Set();
-	const buildNode = async (parent, currentTableName, currentTable, key, value, path) => {
-		const relMeta = relations[currentTableName]?.relations?.[key];
-		if (!relMeta) throw new Error(`Unknown relation '${key}' on table '${currentTableName}'.`);
-		const targetTableName = relMeta.targetTableName;
-		const targetTable = schema[targetTableName];
-		const aliasKeyBase = [...path, key].join("__");
-		let aliasKey = aliasKeyBase;
-		let idx = 1;
-		while (usedAliasKeys.has(aliasKey)) aliasKey = `${aliasKeyBase}_${idx++}`;
-		usedAliasKeys.add(aliasKey);
-		const needsAlias = targetTableName === currentTableName || usedAliasKeys.has(`table:${targetTableName}`);
-		usedAliasKeys.add(`table:${targetTableName}`);
-		const targetAliasTable = needsAlias ? await aliasTable(targetTable, aliasKey, dialect) : targetTable;
-		const node = {
-			path: [...path, key],
-			key,
-			relationType: relMeta.relationType,
-			sourceTableName: currentTableName,
-			targetTableName,
-			sourceTable: currentTable,
-			targetTable,
-			targetAliasTable,
-			aliasKey,
-			sourceColumns: relMeta.sourceColumns ?? [],
-			targetColumns: relMeta.targetColumns ?? [],
-			pkField: getPrimaryKeyField(targetAliasTable),
-			parent,
-			children: []
-		};
-		if (value && typeof value === "object" && value !== true) for (const [childKey, childVal] of Object.entries(value)) {
-			if (childVal !== true && (typeof childVal !== "object" || childVal == null)) continue;
-			const child = await buildNode(node, targetTableName, targetAliasTable, childKey, childVal, [...path, key]);
-			node.children.push(child);
-		}
-		return node;
-	};
-	const root = {
-		path: [],
-		key: "$root",
-		relationType: "one",
-		sourceTableName: baseTableName,
-		targetTableName: baseTableName,
-		sourceTable: baseTable,
-		targetTable: baseTable,
-		targetAliasTable: baseTable,
-		aliasKey: "$base",
-		sourceColumns: [],
-		targetColumns: [],
-		pkField: getPrimaryKeyField(baseTable),
-		children: []
-	};
-	for (const [key, value] of Object.entries(withValue)) {
-		if (value !== true && (typeof value !== "object" || value == null)) continue;
-		const child = await buildNode(void 0, baseTableName, baseTable, key, value, []);
-		root.children.push(child);
-	}
-	const nodes = [];
-	const walk = (n) => {
-		for (const c of n.children) {
-			nodes.push(c);
-			walk(c);
-		}
-	};
-	walk(root);
-	const selectMap = { base: buildSelectMapForTable(baseTable) };
-	for (const n of nodes) selectMap[n.aliasKey] = buildSelectMapForTable(n.targetAliasTable);
-	let q = db.select(selectMap).from(baseTable);
-	if (whereSql) q = q.where(whereSql);
-	for (const n of nodes) {
-		const on = buildJoinOn(n);
-		q = q.leftJoin(n.targetAliasTable, on);
-	}
-	if (limitOne) q = q.limit(1);
-	const rows = await q;
-	const basePk = root.pkField;
-	const baseMap = /* @__PURE__ */ new Map();
-	const ensureManyContainer = (obj, key) => {
-		if (!Array.isArray(obj[key])) obj[key] = [];
-	};
-	const ensureOneContainer = (obj, key) => {
-		if (!(key in obj)) obj[key] = null;
-	};
-	const manyIndexByPath = /* @__PURE__ */ new Map();
-	for (const row of rows) {
-		const baseRow = row.base;
-		const baseId = baseRow[basePk];
-		if (baseId === void 0) continue;
-		const baseObj = (() => {
-			const existing = baseMap.get(baseId);
-			if (existing) return existing;
-			const created = { ...baseRow };
-			baseMap.set(baseId, created);
-			return created;
-		})();
-		for (const n of nodes) {
-			const data = row[n.aliasKey];
-			const relPath = n.path.join(".");
-			const parentPath = n.parent ? n.parent.path.join(".") : "";
-			let parentObj = baseObj;
-			if (parentPath) {
-				const parentIndex = manyIndexByPath.get(parentPath);
-				if (parentIndex && parentIndex.size) parentObj = Array.from(parentIndex.values()).at(-1);
-				else {
-					const parentKey = n.parent?.key;
-					parentObj = parentKey ? baseObj[parentKey] ?? baseObj : baseObj;
-				}
-			}
-			if (isAllNullRow(data)) {
-				if (n.relationType === "one") ensureOneContainer(parentObj, n.key);
-				else ensureManyContainer(parentObj, n.key);
-				continue;
-			}
-			const pk = data[n.pkField];
-			if (n.relationType === "one") parentObj[n.key] = { ...data };
-			else {
-				ensureManyContainer(parentObj, n.key);
-				const indexKey = relPath;
-				if (!manyIndexByPath.has(indexKey)) manyIndexByPath.set(indexKey, /* @__PURE__ */ new Map());
-				const idxMap = manyIndexByPath.get(indexKey);
-				if (!idxMap.has(pk)) {
-					const obj = { ...data };
-					idxMap.set(pk, obj);
-					parentObj[n.key].push(obj);
-				}
-			}
-		}
-	}
-	const out = Array.from(baseMap.values());
-	return limitOne ? out[0] : out;
-}
-
-//#endregion
-export { executeWithJoins };

package/dist/model/core/projection.mjs

@@ -1,28 +0,0 @@
-//#region src/model/core/projection.ts
-function isDrizzleColumn(value) {
-	return !!value && typeof value === "object" && typeof value.getSQL === "function";
-}
-function getTableColumnsMap(table) {
-	const out = {};
-	for (const [key, value] of Object.entries(table)) if (isDrizzleColumn(value)) out[key] = value;
-	return out;
-}
-function buildSelectProjection(table, select, exclude) {
-	const all = getTableColumnsMap(table);
-	if (select && typeof select === "object") {
-		const picked = {};
-		for (const [key, value] of Object.entries(select)) if (value === true && key in all) picked[key] = all[key];
-		if (Object.keys(picked).length) return { selectMap: picked };
-		return { selectMap: all };
-	}
-	if (exclude && typeof exclude === "object") {
-		const omitted = { ...all };
-		for (const [key, value] of Object.entries(exclude)) if (value === true) delete omitted[key];
-		if (Object.keys(omitted).length) return { selectMap: omitted };
-		return { selectMap: all };
-	}
-	return { selectMap: all };
-}
-
-//#endregion
-export { buildSelectProjection };