yizhi-postgres/src/database.ts

import { Pool, PoolClient, QueryResult, QueryResultRow } from "pg";
import { BasicEntity } from "./entity";
import { DeleteBuilder, IDeleteBuilder, IInsertBuilder, InsertBuilder, ISelectBuilder, IUpdateBuilder, SelectBuilder, UpdateBuilder } from "./query";
import { Class } from "./types";
import { formatSQL } from "./util";

interface IQueryResult<R extends QueryResultRow = any> extends QueryResult<R> {
	sql: string;
}

/**
 * 数据库基本操作
 */
export interface IDatabase {
	/**
	 * 执行sql语句
	 * @param sql sql语句
	 * @param args sql参数
	 */
	query<R extends QueryResultRow = any>(sql: string, args?: any[] | Record<string, any>): Promise<IQueryResult<R>>

	/**
	 * 查询实体
	 * @param Entity 实体
	 * @param alias 别名 
	 */
	select<E extends BasicEntity>(Entity: Class<E>, alias?: string): ISelectBuilder<E>

	/**
	 * 进行数据库插入
	 * @param Entity 实体
	 */
	insert<E extends BasicEntity>(Entity: Class<E>): IInsertBuilder<E>

	/**
	 * 进行实体更新
	 * @param Entity 实体
	 */
	update<E extends BasicEntity>(Entity: Class<E>): IUpdateBuilder<E>

	/**
	 * 进行实体删除
	 * @param Entity 实体
	 */
	delete<E extends BasicEntity>(Entity: Class<E>): IDeleteBuilder<E>
}

interface IPostgresClient extends IDatabase {
	/** 开启事务 */
	trans(): Promise<void>;
}

interface IPostgresDatabase extends IDatabase {
	/**
	 * 配置数据库
	 * @param loader 配置加载器
	 */
	config(loader: () => IPostgresConfig): void;
	/**
	 * 连接数据库
	 * @param callback 回调
	 */
	connect<R>(callback: (client: IPostgresClient) => R | Promise<R>): Promise<R>
	/**
	 * 执行事务操作
	 * @param callback 回调
	 */
	transaction<R>(callback: (client: IPostgresClient) => R | Promise<R>): Promise<R>
	/** 关闭数据库 */
	close(): void;
}

interface IPostgresConfig {
	host?: string
	port?: number
	user?: string
	password?: string
	database?: string
	max?: number
}

class PostgresClient implements IPostgresClient {
	#client: PoolClient
	#transOn = false;

	constructor(client: PoolClient) {
		this.#client = client;
	}

	public async query<R extends QueryResultRow = any>(sql: string, args?: any[] | Record<string, any>): Promise<IQueryResult<R>> {
		sql = args ? formatSQL(sql, args) : sql;
		try {
			const res = await this.#client.query(sql);
			return { ...res, sql };
		} catch (err: any) {
			err.sql = sql;
			throw err;
		}
	}


	public async trans() {
		await this.query("begin")
		this.#transOn = true;
	}

	public get transOn() { return this.#transOn; }

	public release() { this.#client.release(); }
	public select<E extends BasicEntity>(Entity: Class<E>, alias?: string): ISelectBuilder<E> { return new SelectBuilder(this.query.bind(this), Entity, alias); }
	public insert<E extends BasicEntity>(Entity: Class<E>): IInsertBuilder<E> { return new InsertBuilder(this.query.bind(this), Entity); }
	public update<E extends BasicEntity>(Entity: Class<E>): IUpdateBuilder<E> { return new UpdateBuilder(this.query.bind(this), Entity); }
	public delete<E extends BasicEntity>(Entity: Class<E>): IDeleteBuilder<E> { return new DeleteBuilder(this.query.bind(this), Entity); }
};


let pool: Pool | null = null;
let confLoader!: () => IPostgresConfig
function getPool() {
	if (!pool) {
		const conf = confLoader();
		pool = new Pool({
			host: conf.host,
			port: conf.port,
			user: conf.user,
			password: conf.password,
			database: conf.database,
			max: conf.max,
			ssl: false,
		});
	}
	return pool;
}


export const database: IPostgresDatabase = {
	config(loader: () => IPostgresConfig) {
		confLoader = loader;
	},

	async connect(callback) {
		const conn = await getPool().connect();
		const client = new PostgresClient(conn);
		let ret: any;
		try {
			ret = await callback(client);
			if (client.transOn) await client.query("commit");
		}
		catch (err) {
			if (client.transOn) await client.query("rollback");
			throw err;
		}
		finally {
			client.release();
		}

		return ret;
	},

	async transaction(callback) {
		return this.connect(async (client) => {
			await client.trans();
			return callback(client);
		});
	},

	close() {
		if (pool) pool.end().catch((err) => console.error(err));
	},

	async query(sql, args) {
		sql = args ? formatSQL(sql, args) : sql;
		try {
			const res = await getPool().query(sql);
			return { ...res, sql };
		} catch (err: any) {
			err.sql = sql;
			throw err;
		}
	},
	select(Entity, alias) { return new SelectBuilder(this.query.bind(this), Entity, alias); },
	insert(Entity) { return new InsertBuilder(this.query.bind(this), Entity); },
	update(Entity) { return new UpdateBuilder(this.query.bind(this), Entity); },
	delete(Entity) { return new DeleteBuilder(this.query.bind(this), Entity); },
};