Add OpenTelemetry integration for Postgres

kansson · May 27, 2023 · 28c28bc · 28c28bc
1 parent 1f1c16b
commit 28c28bc
Show file tree

Hide file tree

Showing 24 changed files with 456 additions and 166 deletions.
diff --git a/drizzle-orm/package.json b/drizzle-orm/package.json
@@ -47,10 +47,12 @@
     "@cloudflare/workers-types": ">=3",
     "@libsql/client": "*",
     "@neondatabase/serverless": ">=0.1",
+    "@opentelemetry/api": "^1.4.1",
     "@planetscale/database": ">=1",
     "@types/better-sqlite3": "*",
     "@types/pg": "*",
     "@types/sql.js": "*",
+    "@vercel/postgres": "*",
     "better-sqlite3": ">=7",
     "bun-types": "*",
     "knex": "*",
@@ -59,8 +61,7 @@
     "pg": ">=8",
     "postgres": ">=3",
     "sql.js": ">=1",
-    "sqlite3": ">=5",
-    "@vercel/postgres": "*"
+    "sqlite3": ">=5"
   },
   "peerDependenciesMeta": {
     "mysql2": {
@@ -116,16 +117,21 @@
     },
     "@libsql/client": {
       "optional": true
+    },
+    "@opentelemetry/api": {
+      "optional": true
     }
   },
   "devDependencies": {
     "@aws-sdk/client-rds-data": "^3.303.0",
     "@cloudflare/workers-types": "^4.20230321.0",
     "@libsql/client": "^0.1.1",
     "@neondatabase/serverless": "^0.2.9",
+    "@opentelemetry/api": "^1.4.1",
     "@originjs/vite-plugin-commonjs": "^1.0.3",
     "@planetscale/database": "^1.5.0",
     "@rollup/plugin-json": "^6.0.0",
+    "@rollup/plugin-replace": "^5.0.2",
     "@rollup/plugin-typescript": "^11.1.0",
     "@types/better-sqlite3": "^7.6.2",
     "@types/node": "^18.15.11",
@@ -134,6 +140,7 @@
     "@vercel/postgres": "^0.1.2",
     "better-sqlite3": "^8.0.0",
     "bun-types": "^0.5.8",
+    "concurrently": "^8.0.1",
     "knex": "^2.4.2",
     "kysely": "^0.24.2",
     "mysql2": "^3.2.0",

diff --git a/drizzle-orm/rollup.config.ts → drizzle-orm/rollup.cjs.config.ts b/drizzle-orm/rollup.config.ts → drizzle-orm/rollup.cjs.config.ts
@@ -1,4 +1,5 @@
 import json from '@rollup/plugin-json';
+import replace from '@rollup/plugin-replace';
 import typescript from '@rollup/plugin-typescript';
 import { defineConfig } from 'rollup';
 import { entries, external } from './rollup.common';
@@ -11,29 +12,25 @@ export default defineConfig([
 			acc[to] = from;
 			return acc;
 		}, {}),
-		output: [
-			{
-				format: 'esm',
-				dir: 'dist.new',
-				entryFileNames: '[name].mjs',
-				chunkFileNames: '[name]-[hash].mjs',
-				sourcemap: true,
-			},
-			{
-				format: 'cjs',
-				dir: 'dist.new',
-				entryFileNames: '[name].cjs',
-				chunkFileNames: '[name]-[hash].cjs',
-				sourcemap: true,
-			},
-		],
+		output: {
+			format: 'cjs',
+			dir: 'dist.new',
+			entryFileNames: '[name].cjs',
+			chunkFileNames: '[name]-[hash].cjs',
+			sourcemap: true,
+		},
 		external,
 		plugins: [
+			replace({
+				'await import': 'require',
+				preventAssignment: true,
+			}),
 			json({
 				preferConst: true,
 			}),
 			typescript({
-				tsconfig: 'tsconfig.build.json',
+				tsconfig: 'tsconfig.cjs.json',
+				outputToFilesystem: true,
 			}),
 		],
 	},

diff --git a/drizzle-orm/rollup.common.ts b/drizzle-orm/rollup.common.ts
@@ -51,4 +51,6 @@ export const external = [
 	'postgres',
 	'sqlite3',
 	'bun:sqlite',
+	'@opentelemetry/api',
+	'@vercel/postgres',
 ];
