Improve & fix a number of things with table queries

- Log all queries - Add SELECT ... LIMIT - Add Select.queryOne - Fix INSERT ... ON CONFLICT DO UPDATE - Add support for EXCLUDED vars - Fix Table.select() overload precedence - Fix duplicate vars being sent - Fix some vars not being sent as vars when they should be - Fix raw iso strings being combined with query text instead of as string - Clean up logging
fluff4me · Mar 8, 2024 · 1756670 · 1756670
1 parent 56260f9
commit 1756670
Show file tree

Hide file tree

Showing 9 changed files with 191 additions and 122 deletions.
diff --git a/src/History.ts b/src/History.ts
@@ -1,26 +1,10 @@
 import { DatabaseError, Pool, PoolClient } from "pg";
 import { StackUtil } from "./IStrongPG";
+import log, { color } from "./Log";
 import Migration, { MigrationVersion } from "./Migration";
 import { DatabaseSchema } from "./Schema";
 import Transaction from "./Transaction";
 
-let ansicolor: typeof import("ansicolor") | undefined;
-function color (color: keyof typeof import("ansicolor"), text: string) {
-	if (!ansicolor) {
-		try {
-			// eslint-disable-next-line @typescript-eslint/no-unsafe-assignment
-			ansicolor = require("ansicolor");
-			// eslint-disable-next-line no-empty
-		} catch { }
-
-		if (!ansicolor)
-			return text;
-	}
-
-	// eslint-disable-next-line @typescript-eslint/no-unsafe-member-access, @typescript-eslint/no-unsafe-call
-	return (ansicolor as any)[color](text) as string;
-}
-
 export class History<SCHEMA extends DatabaseSchema | null = null> {
 
 	private migrations: Migration<DatabaseSchema | null, DatabaseSchema>[] = [];
@@ -60,26 +44,26 @@ export class History<SCHEMA extends DatabaseSchema | null = null> {
 			return -1;
 
 		const targetVersion = commits[commits.length - 1].version!;
-		this.log(color("lightYellow", `Found migrations up to v${targetVersion}`));
+		log(color("lightYellow", `Found migrations up to v${targetVersion}`));
 
 		let migratedVersion: MigrationVersion | undefined;
 		let migratedCommitIndex: number | undefined;
 		let rolledBack = false;
 		for (let i = startCommitIndex + 1; i < commits.length; i++) {
 			const commit = commits[i];
-			this.log(`Beginning migration ${commit.version!} ${commit.file ? color("lightBlue", commit.file) : ""}`);
+			log(`Beginning migration ${commit.version!} ${commit.file ? color("lightBlue", commit.file) : ""}`);
 
 			const statements = commit.compile();
 			if (!statements.length) {
-				this.log("Migration contains no statements");
+				log("Migration contains no statements");
 				continue;
 			}
 
 			let stack: StackUtil.Stack | undefined;
 			try {
 				await Transaction.execute(pool, async client => {
 					for (const statement of statements) {
-						this.log("  >", color("darkGray", statement.text));
+						log("  > ", color("darkGray", statement.text));
 						stack = statement.stack;
 						await client.query(statement);
 					}
@@ -90,7 +74,7 @@ export class History<SCHEMA extends DatabaseSchema | null = null> {
 			} catch (e) {
 				const err = e as DatabaseError;
 				const formattedStack = stack?.format();
-				this.log([
+				log([
 					`${color("lightRed", `Encountered an error: ${err.message[0].toUpperCase()}${err.message.slice(1)}`)}`,
 					err.hint ? `\n  ${err.hint}` : "",
 					formattedStack ? `\n${formattedStack}` : "",
@@ -104,44 +88,15 @@ export class History<SCHEMA extends DatabaseSchema | null = null> {
 		const version = commits[commitIndex].version!;
 
 		if (migratedVersion === undefined && !rolledBack) {
-			this.log(color("lightGreen", `Already on v${version}, no migrations necessary`));
+			log(color("lightGreen", `Already on v${version}, no migrations necessary`));
 			return startCommitIndex;
 		}
 
 		if (migratedVersion !== undefined)
 			await pool.query("INSERT INTO migrations VALUES ($1, $2)", [startCommitIndex, migratedCommitIndex]);
 
-		this.log(color(rolledBack ? "lightYellow" : "lightGreen", `${rolledBack ? "Rolled back" : "Migrated"} to v${version}`));
+		log(color(rolledBack ? "lightYellow" : "lightGreen", `${rolledBack ? "Rolled back" : "Migrated"} to v${version}`));
 
 		return commitIndex;
 	}
-
-	private log (text: string): void;
-	private log (prefix: string, text: string): void;
-	private log (prefix: string, text?: string) {
-		if (!process.env.DEBUG_PG)
-			return;
-
-		if (text === undefined)
-			text = prefix, prefix = "";
-
-		prefix = prefix ? prefix.slice(0, 20).trimEnd() + " " : prefix; // cap prefix length at 20
-
-		const maxLineLength = 150 - prefix.length;
-		text = text.split("\n")
-			.flatMap(line => {
-				const lines = [];
-				while (line.length > maxLineLength) {
-					lines.push(line.slice(0, maxLineLength));
-					line = line.slice(maxLineLength);
-				}
-				lines.push(line.trimEnd());
-				return lines;
-			})
-			.filter(line => line)
-			.map((line, i) => i ? line.padStart(line.length + prefix.length, " ") : `${prefix}${line}`)
-			.join("\n");
-
-		console.log(text);
-	}
 }
diff --git a/src/Log.ts b/src/Log.ts
@@ -0,0 +1,47 @@
+function log (text: string): void;
+function log (prefix: string, text: string): void;
+function log (prefix: string, text?: string) {
+	if (!process.env.DEBUG_PG)
+		return;
+
+	if (text === undefined)
+		text = prefix, prefix = "";
+
+	// prefix = prefix ? prefix.slice(0, 20).trimEnd() + " " : prefix; // cap prefix length at 20
+
+	const maxLineLength = 150 - prefix.length;
+	text = text.split("\n")
+		.flatMap(line => {
+			const lines = [];
+			while (line.length > maxLineLength) {
+				lines.push(line.slice(0, maxLineLength));
+				line = line.slice(maxLineLength);
+			}
+			lines.push(line.trimEnd());
+			return lines;
+		})
+		.filter(line => line)
+		.map((line, i) => i ? line.padStart(line.length + prefix.length, " ") : `${prefix}${line}`)
+		.join("\n");
+
+	console.log(text);
+}
+
+export default log;
+
+let ansicolor: typeof import("ansicolor") | undefined;
+export function color (color: keyof typeof import("ansicolor"), text: string) {
+	if (!ansicolor) {
+		try {
+			// eslint-disable-next-line @typescript-eslint/no-unsafe-assignment
+			ansicolor = require("ansicolor");
+			// eslint-disable-next-line no-empty
+		} catch { }
+
+		if (!ansicolor)
+			return text;
+	}
+
+	// eslint-disable-next-line @typescript-eslint/no-unsafe-member-access, @typescript-eslint/no-unsafe-call
+	return (ansicolor as any)[color](text) as string;
+}
diff --git a/src/Schema.ts b/src/Schema.ts
@@ -133,6 +133,26 @@ class Schema {
 	public static primaryKey<KEYS extends string[]> (...keys: KEYS): KEYS[number][] {
 		return keys;
 	}
+
+	public static getSingleColumnPrimaryKey<SCHEMA extends TableSchema> (schema: SCHEMA) {
+		const primaryKey = schema["PRIMARY_KEY"] as Schema.Column<SCHEMA>[] | undefined;
+		// eslint-disable-next-line @typescript-eslint/no-unsafe-assignment, @typescript-eslint/no-unsafe-member-access
+		// const primaryKey = ?.[0];
+		if (!primaryKey || primaryKey.length !== 1)
+			throw new Error("No primary key or primary key is multiple columns");
+
+		return primaryKey[0];
+	}
+
+	public static isColumn<SCHEMA extends TableSchema> (schema: SCHEMA, column: keyof SCHEMA, type: TypeString) {
+		const columnType = schema[column] as TypeString;
+		switch (type) {
+			case "TIMESTAMP":
+				return columnType.startsWith("TIMESTAMP");
+			default:
+				return columnType === type;
+		}
+	}
 }
 
 export default Schema;
@@ -154,7 +174,7 @@ namespace Schema {
 		keyof { [COLUMN in keyof SCHEMA as COLUMN extends keyof SpecialKeys<any> ? never : SCHEMA[COLUMN] extends Vaguify<TYPE> ? COLUMN : never]: SCHEMA[COLUMN] };
 	export type Columns<SCHEMA> = { [COLUMN in keyof SCHEMA as COLUMN extends keyof SpecialKeys<any> ? never : COLUMN]: SCHEMA[COLUMN] };
 	export type RowOutput<SCHEMA> = { [COLUMN in keyof SCHEMA as COLUMN extends keyof SpecialKeys<any> ? never : COLUMN]: OutputTypeFromString<Extract<SCHEMA[COLUMN], TypeString>> };
-	export type RowInput<SCHEMA> = { [COLUMN in keyof SCHEMA as COLUMN extends keyof SpecialKeys<any> ? never : COLUMN]: InputTypeFromString<Extract<SCHEMA[COLUMN], TypeString>> };
+	export type RowInput<SCHEMA, VARS = {}> = { [COLUMN in keyof SCHEMA as COLUMN extends keyof SpecialKeys<any> ? never : COLUMN]: InputTypeFromString<Extract<SCHEMA[COLUMN], TypeString>, VARS> };
 
 	type Vaguify<T> = T extends TypeStringMap[DataTypeID.BIGINT] ? TypeStringMap[DataTypeID.BIGINT] | TypeStringMap[DataTypeID.BIGSERIAL]
 		: T;

diff --git a/src/Table.ts b/src/Table.ts
@@ -12,22 +12,25 @@ export default class Table<SCHEMA extends TableSchema> {
 	 * SELECT *
 	 */
 	public select (): SelectFromTable<SCHEMA, "*"[]>;
+	/**
+	 * SELECT columns
+	 */
+	public select<COLUMNS extends Schema.Column<SCHEMA>[]> (...columns: COLUMNS): SelectFromTable<SCHEMA, COLUMNS>;
 	/**
 	 * SELECT *
 	 * ...then provide an initialiser for tweaking the query
 	 */
 	public select<RETURN extends SelectFromTable<SCHEMA, "*"[], any> = SelectFromTable<SCHEMA, "*"[]>> (initialiser: Initialiser<SelectFromTable<SCHEMA, "*"[]>, RETURN>): RETURN;
-	/**
-	 * SELECT columns
-	 */
-	public select<COLUMNS extends Schema.Column<SCHEMA>[]> (...columns: COLUMNS): SelectFromTable<SCHEMA, COLUMNS>;
 	/**
 	 * SELECT columns
 	 * ...then provide an initialiser for tweaking the query
 	 */
 	public select<COLUMNS extends Schema.Column<SCHEMA>[], RETURN extends SelectFromTable<SCHEMA, COLUMNS, any>> (...columnsAndInitialiser: [...COLUMNS, Initialiser<SelectFromTable<SCHEMA, COLUMNS>, RETURN>]): RETURN;
 	public select (...params: (Schema.Column<SCHEMA> | "*" | Initialiser<SelectFromTable<SCHEMA>> | Initialiser<SelectFromTable<SCHEMA, "*"[]>>)[]): SelectFromTable<SCHEMA, Schema.Column<SCHEMA>[]> | SelectFromTable<SCHEMA, "*"[]> {
 		const initialiser = typeof params[params.length - 1] === "function" ? params.pop() as Initialiser<SelectFromTable<SCHEMA>> : undefined;
+		if (params.length === 0)
+			params.push("*");
+
 		const query = new SelectFromTable<SCHEMA>(this.name, this.schema, params as Schema.Column<SCHEMA>[]);
 		// eslint-disable-next-line @typescript-eslint/no-unsafe-return
 		return initialiser?.(query) ?? query;

diff --git a/src/expressions/Expression.ts b/src/expressions/Expression.ts
@@ -33,21 +33,26 @@ export type ImplementableExpression = { [KEY in keyof ExpressionValues | keyof E
 
 export default class Expression<VARS = never> implements ImplementableExpression {
 
-	public static stringifyValue<VARS = never> (value: ExpressionOr<VARS, ValidType>, vars?: any[], enableStringConcatenation = false) {
-		let result: string;
+	public static stringifyValue<VARS = never> (value: ExpressionOr<VARS, ValidType>, vars: any[], enableStringConcatenation = false) {
 		if (typeof value === "function") {
 			const expr = new Expression(vars, enableStringConcatenation);
 			value(expr as any as ExpressionValues<VARS, null, null>);
-			result = `(${expr.compile()})`;
-		} else if (typeof value === "string" && !enableStringConcatenation) {
-			vars ??= [];
-			vars.push(value);
-			result = `$${vars.length}`;
-		} else {
-			result = Expression.stringifyValueRaw(value);
+			return `(${expr.compile()})`;
 		}
 
-		return result;
+		const shouldPassAsVariable = false
+			|| (typeof value === "string" && !enableStringConcatenation)
+			|| (value && typeof value === "object" && !(value instanceof Date) && !(value instanceof RegExp));
+		if (!shouldPassAsVariable)
+			return Expression.stringifyValueRaw(value);
+
+		const index = vars.indexOf(value);
+		if (index !== undefined && index !== -1)
+			// already in vars
+			return `$${index + 1}`;
+
+		vars.push(value);
+		return `$${vars.length}`;
 	}
 
 	/**
@@ -73,23 +78,23 @@ export default class Expression<VARS = never> implements ImplementableExpression
 				else if (value instanceof RegExp)
 					return `'${value.source.replace(/'/g, "''")}'`;
 				else
-					return value.toISOString();
+					return `'${value.toISOString()}'`;
 
 			case "number":
 				return `${value}`;
 		}
 	}
 
 	public static compile (initialiser: ExpressionInitialiser<any, any>, enableStringConcatenation = false, vars?: any[]) {
-		const expr = new Expression(vars, enableStringConcatenation);
+		const expr = new Expression(vars ?? [], enableStringConcatenation);
 		// eslint-disable-next-line @typescript-eslint/no-unsafe-argument
 		initialiser(expr as any);
 		return new Statement.Queryable(expr.compile(), undefined, expr.vars);
 	}
 
 	public readonly parts: (() => string)[] = [];
 
-	private constructor (public vars?: any[], private readonly enableStringConcatenation = false) { }
+	private constructor (public vars: any[], private readonly enableStringConcatenation = false) { }
 
 	public compile () {
 		return this.parts.map(part => part()).join("");

diff --git a/src/statements/Insert.ts b/src/statements/Insert.ts
@@ -1,5 +1,5 @@
 import { QueryResult } from "pg";
-import { Initialiser, InputTypeFromString, Value } from "../IStrongPG";
+import { Initialiser, InputTypeFromString, ValidType, Value } from "../IStrongPG";
 import Schema, { TableSchema } from "../Schema";
 import Expression from "../expressions/Expression";
 import Statement from "./Statement";
@@ -9,17 +9,24 @@ export interface InsertIntoTableFactory<SCHEMA extends TableSchema, COLUMNS exte
 	values (...values: { [I in keyof COLUMNS]: InputTypeFromString<SCHEMA[COLUMNS[I]]> }): InsertIntoTable<SCHEMA, COLUMNS>;
 }
 
-export default class InsertIntoTable<SCHEMA extends TableSchema, RESULT = []> extends Statement<RESULT> {
+export interface InsertIntoTableConflictActionFactory<SCHEMA extends TableSchema, COLUMNS extends Schema.Column<SCHEMA>[] = Schema.Column<SCHEMA>[], RESULT = []> {
+	doNothing (): InsertIntoTable<SCHEMA, COLUMNS, RESULT>;
+	doUpdate (initialiser: Initialiser<UpdateTable<SCHEMA, any, { [KEY in COLUMNS[number]as `EXCLUDED.${KEY & string}`]: SCHEMA[KEY] }>>): InsertIntoTable<SCHEMA, COLUMNS, RESULT>;
+}
+
+export default class InsertIntoTable<SCHEMA extends TableSchema, COLUMNS extends Schema.Column<SCHEMA>[] = Schema.Column<SCHEMA>[], RESULT = []> extends Statement<RESULT> {
 
 	public static columns<SCHEMA extends TableSchema, COLUMNS extends Schema.Column<SCHEMA>[] = Schema.Column<SCHEMA>[]> (tableName: string, schema: SCHEMA, columns: COLUMNS, isUpsert = false): InsertIntoTableFactory<SCHEMA, COLUMNS> {
+		const primaryKey = !isUpsert ? undefined : Schema.getSingleColumnPrimaryKey(schema);
+
 		return {
 			values: (...values: any[]) => {
 				const query = new InsertIntoTable<SCHEMA, COLUMNS>(tableName, schema, columns, values as never);
 				if (isUpsert) {
-					query.onConflictDoUpdate(update => {
+					query.onConflict(primaryKey!).doUpdate(update => {
 						for (let i = 0; i < columns.length; i++) {
-							// eslint-disable-next-line @typescript-eslint/no-unsafe-argument
-							update.set(columns[i], values[i]);
+							// eslint-disable-next-line @typescript-eslint/no-unsafe-argument, @typescript-eslint/no-unsafe-return, @typescript-eslint/restrict-template-expressions, @typescript-eslint/no-unsafe-member-access, @typescript-eslint/no-unsafe-call
+							update.set(columns[i], ((expr: any) => expr.var(`EXCLUDED.${String(columns[i])}`)) as never);
 						}
 					});
 				}
@@ -34,30 +41,42 @@ export default class InsertIntoTable<SCHEMA extends TableSchema, RESULT = []> ex
 		super();
 	}
 
-	private onConflict?: null | UpdateTable<SCHEMA, any>;
-	public onConflictDoNothing () {
-		this.onConflict = null;
-		return this;
-	}
-
-	public onConflictDoUpdate (initialiser: Initialiser<UpdateTable<SCHEMA, any>>) {
-		this.onConflict = new UpdateTable(undefined, this.schema, this.vars);
-		initialiser(this.onConflict);
-		return this;
+	private conflictTarget?: Schema.Column<SCHEMA>[];
+	private conflictAction?: null | UpdateTable<SCHEMA, any>;
+	public onConflict (...columns: Schema.Column<SCHEMA>[]): InsertIntoTableConflictActionFactory<SCHEMA, COLUMNS, RESULT> {
+		this.conflictTarget = columns;
+		return {
+			doNothing: () => {
+				this.conflictAction = null;
+				return this;
+			},
+			doUpdate: initialiser => {
+				this.conflictAction = new UpdateTable(undefined, this.schema, this.vars);
+				initialiser(this.conflictAction);
+				return this;
+			},
+		}
 	}
 
 	public compile () {
-		const values = this.values.map(value => Expression.stringifyValue(value, this.vars)).join(",");
-		let onConflict = this.onConflict === undefined ? " "
-			: this.onConflict === null ? "ON CONFLICT DO NOTHING"
+		const values = this.values.map((value: ValidType, i) => {
+			const column = this.columns[i];
+			if (Schema.isColumn(this.schema, column, "TIMESTAMP") && typeof value === "number")
+				value = new Date(value);
+			return Expression.stringifyValue(value, this.vars);
+		}).join(",");
+
+		const conflictTarget = this.conflictTarget?.length ? `(${this.conflictTarget.join(",")})` : "";
+		let conflictAction = this.conflictAction === undefined ? " "
+			: this.conflictAction === null ? `ON CONFLICT ${conflictTarget} DO NOTHING`
 				: undefined;
 
-		if (this.onConflict) {
-			const compiled = this.onConflict.compile()[0];
-			onConflict = `ON CONFLICT DO ${compiled.text}`;
+		if (this.conflictAction) {
+			const compiled = this.conflictAction.compile()[0];
+			conflictAction = `ON CONFLICT ${conflictTarget} DO ${compiled.text}`;
 		}
 
-		return this.queryable(`INSERT INTO ${this.tableName} (${this.columns.join(",")}) VALUES (${values}) ${onConflict!}`, undefined, this.vars);
+		return this.queryable(`INSERT INTO ${this.tableName} (${this.columns.join(",")}) VALUES (${values}) ${conflictAction!}`, undefined, this.vars);
 	}
 
 	protected override resolveQueryOutput (output: QueryResult<any>) {