R524809
60 lines
2.0 KiB
TypeScript
60 lines
2.0 KiB
TypeScript
import { Pool, type PoolClient, type QueryResultRow } from 'pg';
|
|
|
|
export type DatabaseTransaction = Pick<PoolClient, 'query'>;
|
|
|
|
export class Database {
|
|
public constructor(private readonly pool: Pool) {}
|
|
|
|
public async transaction<T>(
|
|
work: (transaction: DatabaseTransaction) => Promise<T>,
|
|
): Promise<T> {
|
|
const client = await this.pool.connect();
|
|
try {
|
|
await client.query('BEGIN');
|
|
// atlas_app_login 是 NOINHERIT 登录角色；应用 DML 必须降为最小权限的 atlas_app。
|
|
await client.query('SET LOCAL ROLE atlas_app');
|
|
const result = await work(client);
|
|
await client.query('COMMIT');
|
|
return result;
|
|
} catch (error) {
|
|
await client.query('ROLLBACK');
|
|
throw error;
|
|
} finally {
|
|
client.release();
|
|
}
|
|
}
|
|
|
|
/** 每个带工作空间边界的业务读写都必须通过此方法执行。 */
|
|
public async workspaceTransaction<T>(
|
|
workspaceId: string,
|
|
work: (transaction: DatabaseTransaction) => Promise<T>,
|
|
): Promise<T> {
|
|
return this.transaction(async (transaction) => {
|
|
// SET LOCAL 仅作用于当前事务，连接归还池后不会泄漏租户 ID。
|
|
await transaction.query(
|
|
"SELECT set_config('app.workspace_id', $1, true)",
|
|
[workspaceId],
|
|
);
|
|
return work(transaction);
|
|
});
|
|
}
|
|
|
|
public async query<Row extends QueryResultRow = QueryResultRow>(
|
|
text: string,
|
|
values?: unknown[],
|
|
) {
|
|
// 即使不带租户上下文的查询也使用受限角色，调用方不会误用登录角色权限。
|
|
return this.transaction((transaction) =>
|
|
transaction.query<Row>(text, values),
|
|
);
|
|
}
|
|
|
|
public async close() {
|
|
await this.pool.end();
|
|
}
|
|
}
|
|
|
|
export function createApplicationDatabase(connectionString: string): Database {
|
|
return new Database(new Pool({ connectionString, max: 10 }));
|
|
}
|