preparePostgresDB() — drizzle-orm Function Reference

Source Code

drizzle-kit/src/cli/connections.ts lines 28–534

export const preparePostgresDB = async (
	credentials: PostgresCredentials | {
		driver: 'pglite';
		client: PGlite;
	},
): Promise<
	DB & {
		packageName:
			| '@aws-sdk/client-rds-data'
			| 'pglite'
			| 'pg'
			| 'postgres'
			| '@vercel/postgres'
			| '@neondatabase/serverless';
		proxy: Proxy;
		transactionProxy: TransactionProxy;
		migrate: (config: string | MigrationConfig) => Promise<void>;
	}
> => {
	if ('driver' in credentials) {
		const { driver } = credentials;
		if (driver === 'aws-data-api') {
			assertPackages('@aws-sdk/client-rds-data');
			const { RDSDataClient, ExecuteStatementCommand, TypeHint } = await import(
				'@aws-sdk/client-rds-data'
			);
			const { AwsDataApiSession, drizzle } = await import(
				'drizzle-orm/aws-data-api/pg'
			);
			const { migrate } = await import('drizzle-orm/aws-data-api/pg/migrator');
			const { PgDialect } = await import('drizzle-orm/pg-core');

			const config: AwsDataApiSessionOptions = {
				database: credentials.database,
				resourceArn: credentials.resourceArn,
				secretArn: credentials.secretArn,
			};
			const rdsClient = new RDSDataClient();
			const session = new AwsDataApiSession(
				rdsClient,
				new PgDialect(),
				undefined,
				config,
				undefined,
			);

			const db = drizzle(rdsClient, config);
			const migrateFn = async (config: MigrationConfig) => {
				return migrate(db, config);
			};

			const query = async (sql: string, params: any[]) => {
				const prepared = session.prepareQuery(
					{ sql, params: params ?? [] },
					undefined,
					undefined,
					false,
				);
				const result = await prepared.all();
				return result as any[];
			};
			const proxy = async (params: ProxyParams) => {
				const prepared = session.prepareQuery<
					PreparedQueryConfig & {
						execute: AwsDataApiPgQueryResult<unknown>;
						values: AwsDataApiPgQueryResult<unknown[]>;
					}
				>(
					{
						sql: params.sql,
						params: params.params ?? [],
						typings: params.typings,
					},
					undefined,
					undefined,
					params.mode === 'array',
				);
				if (params.mode === 'array') {
					const result = await prepared.values();
					return result.rows;
				}