diff --git a/drizzle-orm/rollup.esm.config.ts b/drizzle-orm/rollup.esm.config.ts
@@ -0,0 +1,32 @@
+import json from '@rollup/plugin-json';
+import typescript from '@rollup/plugin-typescript';
+import { defineConfig } from 'rollup';
+import { entries, external } from './rollup.common';
+
+export default defineConfig([
+	{
+		input: entries.reduce<Record<string, string>>((acc, entry) => {
+			const from = 'src/' + entry + '.ts';
+			const to = entry;
+			acc[to] = from;
+			return acc;
+		}, {}),
+		output: {
+			format: 'esm',
+			dir: 'dist.new',
+			entryFileNames: '[name].mjs',
+			chunkFileNames: '[name]-[hash].mjs',
+			sourcemap: true,
+		},
+		external,
+		plugins: [
+			json({
+				preferConst: true,
+			}),
+			typescript({
+				tsconfig: 'tsconfig.esm.json',
+				outputToFilesystem: true,
+			}),
+		],
+	},
+]);
diff --git a/drizzle-orm/scripts/build.ts b/drizzle-orm/scripts/build.ts
@@ -1,5 +1,6 @@
 #!/usr/bin/env -S pnpm tsx
 import 'zx/globals';
+import concurrently from 'concurrently';
 import { entries } from '../rollup.common';
 
 function updateAndCopyPackageJson() {
@@ -11,7 +12,12 @@ function updateAndCopyPackageJson() {
 			const importEntry = `./${entry}.mjs`;
 			const requireEntry = `./${entry}.cjs`;
 			const typesEntry = `./${entry}.d.ts`;
-			acc[exportsEntry] = { import: importEntry, require: requireEntry, default: importEntry, types: typesEntry };
+			acc[exportsEntry] = {
+				types: typesEntry,
+				import: importEntry,
+				require: requireEntry,
+				default: importEntry,
+			};
 			return acc;
 		},
 		{},
@@ -20,11 +26,23 @@ function updateAndCopyPackageJson() {
 	fs.writeJSONSync('dist.new/package.json', pkg, { spaces: 2 });
 }
 
-await $`rollup --config rollup.config.ts --configPlugin typescript`;
-fs.copySync('../README.md', 'dist/README.md');
+await concurrently([
+	{
+		command: 'rollup --config rollup.cjs.config.ts --configPlugin typescript',
+		name: 'cjs',
+	},
+	{
+		command: 'rollup --config rollup.esm.config.ts --configPlugin typescript',
+		name: 'esm',
+	},
+	{
+		command: `tsc -p tsconfig.esm.json --declaration --outDir dist-dts --emitDeclarationOnly &&
+resolve-tspaths --out dist-dts &&
+rollup --config rollup.dts.config.ts --configPlugin typescript`,
+		name: 'dts',
+	},
+]).result;
+fs.copySync('../README.md', 'dist.new/README.md');
 updateAndCopyPackageJson();
-await $`tsc -p tsconfig.build.json --declaration --outDir dist-dts --emitDeclarationOnly`;
-await $`resolve-tspaths --out dist-dts`;
-await $`rollup --config rollup.dts.config.ts --configPlugin typescript`;
 fs.removeSync('dist');
 fs.renameSync('dist.new', 'dist');
diff --git a/drizzle-orm/src/node-postgres/session.ts b/drizzle-orm/src/node-postgres/session.ts
@@ -8,6 +8,7 @@ import type { PgTransactionConfig, PreparedQueryConfig, QueryResultHKT } from '~
 import { PgSession, PreparedQuery } from '~/pg-core/session';
 import { type RelationalSchemaConfig, type TablesRelationalConfig } from '~/relations';
 import { fillPlaceholders, type Query, sql } from '~/sql';
+import { tracer } from '~/tracing';
 import { type Assume, mapResultRow } from '~/utils';
 
 const { Pool } = pg;
@@ -40,32 +41,53 @@ export class NodePgPreparedQuery<T extends PreparedQueryConfig> extends Prepared
 	}
 
 	async execute(placeholderValues: Record<string, unknown> | undefined = {}): Promise<T['execute']> {
-		const params = fillPlaceholders(this.params, placeholderValues);
-
-		this.logger.logQuery(this.rawQuery.text, params);
-
-		const { fields, rawQuery, client, query, joinsNotNullableMap, customResultMapper } = this;
-		if (!fields && !customResultMapper) {
-			return client.query(rawQuery, params);
-		}
-
-		const result = await client.query(query, params);
+		return tracer.startActiveSpan('drizzle.execute', async () => {
+			const params = fillPlaceholders(this.params, placeholderValues);
+
+			this.logger.logQuery(this.rawQuery.text, params);
+
+			const { fields, rawQuery, client, query, joinsNotNullableMap, customResultMapper } = this;
+			if (!fields && !customResultMapper) {
+				return tracer.startActiveSpan('drizzle.driver.execute', async (span) => {
+					span?.setAttributes({
+						'drizzle.query.name': rawQuery.name,
+						'drizzle.query.text': rawQuery.text,
+						'drizzle.query.params': JSON.stringify(params),
+					});
+					return client.query(rawQuery, params);
+				});
+			}
 
-		return customResultMapper
-			? customResultMapper(result.rows)
-			: result.rows.map((row) => mapResultRow<T['execute']>(fields!, row, joinsNotNullableMap));
+			const result = await tracer.startActiveSpan('drizzle.driver.execute', (span) => {
+				span?.setAttributes({
+					'drizzle.query.name': query.name,
+					'drizzle.query.text': query.text,
+					'drizzle.query.params': JSON.stringify(params),
+				});
+				return client.query(query, params);
+			});
+
+			return tracer.startActiveSpan('drizzle.mapResponse', () => {
+				return customResultMapper
+					? customResultMapper(result.rows)
+					: result.rows.map((row) => mapResultRow<T['execute']>(fields!, row, joinsNotNullableMap));
+			});
+		});
 	}
 
 	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);
