Merge pull request #64 from drizzle-team:dri-111-implement-neon-serve…

…rless-driver-support

Implement NeonDB serverless driver support

changelogs/drizzle-orm-pg/0.13.3.md

@@ -0,0 +1,4 @@
+# drizzle-orm-pg 0.13.3-beta.1
+
+- Implemented NeonDB serverless driver support.
+- (internal) Added `session.all()` and `session.values()` methods.

drizzle-orm-pg/package.json

@@ -1,6 +1,6 @@
 {
 	"name": "drizzle-orm-pg",
-	"version": "0.13.2",
+	"version": "0.13.3-beta.1",
 	"description": "Drizzle ORM package for PostgreSQL database",
 	"main": "index.js",
 	"types": "index.d.ts",
@@ -49,6 +49,7 @@
 		}
 	},
 	"devDependencies": {
+		"@neondatabase/serverless": "^0.1.13",
 		"@types/pg": "^8.6.5",
 		"drizzle-orm": "link:../drizzle-orm/src",
 		"pg": "^8.8.0"

drizzle-orm-pg/src/dialect.ts

@@ -18,11 +18,11 @@ export class PgDialect {
 		await session.execute(sql`CREATE SCHEMA IF NOT EXISTS "drizzle"`);
 		await session.execute(migrationTableCreate);
 
-		const dbMigrations = await session.execute<{ id: number; hash: string; created_at: string }>(
+		const dbMigrations = await session.all<{ id: number; hash: string; created_at: string }>(
 			sql`SELECT id, hash, created_at FROM "drizzle"."__drizzle_migrations" ORDER BY created_at DESC LIMIT 1`,
 		);
 
-		const lastDbMigration = dbMigrations.rows[0];
+		const lastDbMigration = dbMigrations[0];
 		await session.execute(sql`BEGIN`);
 
 		try {

drizzle-orm-pg/src/index.ts

@@ -1,12 +1,11 @@
 export * from './alias';
 export * from './checks';
 export * from './columns';
-export * from './connector';
 export * from './db';
 export * from './dialect';
-export * from './driver';
 export * from './foreign-keys';
 export * from './indexes';
+export * from './node-pg';
 export * from './operations';
 export * from './session';
 export * from './table';

drizzle-orm-pg/src/neondb-serverless/connector.ts

@@ -0,0 +1,37 @@
+import { Logger, MigrationConfig, readMigrationFiles } from 'drizzle-orm';
+import { PgDialect } from '~/dialect';
+import { PgSession } from '~/session';
+import { NeonDriver } from './driver';
+import { NeonClient, NeonSession } from './session';
+
+export interface PgConnectorOptions {
+	logger?: Logger;
+	dialect?: PgDialect;
+	driver?: NeonDriver;
+}
+
+export class NeonConnector {
+	dialect: PgDialect;
+	driver: NeonDriver;
+	private session: NeonSession | undefined;
+
+	constructor(client: NeonClient, options: PgConnectorOptions = {}) {
+		this.dialect = new PgDialect();
+		this.driver = new NeonDriver(client, this.dialect, { logger: options.logger });
+	}
+
+	private async getSession() {
+		return this.session ?? (this.session = await this.driver.connect());
+	}
+
+	async connect() {
+		const session = await this.getSession();
+		return this.dialect.createDB(session);
+	}
+
+	async migrate(config: string | MigrationConfig) {
+		const migrations = readMigrationFiles(config);
+		const session = await this.getSession();
+		await this.dialect.migrate(migrations, session);
+	}
+}

drizzle-orm-pg/src/neondb-serverless/driver.ts

@@ -0,0 +1,32 @@
+import { types } from '@neondatabase/serverless';
+import { Logger } from 'drizzle-orm';
+import { PgDialect } from '~/dialect';
+import { NeonClient, NeonSession } from './session';
+
+export interface NeonDriverOptions {
+	logger?: Logger;
+}
+
+export class NeonDriver {
+	constructor(
+		private client: NeonClient,
+		private dialect: PgDialect,
+		private options: NeonDriverOptions = {},
+	) {
+		this.initMappers();
+	}
+
+	async connect(): Promise<NeonSession> {
+		return new NeonSession(this.client, this.dialect, { logger: this.options.logger });
+	}
+
+	initMappers() {
+		types.setTypeParser(types.builtins.TIMESTAMPTZ, (val) => val);
+		types.setTypeParser(types.builtins.TIMESTAMP, (val) => val);
+		types.setTypeParser(types.builtins.DATE, (val) => val);
+	}
+}
+
+export function pg(client: NeonClient, options: NeonDriverOptions = {}) {
+	return new NeonDriver(client, new PgDialect(), options);
+}

drizzle-orm-pg/src/neondb-serverless/index.ts

@@ -0,0 +1,3 @@
+export * from './connector';
+export * from './driver';
+export * from './session';

drizzle-orm-pg/src/neondb-serverless/session.ts

@@ -0,0 +1,111 @@
+import {
+	Client,
+	Pool,
+	PoolClient,
+	QueryArrayConfig,
+	QueryConfig,
+	QueryResult,
+	QueryResultRow,
+} from '@neondatabase/serverless';
+import { Logger, NoopLogger } from 'drizzle-orm';
+import { fillPlaceholders, Query } from 'drizzle-orm/sql';
+import { mapResultRow } from 'drizzle-orm/utils';
+import { PgDialect } from '~/dialect';
+import { SelectFieldsOrdered } from '~/operations';
+import { PgSession, PreparedQuery, PreparedQueryConfig } from '~/session';
+
+export type NeonClient = Pool | PoolClient | Client;
+
+export class NeonPreparedQuery<T extends PreparedQueryConfig> extends PreparedQuery<T> {
+	private rawQuery: QueryConfig;
+	private query: QueryArrayConfig;
+
+	constructor(
+		private client: NeonClient,
+		queryString: string,
+		private params: unknown[],
+		private logger: Logger,
+		private fields: SelectFieldsOrdered | undefined,
+		name: string | undefined,
+	) {
+		super();
+		this.rawQuery = {
+			name,
+			text: queryString,
+		};
+		this.query = {
+			name,
+			text: queryString,
+			rowMode: 'array',
+		};
+	}
+
+	execute(placeholderValues: Record<string, unknown> | undefined = {}): Promise<T['execute']> {
+		const params = fillPlaceholders(this.params, placeholderValues);
+
+		this.logger.logQuery(this.rawQuery.text, params);
+
+		const { fields } = this;
+		if (!fields) {
+			return this.client.query(this.rawQuery, params);
+		}
+
+		const result = this.client.query(this.query, params);
+
+		return result.then((result) => result.rows.map((row) => mapResultRow<T['execute']>(fields, row)));
+	}
+
+	all(placeholderValues: Record<string, unknown> | undefined = {}): Promise<T['all']> {
+		const params = fillPlaceholders(this.params, placeholderValues);
+		this.logger.logQuery(this.rawQuery.text, params);
+		return this.client.query(this.rawQuery, params).then((result) => result.rows);
+	}
+
+	values(placeholderValues: Record<string, unknown> | undefined = {}): Promise<T['values']> {
+		const params = fillPlaceholders(this.params, placeholderValues);
+		this.logger.logQuery(this.rawQuery.text, params);
+		return this.client.query(this.query, params).then((result) => result.rows);
+	}
+}
+
+export interface NeonSessionOptions {
+	logger?: Logger;
+}
+
+export class NeonSession extends PgSession {
+	private logger: Logger;
+
+	constructor(
+		private client: NeonClient,
+		dialect: PgDialect,
+		options: NeonSessionOptions = {},
+	) {
+		super(dialect);
+		this.logger = options.logger ?? new NoopLogger();
+	}
+
+	prepareQuery<T extends PreparedQueryConfig = PreparedQueryConfig>(
+		query: Query,
+		fields: SelectFieldsOrdered | undefined,
+		name: string | undefined,
+	): PreparedQuery<T> {
+		return new NeonPreparedQuery(this.client, query.sql, query.params, this.logger, fields, name);
+	}
+
+	async query(query: string, params: unknown[]): Promise<QueryResult> {
+		this.logger.logQuery(query, params);
+		const result = await this.client.query({
+			rowMode: 'array',
+			text: query,
+			values: params,
+		});
+		return result;
+	}
+
+	async queryObjects<T extends QueryResultRow>(
+		query: string,
+		params: unknown[],
+	): Promise<QueryResult<T>> {
+		return this.client.query<T>(query, params);
+	}
+}

drizzle-orm-pg/src/connector.ts → drizzle-orm-pg/src/node-pg/connector.ts

@@ -1,22 +1,23 @@
 import { Logger, MigrationConfig, readMigrationFiles } from 'drizzle-orm';
 import { PgDialect } from '~/dialect';
-import { PgDriver } from './driver';
-import { PgClient, PgSession } from './session';
+import { PgSession } from '~/session';
+import { NodePgDriver } from './driver';
+import { NodePgClient } from './session';
 
 export interface PgConnectorOptions {
 	logger?: Logger;
 	dialect?: PgDialect;
-	driver?: PgDriver;
+	driver?: NodePgDriver;
 }
 
 export class PgConnector {
 	dialect: PgDialect;
-	driver: PgDriver;
+	driver: NodePgDriver;
 	private session: PgSession | undefined;
 
-	constructor(client: PgClient, options: PgConnectorOptions = {}) {
+	constructor(client: NodePgClient, options: PgConnectorOptions = {}) {
 		this.dialect = new PgDialect();
-		this.driver = new PgDriver(client, this.dialect, { logger: options.logger });
+		this.driver = new NodePgDriver(client, this.dialect, { logger: options.logger });
 	}
 
 	private async getSession() {

drizzle-orm-pg/src/driver.ts → drizzle-orm-pg/src/node-pg/driver.ts

@@ -1,22 +1,22 @@
 import { Logger } from 'drizzle-orm';
 import { types } from 'pg';
-import { PgDialect } from './dialect';
-import { NodePgSession, PgClient, PgSession } from './session';
+import { PgDialect } from '~/dialect';
+import { NodePgClient, NodePgSession } from './session';
 
 export interface PgDriverOptions {
 	logger?: Logger;
 }
 
-export class PgDriver {
+export class NodePgDriver {
 	constructor(
-		private client: PgClient,
+		private client: NodePgClient,
 		private dialect: PgDialect,
 		private options: PgDriverOptions = {},
 	) {
 		this.initMappers();
 	}
 
-	async connect(): Promise<PgSession> {
+	async connect(): Promise<NodePgSession> {
 		return new NodePgSession(this.client, this.dialect, { logger: this.options.logger });
 	}
 
@@ -26,3 +26,7 @@ export class PgDriver {
 		types.setTypeParser(types.builtins.DATE, (val) => val);
 	}
 }
+
+export function pg(client: NodePgClient, options: PgDriverOptions = {}) {
+	return new NodePgDriver(client, new PgDialect(), options);
+}

drizzle-orm-pg/src/node-pg/index.ts

@@ -0,0 +1,3 @@
+export * from './connector';
+export * from './driver';
+export * from './session';