@lobomfz/db 0.3.2 → 0.3.4

package/package.json

@@ -1,6 +1,6 @@
 {
 	"name": "@lobomfz/db",
-	"version": "0.3.2",
+	"version": "0.3.4",
 	"description": "Bun SQLite database with Arktype schemas and typed Kysely client",
 	"keywords": [
 		"arktype",

package/src/database.ts

@@ -63,6 +63,7 @@ export class Database<T extends SchemaRecord> {
 	private sqlite: BunDatabase;
 
 	private columns: ColumnsMap = new Map();
+	private tableColumns = new Map<string, Set<string>>();
 
 	readonly infer: TablesFromSchemas<T> = undefined as any;
 
@@ -82,7 +83,7 @@ export class Database<T extends SchemaRecord> {
 
 		this.kysely = new Kysely<TablesFromSchemas<T>>({
 			dialect: new BunSqliteDialect({ database: this.sqlite }),
-			plugins: [new DeserializePlugin(this.columns, validation), new ParseJSONResultsPlugin()],
+			plugins: [new DeserializePlugin(this.columns, this.tableColumns, validation), new ParseJSONResultsPlugin()],
 		});
 	}
 
@@ -264,6 +265,8 @@ export class Database<T extends SchemaRecord> {
 	}
 
 	private registerColumns(tableName: string, props: Prop[]) {
+		this.tableColumns.set(tableName, new Set(props.map((p) => p.key)));
+
 		const colMap = new Map<string, ColumnCoercion>();
 
 		for (const prop of props) {

package/src/plugin.ts

@@ -1,28 +1,41 @@
+import { type } from "arktype";
+import type { Type } from "arktype";
 import {
 	type KyselyPlugin,
+	type OperationNode,
 	type RootOperationNode,
 	type UnknownRow,
 	type QueryId,
+	AggregateFunctionNode,
 	TableNode,
 	AliasNode,
 	ValuesNode,
 	ValueNode,
 	ColumnNode,
+	IdentifierNode,
+	ReferenceNode,
+	ParensNode,
+	CastNode,
+	SelectQueryNode,
 } from "kysely";
-import { type } from "arktype";
-import type { Type } from "arktype";
+
 import { JsonParseError } from "./errors";
-import { JsonValidationError } from "./validation-error";
 import type { JsonValidation } from "./types";
+import { JsonValidationError } from "./validation-error";
 
 export type ColumnCoercion = "boolean" | "date" | { type: "json"; schema: Type };
 export type ColumnsMap = Map<string, Map<string, ColumnCoercion>>;
 
+type ResolvedCoercion = { table: string; coercion: ColumnCoercion };
+
+const typePreservingAggregateFunctions = new Set(["max", "min"]);
+
 export class DeserializePlugin implements KyselyPlugin {
 	private queryNodes = new WeakMap<QueryId, RootOperationNode>();
 
 	constructor(
 		private columns: ColumnsMap,
+		private tableColumns: Map<string, Set<string>>,
 		private validation: Required<JsonValidation>,
 	) {}
 
@@ -149,48 +162,231 @@ export class DeserializePlugin implements KyselyPlugin {
 		}
 	}
 
+	private coerceSingle(table: string, row: UnknownRow, col: string, coercion: ColumnCoercion) {
+		if (coercion === "boolean") {
+			if (typeof row[col] === "number") {
+				row[col] = row[col] === 1;
+			}
+
+			return;
+		}
+
+		if (coercion === "date") {
+			if (typeof row[col] === "number") {
+				row[col] = new Date(row[col] * 1000);
+			}
+
+			return;
+		}
+
+		if (typeof row[col] !== "string") {
+			return;
+		}
+
+		const value = row[col];
+
+		let parsed: unknown;
+
+		try {
+			parsed = JSON.parse(value);
+		} catch (e) {
+			throw new JsonParseError(table, col, value, e);
+		}
+
+		if (this.validation.onRead) {
+			this.validateJsonValue(table, col, parsed, coercion.schema);
+		}
+
+		row[col] = parsed;
+	}
+
 	private coerceRow(table: string, row: UnknownRow, cols: Map<string, ColumnCoercion>) {
 		for (const [col, coercion] of cols) {
 			if (!(col in row)) {
 				continue;
 			}
 
-			if (coercion === "boolean") {
-				if (typeof row[col] === "number") {
-					row[col] = row[col] === 1;
-				}
+			this.coerceSingle(table, row, col, coercion);
+		}
+	}
 
-				continue;
+	private getIdentifierName(node: OperationNode | undefined) {
+		if (!node || !IdentifierNode.is(node)) {
+			return null;
+		}
+
+		return node.name;
+	}
+
+	private addTableScopeEntry(scope: Map<string, string>, node: OperationNode) {
+		if (AliasNode.is(node) && TableNode.is(node.node)) {
+			const alias = this.getIdentifierName(node.alias);
+			const table = node.node.table.identifier.name;
+
+			scope.set(table, table);
+
+			if (alias) {
+				scope.set(alias, table);
 			}
 
-			if (coercion === "date") {
-				if (typeof row[col] === "number") {
-					row[col] = new Date(row[col] * 1000);
-				}
+			return;
+		}
+
+		if (TableNode.is(node)) {
+			const table = node.table.identifier.name;
+			scope.set(table, table);
+		}
+	}
+
+	private getTableScope(node: SelectQueryNode) {
+		const scope = new Map<string, string>();
+
+		for (const fromNode of node.from?.froms ?? []) {
+			this.addTableScopeEntry(scope, fromNode);
+		}
+
+		for (const join of node.joins ?? []) {
+			this.addTableScopeEntry(scope, join.table);
+		}
+
+		return scope;
+	}
+
+	private resolveReferenceCoercion(node: ReferenceNode | ColumnNode, scope: Map<string, string>) {
+		const column = ColumnNode.is(node)
+			? node.column.name
+			: ColumnNode.is(node.column)
+				? node.column.column.name
+				: null;
+
+		if (!column) {
+			return null;
+		}
+
+		if (ReferenceNode.is(node) && node.table) {
+			const tableRef = node.table.table.identifier.name;
+			const table = scope.get(tableRef) ?? tableRef;
+			const coercion = this.columns.get(table)?.get(column);
+
+			if (!coercion) {
+				return null;
+			}
+
+			return { table, coercion } satisfies ResolvedCoercion;
+		}
+
+		let match: ResolvedCoercion | null = null;
+		const resolvedTables = new Set<string>();
 
+		for (const table of scope.values()) {
+			if (resolvedTables.has(table)) {
 				continue;
 			}
 
-			if (typeof row[col] !== "string") {
+			resolvedTables.add(table);
+
+			const coercion = this.columns.get(table)?.get(column);
+
+			if (!coercion) {
 				continue;
 			}
 
-			const value = row[col];
+			if (match) {
+				return null;
+			}
+
+			match = { table, coercion };
+		}
+
+		return match;
+	}
+
+	private resolveSelectionCoercion(node: OperationNode, scope: Map<string, string>): ResolvedCoercion | null {
+		if (AliasNode.is(node)) {
+			return this.resolveSelectionCoercion(node.node, scope);
+		}
+
+		if (ReferenceNode.is(node) || ColumnNode.is(node)) {
+			return this.resolveReferenceCoercion(node, scope);
+		}
 
-			let parsed: unknown;
+		if (SelectQueryNode.is(node)) {
+			return this.resolveScalarSubqueryCoercion(node);
+		}
 
-			try {
-				parsed = JSON.parse(value);
-			} catch (e) {
-				throw new JsonParseError(table, col, value, e);
+		if (AggregateFunctionNode.is(node)) {
+			if (
+				node.aggregated.length !== 1 ||
+				!typePreservingAggregateFunctions.has(node.func.toLowerCase())
+			) {
+				return null;
 			}
 
-			if (this.validation.onRead) {
-				this.validateJsonValue(table, col, parsed, coercion.schema);
+			return this.resolveSelectionCoercion(node.aggregated[0]!, scope);
+		}
+
+		if (ParensNode.is(node)) {
+			return this.resolveSelectionCoercion(node.node, scope);
+		}
+
+		if (CastNode.is(node)) {
+			return null;
+		}
+
+		return null;
+	}
+
+	private resolveScalarSubqueryCoercion(node: SelectQueryNode) {
+		if (!node.selections || node.selections.length !== 1) {
+			return null;
+		}
+
+		return this.resolveSelectionCoercion(
+			node.selections[0]!.selection,
+			this.getTableScope(node),
+		);
+	}
+
+	private getSelectionOutputName(node: OperationNode) {
+		if (AliasNode.is(node)) {
+			return this.getIdentifierName(node.alias);
+		}
+
+		if (ReferenceNode.is(node) && ColumnNode.is(node.column)) {
+			return node.column.column.name;
+		}
+
+		if (ColumnNode.is(node)) {
+			return node.column.name;
+		}
+
+		return null;
+	}
+
+	private getSelectCoercions(node: RootOperationNode) {
+		const result = new Map<string, ResolvedCoercion>();
+
+		if (node.kind !== "SelectQueryNode" || !node.selections) {
+			return result;
+		}
+
+		const scope = this.getTableScope(node);
+
+		for (const selectionNode of node.selections) {
+			const output = this.getSelectionOutputName(selectionNode.selection);
+
+			if (!output) {
+				continue;
 			}
 
-			row[col] = parsed;
+			const resolved = this.resolveSelectionCoercion(selectionNode.selection, scope);
+
+			if (resolved) {
+				result.set(output, resolved);
+			}
 		}
+
+		return result;
 	}
 
 	transformResult: KyselyPlugin["transformResult"] = async (args) => {
@@ -206,14 +402,40 @@ export class DeserializePlugin implements KyselyPlugin {
 			return args.result;
 		}
 
-		const cols = this.columns.get(table);
-
-		if (!cols) {
-			return args.result;
-		}
+		const mainCols = this.columns.get(table);
+		const mainTableColumns = this.tableColumns.get(table);
+		const selectCoercions = this.getSelectCoercions(node);
 
 		for (const row of args.result.rows) {
-			this.coerceRow(table, row, cols);
+			if (mainCols) {
+				this.coerceRow(table, row, mainCols);
+			}
+
+			for (const col of Object.keys(row)) {
+				const resolved = selectCoercions.get(col);
+
+				if (resolved) {
+					this.coerceSingle(resolved.table, row, col, resolved.coercion);
+					continue;
+				}
+
+				if (mainTableColumns?.has(col)) {
+					continue;
+				}
+
+				for (const [otherTable, otherCols] of this.columns) {
+					if (otherTable === table) {
+						continue;
+					}
+
+					const coercion = otherCols.get(col);
+
+					if (coercion) {
+						this.coerceSingle(otherTable, row, col, coercion);
+						break;
+					}
+				}
+			}
 		}
 
 		return { ...args.result };