+		return tracer.startActiveSpan('drizzle.execute', () => {
+			const params = fillPlaceholders(this.params, placeholderValues);
+			this.logger.logQuery(this.rawQuery.text, params);
+			return tracer.startActiveSpan('drizzle.driver.execute', (span) => {
+				span?.setAttributes({
+					'drizzle.query.name': this.rawQuery.name,
+					'drizzle.query.text': this.rawQuery.text,
+					'drizzle.query.params': JSON.stringify(params),
+				});
+				return this.client.query(this.rawQuery, params).then((result) => result.rows);
+			});
+		});
 	}
 }
 
@@ -98,23 +120,6 @@ export class NodePgSession<
 		return new NodePgPreparedQuery(this.client, query.sql, query.params, this.logger, fields, name, customResultMapper);
 	}
 
-	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);
-	}
-
 	override async transaction<T>(
 		transaction: (tx: NodePgTransaction<TFullSchema, TSchema>) => Promise<T>,
 		config?: PgTransactionConfig | undefined,

diff --git a/drizzle-orm/src/pg-core/query-builders/delete.ts b/drizzle-orm/src/pg-core/query-builders/delete.ts
@@ -5,6 +5,7 @@ import type { SelectResultFields } from '~/query-builders/select.types';
 import { QueryPromise } from '~/query-promise';
 import type { Query, SQL, SQLWrapper } from '~/sql';
 import { type InferModel, Table } from '~/table';
+import { tracer } from '~/tracing';
 import { orderSelectedFields, type Simplify } from '~/utils';
 import type { SelectedFieldsFlat, SelectedFieldsOrdered } from './select.types';
 
@@ -69,7 +70,13 @@ export class PgDelete<
 			execute: TReturning extends undefined ? QueryResultKind<TQueryResult, never> : TReturning[];
 		}
 	> {
-		return this.session.prepareQuery(this.dialect.sqlToQuery(this.getSQL()), this.config.returning, name);
+		return tracer.startActiveSpan('drizzle.prepareQuery', () => {
+			return this.session.prepareQuery<
+				PreparedQueryConfig & {
+					execute: TReturning extends undefined ? QueryResultKind<TQueryResult, never> : TReturning[];
+				}
+			>(this.dialect.sqlToQuery(this.getSQL()), this.config.returning, name);
+		});
 	}
 
 	prepare(name: string): PreparedQuery<
@@ -81,6 +88,8 @@ export class PgDelete<
 	}
 
 	override execute: ReturnType<this['prepare']>['execute'] = (placeholderValues) => {
-		return this._prepare().execute(placeholderValues);
+		return tracer.startActiveSpan('drizzle.operation', () => {
+			return this._prepare().execute(placeholderValues);
+		});
 	};
 }
diff --git a/drizzle-orm/src/pg-core/query-builders/insert.ts b/drizzle-orm/src/pg-core/query-builders/insert.ts
@@ -7,6 +7,7 @@ import { QueryPromise } from '~/query-promise';
 import type { Placeholder, Query, SQLWrapper } from '~/sql';
 import { Param, SQL, sql } from '~/sql';
 import { type InferModel, Table } from '~/table';
+import { tracer } from '~/tracing';
 import type { Simplify } from '~/utils';
 import { mapUpdateSet, orderSelectedFields } from '~/utils';
 import type { SelectedFieldsFlat, SelectedFieldsOrdered } from './select.types';
@@ -143,7 +144,13 @@ export class PgInsert<
 			execute: TReturning extends undefined ? QueryResultKind<TQueryResult, never> : TReturning[];
 		}
 	> {
-		return this.session.prepareQuery(this.dialect.sqlToQuery(this.getSQL()), this.config.returning, name);
+		return tracer.startActiveSpan('drizzle.prepareQuery', () => {
+			return this.session.prepareQuery<
+				PreparedQueryConfig & {
+					execute: TReturning extends undefined ? QueryResultKind<TQueryResult, never> : TReturning[];
+				}
+			>(this.dialect.sqlToQuery(this.getSQL()), this.config.returning, name);
+		});
 	}
 
 	prepare(name: string): PreparedQuery<
@@ -155,6 +162,8 @@ export class PgInsert<
 	}
 
 	override execute: ReturnType<this['prepare']>['execute'] = (placeholderValues) => {
-		return this._prepare().execute(placeholderValues);
+		return tracer.startActiveSpan('drizzle.operation', () => {
+			return this._prepare().execute(placeholderValues);
+		});
 	};
 }