From 86512224b0770f80d84cd7a8d4d1cccc3de8e1da Mon Sep 17 00:00:00 2001 From: Alex Holovach Date: Sun, 5 Oct 2025 06:58:02 -0500 Subject: [PATCH 1/5] fix(otel-drizzle) jsdoc --- packages/otel-drizzle/src/index.ts | 52 +++++++++++++++++++++--------- 1 file changed, 36 insertions(+), 16 deletions(-) diff --git a/packages/otel-drizzle/src/index.ts b/packages/otel-drizzle/src/index.ts index d6ed67c..c77239a 100644 --- a/packages/otel-drizzle/src/index.ts +++ b/packages/otel-drizzle/src/index.ts @@ -141,41 +141,61 @@ function finalizeSpan(span: Span, error?: unknown): void { } /** - * Instruments a Drizzle database client with OpenTelemetry tracing. - * - * This function wraps the client's `query` method to automatically create - * spans for each database operation. It supports both promise-based and - * callback-based query patterns. + * Instruments a database connection pool with OpenTelemetry tracing. * + * This function wraps the connection pool's `query` method to automatically create + * spans for each database operation. * The instrumentation is idempotent - calling it multiple times on the same - * client will only instrument it once. + * pool will only instrument it once. * - * @typeParam TClient - The type of the Drizzle client - * @param client - The Drizzle client instance to instrument + * @typeParam TClient - The type of the database connection pool or client + * @param client - The database connection pool or client to instrument (e.g., pg Pool, mysql2 Connection) * @param config - Optional configuration for instrumentation behavior - * @returns The instrumented client (same instance, modified in place) + * @returns The instrumented pool/client (same instance, modified in place) * * @example * ```typescript + * // PostgreSQL with node-postgres * import { drizzle } from 'drizzle-orm/node-postgres'; * import { Pool } from 'pg'; * import { instrumentDrizzle } from '@kubiks/otel-drizzle'; * * const pool = new Pool({ connectionString: process.env.DATABASE_URL }); - * const db = drizzle(pool); + * const instrumentedPool = instrumentDrizzle(pool); + * const db = drizzle(instrumentedPool); * - * // Instrument with defaults - * instrumentDrizzle(db); - * - * // Or with custom configuration - * instrumentDrizzle(db, { + * // With custom configuration + * const instrumentedPool = instrumentDrizzle(pool, { * dbSystem: 'postgresql', * dbName: 'myapp', * captureQueryText: true, - * maxQueryTextLength: 500, + * maxQueryTextLength: 1000, * peerName: 'db.example.com', * peerPort: 5432, * }); + * const db = drizzle(instrumentedPool); + * ``` + * + * @example + * ```typescript + * // MySQL with mysql2 + * import { drizzle } from 'drizzle-orm/mysql2'; + * import mysql from 'mysql2/promise'; + * import { instrumentDrizzle } from '@kubiks/otel-drizzle'; + * + * const connection = await mysql.createConnection(process.env.DATABASE_URL); + * const db = drizzle(instrumentDrizzle(connection, { dbSystem: 'mysql' })); + * ``` + * + * @example + * ```typescript + * // SQLite with better-sqlite3 + * import { drizzle } from 'drizzle-orm/better-sqlite3'; + * import Database from 'better-sqlite3'; + * import { instrumentDrizzle } from '@kubiks/otel-drizzle'; + * + * const sqlite = new Database('database.db'); + * const db = drizzle(instrumentDrizzle(sqlite, { dbSystem: 'sqlite' })); * ``` */ export function instrumentDrizzle( From 5edad4514e03922b227799a9b6e0d12b47367c86 Mon Sep 17 00:00:00 2001 From: Alex Holovach Date: Sun, 5 Oct 2025 07:21:12 -0500 Subject: [PATCH 2/5] add instrumentDrizzleClient function --- packages/otel-drizzle/README.md | 41 +++- .../kubiks-otel-drizzle-1.0.0.tgz | Bin 8755 -> 0 bytes packages/otel-drizzle/package.json | 5 +- packages/otel-drizzle/src/index.test.ts | 177 +++++++++++++++++- packages/otel-drizzle/src/index.ts | 170 +++++++++++++++++ pnpm-lock.yaml | 2 +- 6 files changed, 388 insertions(+), 7 deletions(-) delete mode 100644 packages/otel-drizzle/kubiks-otel-drizzle-1.0.0.tgz diff --git a/packages/otel-drizzle/README.md b/packages/otel-drizzle/README.md index e514d88..4e42826 100644 --- a/packages/otel-drizzle/README.md +++ b/packages/otel-drizzle/README.md @@ -42,7 +42,11 @@ Works with any observability platform that supports OpenTelemetry including: ## Usage -Simply wrap your database connection pool with `instrumentDrizzle()` before passing it to Drizzle: +There are two ways to instrument Drizzle ORM with OpenTelemetry: + +### Option 1: Instrument the Connection Pool (Recommended) + +Wrap your database connection pool with `instrumentDrizzle()` before passing it to Drizzle: ```typescript import { drizzle } from "drizzle-orm/node-postgres"; @@ -57,9 +61,32 @@ const db = drizzle(instrumentedPool); const users = await db.select().from(usersTable); ``` -### Optional Configuration +### Option 2: Instrument an Existing Drizzle Client + +If you already have a Drizzle database instance or don't have access to the underlying pool, use `instrumentDrizzleClient()`: ```typescript +import { drizzle } from "drizzle-orm/node-postgres"; +import { Pool } from "pg"; +import { instrumentDrizzleClient } from "@kubiks/otel-drizzle"; +import * as schema from "./schema"; + +const pool = new Pool({ connectionString: process.env.DATABASE_URL }); +const db = drizzle(pool, { schema }); + +// Instrument the existing database instance +instrumentDrizzleClient(db); + +// All queries are now traced automatically +const users = await db.select().from(schema.users); +``` + +### Optional Configuration + +Both instrumentation methods accept the same configuration options: + +```typescript +// Option 1: Instrument the pool const pool = new Pool({ connectionString: process.env.DATABASE_URL }); const instrumentedPool = instrumentDrizzle(pool, { dbSystem: "postgresql", // Database type (default: 'postgresql') @@ -70,6 +97,16 @@ const instrumentedPool = instrumentDrizzle(pool, { peerPort: 5432, // Database server port }); const db = drizzle(instrumentedPool); + +// Option 2: Instrument the Drizzle client +const db = drizzle(pool, { schema }); +instrumentDrizzleClient(db, { + dbSystem: "postgresql", + dbName: "myapp", + captureQueryText: true, + peerName: "db.example.com", + peerPort: 5432, +}); ``` ### Works with All Drizzle-Supported Databases diff --git a/packages/otel-drizzle/kubiks-otel-drizzle-1.0.0.tgz b/packages/otel-drizzle/kubiks-otel-drizzle-1.0.0.tgz deleted file mode 100644 index 30eb6c0eebd8b286e42256668f1f1ce1c0243d8b..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 8755 zcmV-3BFx<%iwFP!000001MNNScH6j)`&&=Jap$CVuc9PP+qpLxH&>RG)SC~dw$o0x z?QImBPBgV8SCaEFsekUjeT3bY+b3B7Bqd6gowU=Ldv>=zJ+4IpAP9mWK!MD}>woml z`PxCtagVz0zx)xO_4V~K7F>nvGMHL`V)r#Z({u)!cE~b zjT0{dR(>y?V(0#sd=6VZcHsAU5Oc$D!pTkKpI;=b+OM(4>yJ0tD<1Xv$S{t1H1^}z z4+G}M?1D%9?1r63UXbvC$%YZpMDADlCf^`Qv^ zAd?Fq6Nkg(%8NJv511FnVc++F9vg)H>6izJmtcECf5c-}O)faAbR|ZmMvV-(H!}Qy zVV$aCSAKF4P7@aKIEnl|(lD7H^heVH@}R0m{@9l-AUgFsHh}0fh8`fFCL4zXe~ABc z>egg>HuB>OlMQ?%eKt*?Fvda>po#rh3nLcu(Z~P_KJZO_&A1~NWPO4HO{9ooEV;S} z$N8T4u`!%R0W`@8XApwOsKsA-Kfw|NABLk*c!gc-hrz(d9>rUR(Ss`QEWG5@7cs0s zm;kGS1ss4$Hmb5JzVJpPcE+W^pj{BJr*$X7=3&kRiSLcrB#fxNVmB>l|AouC?cLsM zyW_G}mmPQ7uUbvF$trdio+~DM-Riw)pY#|&bnK(v8`j=s_R$-*-#TiV%>CuK<955O z-7#8+#|JGJ%34Rx!AY}qw8t8Nchm;mZGi>@(O#QjgOXIs?IN*5x8uBkN4wEFX!YKh z#%`;3gfw^C9cHs*yVGkqCkJ+i9iMcL+g%r$Zvxq)*3oVUT5%8Eqn-t=LK$;k!2|2Q zun!KfDZ@Si?mNgIbK1vmI<37IJ@%q~&~%}&;R0uNY7L8qJpyLD)?rhRDdxrEgQ zDjfsC1b6K93m1#9JsbWxy;l1OyWzBtdL4K)p_iRriu$_Mbxme>T3r;#Zl`@{8YoV{ zXcG~@J8}gT6gA6@2!P=8N!LvYvZiYv0Ie<}V;2>$W&E?v)c4OH^qI)@A1&IN9}M`l z^=tf%+qh%>e+p}VVg28Hyt(;({r@vQ{uoU-yQNk3q;t>=L zQlNgsVS(=OPg5UPZCth{2mF$coLq&%hUIfQ;u(Pt0TCCK6ZrWMmX#pJrCK1?%L!Hm zzBT5FX8}vKZHf2c1F{`>V?@pMU z!Ax$pv96=)&Uwb$PvB35i7I^JlGWI1R}4snmvHR2}hz`M5~Gtnymzho;qs zH63XBwUrhBhkAXwvT}rv<}EO~RfSggTmQYu1MtHl?`*`keyC3Bohk>6+UL&`E328+ zwj-W^o5GTrc_tYJ0@WM(SpnI3(@_E*8#HoLo#x7F)V-j7JKH|12IhWkTa4Qy(5-z^ zl*->h4&Y^uI1FWUv*u0w3LzcVD}VnuJ@Y@tYavvx4kG`z-$uM*+V#pLjFWTl|2~Z> z=5@XD;lm}lv2B2B3a}RFLPSm`A5_ePdS!5C#W!)n#})IKo>2Vc9#+g6+FQ$pOv>U4OMYH5#u&I1o2LvIM!9j5j?9-G-H_pX0_!#R=qT=wnPLY>Wq~h6ZmUi_7hrf24HV8B8 zJh}L}&Uccn(;shv1N0-2Ms*zkXO?B1X4s!lP_l6gpvGM=(@9n7+S)&=J3oh$=kKom zQd^r^{F?Vcv3IQVD4b5>9T>bFD+x~~Foup7^J*=l+!uWDohmxH=*jl6TW%DE(N1DT zydOpb_qxv~2o6-Y<_iOUB@n7z(F2%NH50nT6W}wM#?{+?IN)1>)tz{OC~}~{ayy-N zXI4X+cfz*!XGoI=1D*XJ1^X}_KP2gdm&5miHtnJA?e+Uit3Z~wN#yl;gvMuUT0b-g zXI;|xtx4U62P*f?zBd6B{t^lHK$*5a)_+(x$KG{XazG=w^;>;|Z$lE7f7$gIPAIf-OPhLA_D+-mbr|nWOq(=Tmi1+qw+>0b4hn`e^4=S_$jN+Yu

mCDn5tvmegp?ej?+H;v>b>^j(T7KVKgEfP&<)fVf6q^ zwK-Z@DHI=@qnhc@6}P~wnjcqI24IQ>w?Yze{iHfDe*>-0!|$4LZPuG3(>8lGb39jk zX^w?#y6JGB@?oa}3syxm8I}hEy#zU!QEJbJJ6SbG&DE}nD}&^$8k!gQ)UQTnSTlz! zD?=)pZG)Ls)84hb$z%j80YsZ7j;d$I?`yNZm-H{HVeK}#h{7u-NrbamZL7RVBq^HB zQS?$?VY!4zmcvYq!x8i&u2~R|2GtN2A|MaMYc)$Yi}dpCDvn~Eqqt_qz|kyS-e5lA zRFJw`|6nFFevJnHwuQb>*FCg*y-xQ-v+?1`K6K6FVrl!>?bsNPnY+cRZqM#fxM?2D zRlVs#mD$dfAGy5`$FAEEEnSw^9YeTi{)lycy4=b2Z?1BOeCex{W*Hzs-#&8^!o z0RI_x2@z<+TDR6k=_K+m!6$~&p*M=TC=+|6P8baRbEv@vpxEp8c^o4?0rV?SJR^jG zjF=?(HhB9zJRc7S>-dWDdZlZt#tq#ir2-8^)y@c_AaFN6}D zL!7J$B>+bjoBYSdU#*S5S?lccXZF1Qc-?viI;5!p4ElfveP7_J(#vL{LJ%#92&oC{ zZ&ESmVdBBZRq}jR&pe?M|9EU|Xy8H^ie2RPskO0DsGgh?CRA=%8!4`wD63&Krs{QT z^O^NT)$;gk>W>Dp24s-ds34W9HmxVtQ&j`uPUH=-d=nS;7B%vcB=tcmMCt_}u2byG>zw zDJmo{0Q?+bzqU$@y9vI8wU#qp?{xwWM)G8*-b#OoJrcYB?@HE+cNHY4-Zi`n7?V2z zCjM5NF6l$IQ$|I!x{F_6)CdBeX3Lt}eiyt8#9M);yKA2u^giHwRvA>YI--?$#Dthf zOQUo`%66YM(AH76*Ezum-)(-_J+Su(^p2hBcNGlYi2yOXw#K@AjPKt2Sb7OS^!6Ry z;(?cV=-e<*ZT!UV0x5OD!u+NJ@sRF(6SoU7(#;Enlx~nO$4hrM(#F(2`f|K<4CGF` z9Is^}2zYW8Mj!7Xf!y{9wUXPzU&2{#b0v?7;F3$yLE12NywT_kj4XgwIc^zmv1#xz z2(JPYGgQE@PYY_S{+!|VUxCD~+WQ3(xP=y-$0^R+_k<<4zkZgb&)H`MK(u^xLWVFw zRPf}Cd=T7$|B;KCMmt(t`R2B<;_aN?zTeq`$u{LsJ);46^aubSF%j`irB4)2vSCpH zTOpXU5N$1$bH-zEXG7D4BBa`NhM|og8&iqZRkUzRc@VIW?{PH^Y|wH2|vIWdRM7gie8RK9aBOO zXKG{q9(>W8OduxHbue2e2GbYk$t~~X>kFP-K#zco-27+5>1f2dFAs!bFtox&(z(8% zLg-BvUm`%`XJLpQ?Of}J-nBoTj+xlsCxwyN1Fdz7bQ3d>MqzM{%wBi_B{4}dSH8$w z0#1THW|8F;VgnatOi%Orpk_Jo%)ir8d)-m%20~p0L z2`aV3F^|B!D|xxs-gq+NTf(Gq@nrQwu)9d-U1j5M5R*QG=n{gc>JGVbo2wxw-y~3J-7Kxxfj&nRGY-sEyv6`CcyDlzhrK3|NJUcqiOr{bD89ectejrSmKV<5ar9xx%IvQ+& z0kcv7GA1|p6&M-5nWB**Ltt4>PrNME=)oG52@}G=8!&#N%Mz5JGzPOqV7b2~f$gMy zWV0H*)m_AW%6|!!R6LT#wm_MseA4)?opT5c{r*?53brfHDHMDuqky9GiEpNhY^o&$ zm8!^?B3{m*_Fg;Bs8&%JM?p|yTVp_%O$?^vhFy4+Vu_RK8hbJAqwO*oZeN1okvh!A zVEdqhU|7MB4nsv+co%YYB1MN)UhhGNUOoRT#AB)sJl|FXw4u?dEFXt-B#J00*_k#1 z^d%9}NmiWqUFcro{{{w{vP-2Ej&mjelSjRw%q@xXn#pM3XO&+0XS1uhatc)I8Dqefr~VMoG_H(1=kA!$NyV>9MH3%57%Vu6y8O zS||8eZU_G7N#CjoOR24V3#?*W(V3Ub600cm4HGKHP*l%)oJ+FEAji0=18`}=g`ROC z$q0BUC_-40h*6Z!t`EVb&>K<(%Fkvo1$Qus{4tF*>XdLPy%8ffdQL1uNr>KSBc&}7 z1PbsSeaWC#?5T!OeX0p*$y{=XHh+-(e5MEiDbvL`ZWJ!Zm`@`i617_1{t>R0F2X2Z zgq5Y1fMFA=n`{1l1ls#4HhvKjC7%?BEc4HC(3E9N{fs`v7GfOpZ*x#@p_?>xM5%{K z=;;VvMczcr6G7r(%uY!^Pr-U77a`@qdsA@yLGLkv?M5af5-qYLWMjAoDx0#C&Nqv- z5*AM<^xg*CQDAQ(!X|ntFhC_EN_`o!$zm}TClFNfNLnh%{9H(bXCoj)7*9gn_h73; z3g@8@GYMQ#9~KGJX-uRS#|~nQHmOl1HE?5TA3GzB*dZul)-=73I0k2S%z7e}gSL<# z3L2&3rv^q)7I+iF>FX$YuvF_sP@!1JVT)*N6$V#M$bK&L8_Kz(%jJy2GUOad;LIfeO6U0 z_Ha9ocno1niZwYeR?j6?q_Ab-r@$8DP3Coi)awil0X9TaGYB}%2VoPoP)YIKDa^%7 zt7-S_hTU~PoOBN4Y*nNNXFy1{Q$@P9tTAw4vPmmurhKUwi|!PAQ5~G+(c6(g>x;lU zH{uIPm}Gz+u`Ht90vV=3??W@Kr3^CmaD3xUCTa0pCK`pAs zK`Qd0sCSr_E)}m-@1P#H7i7IsJna#+%oNLVU763F9d(fyb?c;tW zwDGop^jE~R$kwt?%!SinlORddwu`0Iw7QdA5zsfe=%c!VCXIVzE zPPihk6r->lTbXy3awNTc=__D>hdUZ_o+&|m(iCcLYlC=Qf3I9WImT(Mpn2+SXGsW9 zV_P|yWwg*gl2awyxtIOkVOnBr<}ei!WEC{+VCku{a?@MNf`SNYuGqZN?MO?jRcygp z!6?6$71C>!Q<}N4S!QcVnq)fLYbj5t5oxE2Ui?COS+Zs(uUc{Fk+iTK_ohpRNgA9| zdUNJyWknl&w7BU2>#8fHEW>6D?UsblW{DM3yxb}0G*dNet+XG?5Z+HR`GnTgiY}y# zHuVcBQ#x#Fkmo2%Yejj`j(GE!8e^3&J(n#_4cBz_ck%@ z&7*DL2yQWCn(%?MznQ$!%I0Qq)AqP-fNdo)x! zLW{;o8=R$@aGzGnT2R}fa-T%`S&RHx={as^K+%oBn*wC%q6so}gFAj%=Z#SL@Xu$`nqPhrFKY-f@_ ze01!+GkW;cFyGitBx?$yG=du01!Bn1+MM(2!QI|);yhPiL= zT^FJMg2nbrg#J~7e%1hS?Oh_Y<7UOcvZ0Wzp1b9iDv}lx1%$M+mmCg_gfB-bRXN$TM{DNP5C}fa@ zgBr4LJ3l#vX8gu(ZZNPjdqm?!j*>e;sW$sC>^8>^?e8JcHp*wSfxNs#+!iPd3=`10 zF1F;+_Zu1u5dkH9tF{S!cff57IiC^IB%4t(ST^x-8jR`?JDEKZa z@NtVMptA2ZGbL!_&l8YY>$)Kd0+b921F$!Qx{UE*kEQ{#?EI_wnkG5Q^pXJ2uyvf1 zQ1n_9{eq~X2sp-$*lMosp_t8aTJ#B}iE7)GC5|nBl5PS->lhn9NrljRiKSf}^+S@` zF;=0PIDtcT=Z#mQe}}ev*?3JYJJ3&1BWmFV($=OEy~Z2_I>*lVPJm`cF4U;m2s`;{F?E@5tX zLt&afk#fAG`QW_4hIU69)IdzFb4M|~chUTXppG-3e-p-KOr#DG=-MGMhK)0nRUd`> znyi6CvJ=vJv(eD3((68AeO?4+&_FUiC-{YmLi05uj22>S(0pr&DYY-)(0;YZWD_;Atxl$mf_Nt(YkGD~(yM$ATI>5Ljj#5(Ch_fXX)<<+PwHl0*f8B1*?VMv z52=7FC3SwHZghnvqp~;y#I<)dj#6_84A}_MQOgTyR;!iuK5Wc&KVvUl1iD2g_mo(1 zB_U~(@bT-VQYf@r`@{{Hm4gi>GKKhx6be`)LxY-s-B9MmxoF(4Q$i*OBog~;CE}nR zqLFDUgR0wKvLMdW#<-zOnf($incG2|@5sYNm=Gj#X1$63L7fjWS>?lrh`aQNPgsfD zANadu&r&Z=cQvqZ&b!3I`3;a_X}XOu2d?7yVuj+OJTxz_!O$9OirZRPrm`9Ex1Mi+ z&vSbeDS~=PU$1NA+u79&tRi1R<4@hIC+hNl^5H4DxVX)Zg~HXQU*gCM5mj5R#I@W~ z62Dbxai2zVTi8~?Qni_(ejiI!5Ut>^Tu7JaRE7tlFxx&QRg(X!D4ZzivQe9Bx~NG~ zx$qMU3K6xG8vZAOMf^pvUFE$o`+hL^|IA07|40w>{Q+$(KmW1ubbaGV;rz$rjrH&M zKmIA72dwNpwPDe4)=ZaB1tCk*4kPajV<+VAmQG4t!5d6sT_7Bi3sf1XLl;FH$=vF;K%_4_L#C{k|&5n~Zlgt389GM`z66xNsgbrk@)0jv2 zs?odheB{B3Vfgwcj)~8q*qbhZhzO)Rzl(K~{tz-W1FoVqjI7 zma*tCe$CL|dQ-;4DlFV-ljw*F;>D8-8fEQvETp#}H)EdCqO%4@F@xg+=*VD7TBFe^ zIc8AE5zlzUT~p+3K?+DpQym`Ah{#vH87+J_kF*vgjHJZACb{2OnrTubeRq1kF_)J@ z(<-m6O@BD#5l*6XY(Ty66dwUS(6Z0}Xrr^pbYQa%7^N~V!+ zTf~f7tjEhoXqmhphy?ksk|Lc0dYp-mp4H8WRJSJc4>LJt<~B8B1Zs&yM<=Z&<17U{ z3LyZ%-$4q3`MObYdoa=aNcS!mUYJ5yu^viqOc2$rE(U`<>d@g_LHw$ zBhNU#Y+DtFtu`}(r;|gV%z?>Xivp>=3a9!|6kCI{hxgS*fiW)|QU0)OCX4H}SsPp} z)@|+=yAPJ5cxWz(GoeKM+=n(ILQj_>aiWyg4Snu@)r9}z6)(k{=%&J212KZYi^+7V z1JxSIOz%%s-1>l4(eLxzYDd~cYRd|ljW^*mDmoQxaegG!Enc9rMXl81p)>`_$1L?Q zcVU*odL`^IvsJD#<+L!m)0F&+veerLvu1W-5j@FrbnLV7c{Tg!zx5}80%@xQ*QZEr z@d;A^%@#5oCNi@4Ls*nkrpTqEj1=>SOeiB9pu|}dL>8w_@Yl&PsI!Z|notYBQiACj zZ6vsxKMY#Ul8p4DK7(4;44&)gv=602{`!U6aiL={|9;%a+CXK^NkDkO{xHohmNbh# zr6iF%v5kb5`)?(L^UeUuA$J6`*bdq*tsUK&D6 z;vk*OI7dv<_JVD|87|I&iOIm#2@;W$O^0=1#QGmaYyi+*o(S6^Ptu^%EvF;Y;@!gD zIoTH?M@#_(nNQ3o0V5;~6*-eUqIAXpD3@9;h@fGb`sac8D+LghjqtAxTygwH-Vgo2 zAA2Jf!qRlXy#XG65W0lNh^ExkDmwyH@EA}qnf}I5;SPs>pE97SM}28VN&DoMZ^~}r z7`iy3AnncRf-LLPIj`gHqv!@5zQtvt^f0=p_`8fj76fwyv}Tl_L^t#n?mmW|yYiTc zGw8JIOO~Qio8w4>Px+WXcV5oL!b9g8r6s-j((~p8)-eK18(MPEJlEzuPniqgMMup` z4w)Bt!&o!rZP~%{9CKfFs$8I^$G8h;%IQdX{zQ2;&hy8~%h7$x~?p8aL`_OQ|VUqaNo~4o2Hz zt)Z}h${u6Fo`+h_0m@dMoHY$$(%rd@&EqMe`kAP zIM7eF&>6NR`-4k%|Ni^G{)06466Sg|0>j%x{(J6u)>$@c?rh@YPHh>39>J7}i_zsFGTp8K4G9UNC`aGV1 z`keW&F5MUuPG~O@Fs8`%(y{UgkEZ?_%Zl7Jl1$tqNnm;Z-#%_B1lq30hSR@tb=c}r zL5Gj%uT33`-TxSwd8_mv_n)|9p(`t?q&|mR@_UOD}@UwYn@qFEH?XuYyx_zVI%%9J@<{_5#0O9k+i?jBj;GjB)23F)8EU=!O{Mo)^T(cE93!7~}Wc4x=Ogj>}<; zZ*w<{5`vslM)6`8<16ol$>GgsLj0avVI<4n=1!Q5Qt2yCYH7)ZFvi0BU^3igx4|sJ zr~QE?cfqL2{8cc unknown; @@ -296,3 +296,178 @@ describe("instrumentDrizzle", () => { expect(result).toBeNull(); }); }); + +describe("instrumentDrizzleClient", () => { + let provider: BasicTracerProvider; + let exporter: InMemorySpanExporter; + + beforeEach(() => { + exporter = new InMemorySpanExporter(); + provider = new BasicTracerProvider({ + spanProcessors: [new SimpleSpanProcessor(exporter)], + }); + trace.setGlobalTracerProvider(provider); + }); + + afterEach(async () => { + await provider.shutdown(); + exporter.reset(); + trace.disable(); + }); + + it("instruments a db with $client property", async () => { + const mockClient = { + query: vi.fn(() => Promise.resolve({ rows: [{ id: 1 }] })), + }; + + const mockDb = { + $client: mockClient, + select: vi.fn(), + }; + + const instrumented = instrumentDrizzleClient(mockDb); + expect(instrumented).toBe(mockDb); + + // Execute a query through the client + const result = await mockClient.query("SELECT * FROM users"); + expect(result).toEqual({ rows: [{ id: 1 }] }); + + const spans = exporter.getFinishedSpans(); + expect(spans).toHaveLength(1); + expect(spans[0]?.name).toBe("drizzle.select"); + expect(spans[0]?.attributes["db.statement"]).toBe("SELECT * FROM users"); + }); + + it("instruments a db with _.session.execute property", async () => { + const mockSession = { + execute: vi.fn(() => Promise.resolve({ rows: [{ id: 2 }] })), + }; + + const mockDb = { + _: { + session: mockSession, + }, + select: vi.fn(), + }; + + const instrumented = instrumentDrizzleClient(mockDb); + expect(instrumented).toBe(mockDb); + + // Execute a query through the session + const result = await mockSession.execute("INSERT INTO users (name) VALUES ('test')"); + expect(result).toEqual({ rows: [{ id: 2 }] }); + + const spans = exporter.getFinishedSpans(); + expect(spans).toHaveLength(1); + expect(spans[0]?.name).toBe("drizzle.insert"); + expect(spans[0]?.attributes["db.operation"]).toBe("INSERT"); + }); + + it("only instruments once when called multiple times", async () => { + const mockClient = { + query: vi.fn(() => Promise.resolve({ rows: [] })), + }; + + const mockDb = { + $client: mockClient, + }; + + const firstInstrumented = instrumentDrizzleClient(mockDb); + const wrappedQuery = mockClient.query; + + const secondInstrumented = instrumentDrizzleClient(mockDb); + + expect(firstInstrumented).toBe(mockDb); + expect(secondInstrumented).toBe(mockDb); + expect(mockClient.query).toBe(wrappedQuery); + }); + + it("respects custom configuration", async () => { + const mockClient = { + query: vi.fn(() => Promise.resolve({ rows: [] })), + }; + + const mockDb = { + $client: mockClient, + }; + + const config: InstrumentDrizzleConfig = { + dbSystem: "mysql", + dbName: "test_db", + peerName: "db.example.com", + peerPort: 3306, + }; + + instrumentDrizzleClient(mockDb, config); + + await mockClient.query("SELECT 1"); + + const spans = exporter.getFinishedSpans(); + expect(spans).toHaveLength(1); + + const span = spans[0]; + if (!span) { + throw new Error("Expected a recorded span"); + } + + expect(span.attributes["db.system"]).toBe("mysql"); + expect(span.attributes["db.name"]).toBe("test_db"); + expect(span.attributes["net.peer.name"]).toBe("db.example.com"); + expect(span.attributes["net.peer.port"]).toBe(3306); + }); + + it("returns db unchanged if db is null", () => { + const result = instrumentDrizzleClient(null as any); + expect(result).toBeNull(); + }); + + it("returns db unchanged if no instrumentable properties exist", () => { + const mockDb = { + select: vi.fn(), + // No $client or _.session properties + }; + + const result = instrumentDrizzleClient(mockDb); + expect(result).toBe(mockDb); + }); + + it("handles errors in session.execute", async () => { + const error = new Error("database error"); + const mockSession = { + execute: vi.fn(() => Promise.reject(error)), + }; + + const mockDb = { + _: { + session: mockSession, + }, + }; + + instrumentDrizzleClient(mockDb); + + await expect(mockSession.execute("DELETE FROM users")).rejects.toThrow(error); + + const spans = exporter.getFinishedSpans(); + expect(spans).toHaveLength(1); + expect(spans[0]?.status.code).toBe(SpanStatusCode.ERROR); + }); + + it("only instruments session once when called multiple times", () => { + const mockSession = { + execute: vi.fn(() => Promise.resolve({ rows: [] })), + }; + + const mockDb = { + _: { + session: mockSession, + }, + }; + + instrumentDrizzleClient(mockDb); + const wrappedExecute = mockSession.execute; + + instrumentDrizzleClient(mockDb); + + expect(mockSession.execute).toBe(wrappedExecute); + }); +}); diff --git a/packages/otel-drizzle/src/index.ts b/packages/otel-drizzle/src/index.ts index c77239a..5f2d996 100644 --- a/packages/otel-drizzle/src/index.ts +++ b/packages/otel-drizzle/src/index.ts @@ -316,3 +316,173 @@ export function instrumentDrizzle( return client; } + +/** + * Interface for Drizzle database instances with minimal type requirements. + */ +interface DrizzleDbLike { + $client?: DrizzleClientLike; + _?: { + session?: { + execute?: QueryFunction; + [INSTRUMENTED_FLAG]?: true; + }; + }; +} + +/** + * Instruments an already created Drizzle database instance with OpenTelemetry tracing. + * + * This function is useful when you have a Drizzle database instance created with + * `const db = drizzle(pool, { schema })` and want to add instrumentation to it + * without having direct access to the underlying pool. + * + * The instrumentation is idempotent - calling it multiple times on the same + * database will only instrument it once. + * + * @typeParam TDb - The type of the Drizzle database instance + * @param db - The Drizzle database instance to instrument + * @param config - Optional configuration for instrumentation behavior + * @returns The instrumented database instance (same instance, modified in place) + * + * @example + * ```typescript + * // When you have a pre-created Drizzle database instance + * import { drizzle } from 'drizzle-orm/node-postgres'; + * import { Pool } from 'pg'; + * import { instrumentDrizzleClient } from '@kubiks/otel-drizzle'; + * import * as schema from './schema'; + * + * const pool = new Pool({ connectionString: process.env.DATABASE_URL }); + * const db = drizzle(pool, { schema }); + * + * // Instrument the existing db instance + * instrumentDrizzleClient(db, { + * dbSystem: 'postgresql', + * dbName: 'myapp', + * captureQueryText: true, + * peerName: 'db.example.com', + * peerPort: 5432, + * }); + * + * // Now all queries through db are traced + * const users = await db.select().from(schema.users); + * ``` + * + * @example + * ```typescript + * // Works with any Drizzle driver + * import { drizzle } from 'drizzle-orm/mysql2'; + * import mysql from 'mysql2/promise'; + * import { instrumentDrizzleClient } from '@kubiks/otel-drizzle'; + * + * const connection = await mysql.createConnection(process.env.DATABASE_URL); + * const db = drizzle(connection); + * + * instrumentDrizzleClient(db, { dbSystem: 'mysql' }); + * ``` + */ +export function instrumentDrizzleClient( + db: TDb, + config?: InstrumentDrizzleConfig, +): TDb { + if (!db) { + return db; + } + + // Try to instrument via $client first (most common case) + if (db.$client && typeof db.$client.query === "function") { + instrumentDrizzle(db.$client, config); + return db; + } + + // Try to instrument via session.execute as fallback + if (db._ && db._.session && typeof db._.session.execute === "function") { + const session = db._.session; + + // Check if already instrumented + if (session[INSTRUMENTED_FLAG]) { + return db; + } + + const { + tracerName = DEFAULT_TRACER_NAME, + dbSystem = DEFAULT_DB_SYSTEM, + dbName, + captureQueryText = true, + maxQueryTextLength = 1000, + peerName, + peerPort, + } = config ?? {}; + + const tracer = trace.getTracer(tracerName); + const originalExecute = session.execute; + + if (!originalExecute) { + return db; + } + + const instrumentedExecute: QueryFunction = function instrumented( + this: unknown, + ...args: unknown[] + ) { + // Extract query information + const queryText = extractQueryText(args[0]); + const operation = queryText ? extractOperation(queryText) : undefined; + const spanName = operation + ? `drizzle.${operation.toLowerCase()}` + : "drizzle.query"; + + // Start span + const span = tracer.startSpan(spanName, { kind: SpanKind.CLIENT }); + span.setAttribute(SEMATTRS_DB_SYSTEM, dbSystem); + + if (operation) { + span.setAttribute(SEMATTRS_DB_OPERATION, operation); + } + + if (dbName) { + span.setAttribute(SEMATTRS_DB_NAME, dbName); + } + + if (captureQueryText && queryText !== undefined) { + const sanitized = sanitizeQueryText(queryText, maxQueryTextLength); + span.setAttribute(SEMATTRS_DB_STATEMENT, sanitized); + } + + if (peerName) { + span.setAttribute(SEMATTRS_NET_PEER_NAME, peerName); + } + + if (peerPort) { + span.setAttribute(SEMATTRS_NET_PEER_PORT, peerPort); + } + + const activeContext = trace.setSpan(context.active(), span); + + // Promise-based pattern (session.execute is typically promise-based) + return context.with(activeContext, () => { + try { + const result = originalExecute.apply(this, args); + return Promise.resolve(result) + .then((value) => { + finalizeSpan(span); + return value; + }) + .catch((error) => { + finalizeSpan(span, error); + throw error; + }); + } catch (error) { + finalizeSpan(span, error); + throw error; + } + }); + }; + + session[INSTRUMENTED_FLAG] = true; + session.execute = instrumentedExecute; + } + + return db; +} diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 9dfeebe..1234b90 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -66,7 +66,7 @@ importers: specifier: ^0.36.4 version: 0.36.4(@opentelemetry/api@1.9.0)(@types/pg@8.15.5)(@types/react@18.2.46)(kysely@0.28.7)(postgres@3.4.7)(react@18.2.0) postgres: - specifier: ^3.4.5 + specifier: ^3.4.7 version: 3.4.7 rimraf: specifier: 3.0.2 From e77c63390ad13dd1bc0e07540b4630afea764b51 Mon Sep 17 00:00:00 2001 From: Alex Holovach Date: Sun, 5 Oct 2025 08:03:25 -0500 Subject: [PATCH 3/5] trace transactions as well --- packages/otel-drizzle/README.md | 47 ++- packages/otel-drizzle/package.json | 2 +- packages/otel-drizzle/src/index.test.ts | 135 +++++++- packages/otel-drizzle/src/index.ts | 440 ++++++++++++++++++++++-- 4 files changed, 583 insertions(+), 41 deletions(-) diff --git a/packages/otel-drizzle/README.md b/packages/otel-drizzle/README.md index 4e42826..1e8de12 100644 --- a/packages/otel-drizzle/README.md +++ b/packages/otel-drizzle/README.md @@ -63,22 +63,27 @@ const users = await db.select().from(usersTable); ### Option 2: Instrument an Existing Drizzle Client -If you already have a Drizzle database instance or don't have access to the underlying pool, use `instrumentDrizzleClient()`: +If you already have a Drizzle database instance or don't have access to the underlying pool, use `instrumentDrizzleClient()`. This method instruments the database at the session level, capturing all query operations: ```typescript -import { drizzle } from "drizzle-orm/node-postgres"; -import { Pool } from "pg"; +// Works with postgres-js (Postgres.js) +import { drizzle } from "drizzle-orm/postgres-js"; import { instrumentDrizzleClient } from "@kubiks/otel-drizzle"; import * as schema from "./schema"; -const pool = new Pool({ connectionString: process.env.DATABASE_URL }); -const db = drizzle(pool, { schema }); +const db = drizzle(process.env.DATABASE_URL!, { schema }); // Instrument the existing database instance instrumentDrizzleClient(db); // All queries are now traced automatically const users = await db.select().from(schema.users); +// Direct execute calls are also traced +await db.execute("SELECT * FROM users"); +// Transactions are also traced +await db.transaction(async (tx) => { + await tx.insert(schema.users).values({ name: "John" }); +}); ``` ### Optional Configuration @@ -111,16 +116,26 @@ instrumentDrizzleClient(db, { ### Works with All Drizzle-Supported Databases -This package supports **all databases that Drizzle ORM supports**, including PostgreSQL, MySQL, SQLite, Turso, Neon, PlanetScale, and more. +This package automatically detects and instruments **all databases that Drizzle ORM supports**. It works by detecting whether your database driver uses a `query` or `execute` method and instrumenting it appropriately. This includes: + +- **PostgreSQL** (node-postgres, postgres.js, Neon, Vercel Postgres, etc.) +- **MySQL** (mysql2, PlanetScale, TiDB, etc.) +- **SQLite** (better-sqlite3, LibSQL/Turso, Cloudflare D1, etc.) +- **And any other Drizzle-supported database** ```typescript -// PostgreSQL with node-postgres +// PostgreSQL with postgres-js (Postgres.js) - use instrumentDrizzleClient +import { drizzle } from "drizzle-orm/postgres-js"; +const db = drizzle(process.env.DATABASE_URL!); +instrumentDrizzleClient(db); + +// PostgreSQL with node-postgres (pg) - use instrumentDrizzle on pool import { drizzle } from "drizzle-orm/node-postgres"; import { Pool } from "pg"; const pool = new Pool({ connectionString: process.env.DATABASE_URL }); const db = drizzle(instrumentDrizzle(pool)); -// MySQL with mysql2 +// MySQL with mysql2 (uses 'execute' or 'query' method) import { drizzle } from "drizzle-orm/mysql2"; import mysql from "mysql2/promise"; const connection = await mysql.createConnection(process.env.DATABASE_URL); @@ -131,6 +146,12 @@ import { drizzle } from "drizzle-orm/better-sqlite3"; import Database from "better-sqlite3"; const sqlite = new Database("database.db"); const db = drizzle(instrumentDrizzle(sqlite, { dbSystem: "sqlite" })); + +// LibSQL/Turso (uses 'execute' method) +import { drizzle } from "drizzle-orm/libsql"; +import { createClient } from "@libsql/client"; +const client = createClient({ url: "...", authToken: "..." }); +const db = drizzle(instrumentDrizzle(client, { dbSystem: "sqlite" })); ``` ## What You Get @@ -138,12 +159,20 @@ const db = drizzle(instrumentDrizzle(sqlite, { dbSystem: "sqlite" })); Each database query automatically creates a span with rich telemetry data: - **Span name**: `drizzle.select`, `drizzle.insert`, `drizzle.update`, etc. -- **Operation type**: `db.operation` attribute (SELECT, INSERT, UPDATE, DELETE) +- **Operation type**: `db.operation` attribute (SELECT, INSERT, UPDATE, DELETE, SET) - **SQL query text**: Full query statement captured in `db.statement` (configurable) - **Database system**: `db.system` attribute (postgresql, mysql, sqlite, etc.) +- **Transaction tracking**: Transaction queries are marked with `db.transaction` attribute - **Error tracking**: Exceptions are recorded with stack traces and proper span status - **Performance metrics**: Duration and timing information for every query +### Transaction Support + +All queries within transactions are automatically traced, including: +- RLS (Row Level Security) queries like `SET LOCAL role` and `set_config()` +- All nested transaction queries +- Transaction rollbacks and commits + ### Span Attributes The instrumentation adds the following attributes to each span following [OpenTelemetry semantic conventions](https://opentelemetry.io/docs/specs/semconv/database/): diff --git a/packages/otel-drizzle/package.json b/packages/otel-drizzle/package.json index ba77f5f..86273fa 100644 --- a/packages/otel-drizzle/package.json +++ b/packages/otel-drizzle/package.json @@ -1,6 +1,6 @@ { "name": "@kubiks/otel-drizzle", - "version": "2.0.4", + "version": "2.0.9", "private": false, "publishConfig": { "access": "public" diff --git a/packages/otel-drizzle/src/index.test.ts b/packages/otel-drizzle/src/index.test.ts index dee933d..2a8ea16 100644 --- a/packages/otel-drizzle/src/index.test.ts +++ b/packages/otel-drizzle/src/index.test.ts @@ -295,6 +295,56 @@ describe("instrumentDrizzle", () => { const result = instrumentDrizzle(null as any); expect(result).toBeNull(); }); + + it("instruments a client with execute method instead of query", async () => { + const client = { + execute: vi.fn(() => Promise.resolve({ rows: [{ id: 1 }] })), + }; + + instrumentDrizzle(client); + + // Execute with SQL object format (used by various drivers) + const result = await client.execute({ + sql: "SELECT * FROM users WHERE id = ?", + args: [1], + }); + + expect(result).toEqual({ rows: [{ id: 1 }] }); + + const spans = exporter.getFinishedSpans(); + expect(spans).toHaveLength(1); + + const span = spans[0]; + if (!span) { + throw new Error("Expected a recorded span"); + } + + expect(span.name).toBe("drizzle.select"); + expect(span.attributes["db.statement"]).toBe("SELECT * FROM users WHERE id = ?"); + expect(span.attributes["db.operation"]).toBe("SELECT"); + }); + + it("instruments a client with execute method using string query", async () => { + const client = { + execute: vi.fn(() => Promise.resolve({ rows: [] })), + }; + + instrumentDrizzle(client); + + await client.execute("DELETE FROM users"); + + const spans = exporter.getFinishedSpans(); + expect(spans).toHaveLength(1); + + const span = spans[0]; + if (!span) { + throw new Error("Expected a recorded span"); + } + + expect(span.name).toBe("drizzle.delete"); + expect(span.attributes["db.operation"]).toBe("DELETE"); + expect(span.attributes["db.statement"]).toBe("DELETE FROM users"); + }); }); describe("instrumentDrizzleClient", () => { @@ -315,7 +365,35 @@ describe("instrumentDrizzleClient", () => { trace.disable(); }); - it("instruments a db with $client property", async () => { + it("instruments a db with session.prepareQuery method", async () => { + const mockPreparedQuery = { + execute: vi.fn(() => Promise.resolve({ rows: [{ id: 1 }] })), + }; + + const mockSession = { + prepareQuery: vi.fn(() => mockPreparedQuery), + }; + + const mockDb = { + session: mockSession, + select: vi.fn(), + }; + + const instrumented = instrumentDrizzleClient(mockDb); + expect(instrumented).toBe(mockDb); + + // Simulate what happens when db.select().from() is called + const prepared = mockSession.prepareQuery({ sql: "SELECT * FROM users" }); + const result = await prepared.execute(); + expect(result).toEqual({ rows: [{ id: 1 }] }); + + const spans = exporter.getFinishedSpans(); + expect(spans).toHaveLength(1); + expect(spans[0]?.name).toBe("drizzle.select"); + expect(spans[0]?.attributes["db.statement"]).toBe("SELECT * FROM users"); + }); + + it("instruments a db with $client property as fallback", async () => { const mockClient = { query: vi.fn(() => Promise.resolve({ rows: [{ id: 1 }] })), }; @@ -323,6 +401,7 @@ describe("instrumentDrizzleClient", () => { const mockDb = { $client: mockClient, select: vi.fn(), + // No direct execute method }; const instrumented = instrumentDrizzleClient(mockDb); @@ -363,6 +442,26 @@ describe("instrumentDrizzleClient", () => { expect(spans[0]?.attributes["db.operation"]).toBe("INSERT"); }); + it("instruments session.query method", async () => { + const mockSession = { + query: vi.fn(() => Promise.resolve({ rows: [] })), + }; + + const mockDb = { + session: mockSession, + }; + + instrumentDrizzleClient(mockDb); + + // Direct query through session + await mockSession.query("INSERT INTO users (name) VALUES ($1)", ["John"]); + + const spans = exporter.getFinishedSpans(); + expect(spans).toHaveLength(1); + expect(spans[0]?.name).toBe("drizzle.insert"); + expect(spans[0]?.attributes["db.statement"]).toBe("INSERT INTO users (name) VALUES ($1)"); + }); + it("only instruments once when called multiple times", async () => { const mockClient = { query: vi.fn(() => Promise.resolve({ rows: [] })), @@ -452,6 +551,40 @@ describe("instrumentDrizzleClient", () => { expect(spans[0]?.status.code).toBe(SpanStatusCode.ERROR); }); + it("instruments transaction execute calls", async () => { + let txObject: any; + + const mockSession = { + transaction: vi.fn(async (callback: any) => { + // Create a mock transaction object + txObject = { + execute: vi.fn(() => Promise.resolve({ rows: [] })), + }; + return callback(txObject); + }), + }; + + const mockDb = { + session: mockSession, + }; + + instrumentDrizzleClient(mockDb); + + // Execute a transaction with RLS queries + await mockSession.transaction(async (tx: any) => { + await tx.execute({ sql: "SET LOCAL role org_role" }); + await tx.execute({ sql: "SELECT set_config('request.org_id', $1, true)", params: ["org123"] }); + }); + + const spans = exporter.getFinishedSpans(); + expect(spans).toHaveLength(2); + expect(spans[0]?.name).toBe("drizzle.set"); + expect(spans[0]?.attributes["db.statement"]).toBe("SET LOCAL role org_role"); + expect(spans[0]?.attributes["db.transaction"]).toBe(true); + expect(spans[1]?.name).toBe("drizzle.select"); + expect(spans[1]?.attributes["db.transaction"]).toBe(true); + }); + it("only instruments session once when called multiple times", () => { const mockSession = { execute: vi.fn(() => Promise.resolve({ rows: [] })), diff --git a/packages/otel-drizzle/src/index.ts b/packages/otel-drizzle/src/index.ts index 5f2d996..afe30a0 100644 --- a/packages/otel-drizzle/src/index.ts +++ b/packages/otel-drizzle/src/index.ts @@ -22,11 +22,13 @@ export const SEMATTRS_NET_PEER_PORT = "net.peer.port"; type QueryCallback = (error: unknown, result: unknown) => void; -type QueryFunction = (...args: unknown[]) => unknown; +type QueryFunction = (...args: any[]) => any; interface DrizzleClientLike { - query: QueryFunction; + query?: QueryFunction; + execute?: QueryFunction; [INSTRUMENTED_FLAG]?: true; + [key: string]: any; // Allow other properties } /** @@ -82,14 +84,14 @@ function extractQueryText(queryArg: unknown): string | undefined { return queryArg; } if (queryArg && typeof queryArg === "object") { + // Generic SQL object format (used by LibSQL, MySQL, and others) + if (typeof (queryArg as { sql?: unknown }).sql === "string") { + return (queryArg as { sql: string }).sql; + } // PostgreSQL-style query object if (typeof (queryArg as { text?: unknown }).text === "string") { return (queryArg as { text: string }).text; } - // MySQL/generic-style query object - if (typeof (queryArg as { sql?: unknown }).sql === "string") { - return (queryArg as { sql: string }).sql; - } // Drizzle SQL object if ( typeof (queryArg as { queryChunks?: unknown }).queryChunks === "object" @@ -141,15 +143,16 @@ function finalizeSpan(span: Span, error?: unknown): void { } /** - * Instruments a database connection pool with OpenTelemetry tracing. + * Instruments a database connection pool or client with OpenTelemetry tracing. * - * This function wraps the connection pool's `query` method to automatically create - * spans for each database operation. + * This function wraps the connection's `query` or `execute` method to automatically create + * spans for each database operation. It automatically detects which method is available + * and instruments it appropriately for any database driver. * The instrumentation is idempotent - calling it multiple times on the same - * pool will only instrument it once. + * connection will only instrument it once. * * @typeParam TClient - The type of the database connection pool or client - * @param client - The database connection pool or client to instrument (e.g., pg Pool, mysql2 Connection) + * @param client - The database connection pool or client to instrument (e.g., pg Pool, mysql2 Connection, LibSQL Client) * @param config - Optional configuration for instrumentation behavior * @returns The instrumented pool/client (same instance, modified in place) * @@ -197,6 +200,17 @@ function finalizeSpan(span: Span, error?: unknown): void { * const sqlite = new Database('database.db'); * const db = drizzle(instrumentDrizzle(sqlite, { dbSystem: 'sqlite' })); * ``` + * + * @example + * ```typescript + * // LibSQL/Turso (automatically detects 'execute' method) + * import { drizzle } from 'drizzle-orm/libsql'; + * import { createClient } from '@libsql/client'; + * import { instrumentDrizzle } from '@kubiks/otel-drizzle'; + * + * const client = createClient({ url: '...', authToken: '...' }); + * const db = drizzle(instrumentDrizzle(client, { dbSystem: 'sqlite' })); + * ``` */ export function instrumentDrizzle( client: TClient, @@ -205,7 +219,12 @@ export function instrumentDrizzle( if (!client) { return client; } - if (typeof client.query !== "function") { + + // Check if client has query or execute method + const hasQuery = typeof client.query === "function"; + const hasExecute = typeof client.execute === "function"; + + if (!hasQuery && !hasExecute) { return client; } @@ -224,11 +243,18 @@ export function instrumentDrizzle( } = config ?? {}; const tracer = trace.getTracer(tracerName); - const originalQuery = client.query; + + // Store the original method (query or execute) + const methodName = hasQuery ? "query" : "execute"; + const originalMethod = hasQuery ? client.query : client.execute; - const instrumentedQuery: QueryFunction = function instrumented( - this: unknown, - ...incomingArgs: unknown[] + if (!originalMethod) { + return client; + } + + const instrumentedMethod: QueryFunction = function instrumented( + this: any, + ...incomingArgs: any[] ) { const args = [...incomingArgs]; let callback: QueryCallback | undefined; @@ -283,7 +309,7 @@ export function instrumentDrizzle( }; try { - return originalQuery.apply(this, [...args, wrappedCallback]); + return originalMethod.apply(this, [...args, wrappedCallback]); } catch (error) { finalizeSpan(span, error); throw error; @@ -294,7 +320,7 @@ export function instrumentDrizzle( // Promise-based pattern return context.with(activeContext, () => { try { - const result = originalQuery.apply(this, args); + const result = originalMethod.apply(this, args); return Promise.resolve(result) .then((value) => { finalizeSpan(span); @@ -312,7 +338,13 @@ export function instrumentDrizzle( }; client[INSTRUMENTED_FLAG] = true; - client.query = instrumentedQuery; + + // Replace the original method with the instrumented one + if (hasQuery) { + client.query = instrumentedMethod; + } else { + client.execute = instrumentedMethod; + } return client; } @@ -321,21 +353,37 @@ export function instrumentDrizzle( * Interface for Drizzle database instances with minimal type requirements. */ interface DrizzleDbLike { - $client?: DrizzleClientLike; + $client?: DrizzleClientLike | any; // Allow any client type + execute?: QueryFunction; // Direct execute method on db + transaction?: QueryFunction; // Transaction method on db _?: { session?: { execute?: QueryFunction; [INSTRUMENTED_FLAG]?: true; + [key: string]: any; }; + [key: string]: any; }; + [INSTRUMENTED_FLAG]?: true; + [key: string]: any; // Allow other properties } /** * Instruments an already created Drizzle database instance with OpenTelemetry tracing. * - * This function is useful when you have a Drizzle database instance created with - * `const db = drizzle(pool, { schema })` and want to add instrumentation to it - * without having direct access to the underlying pool. + * This function instruments the database at the session level, intercepting: + * - `session.prepareQuery` - Used by all query builders (select, insert, update, delete) + * - `session.query` - Used for direct SQL execution + * - `$client` methods - As a fallback for underlying connection + * + * This ensures all database operations are traced, whether you use: + * - Query builders: `db.select().from(table)` + * - Direct execution: `db.execute(sql)` + * - Transactions: `db.transaction()` + * + * This is useful when you have a Drizzle database instance created with + * `const db = drizzle(connectionString)` or `const db = drizzle(pool, { schema })` + * and want to add instrumentation to it without having direct access to the underlying pool. * * The instrumentation is idempotent - calling it multiple times on the same * database will only instrument it once. @@ -390,14 +438,340 @@ export function instrumentDrizzleClient( return db; } - // Try to instrument via $client first (most common case) - if (db.$client && typeof db.$client.query === "function") { - instrumentDrizzle(db.$client, config); + // Check if already instrumented + if (db[INSTRUMENTED_FLAG]) { return db; } - // Try to instrument via session.execute as fallback - if (db._ && db._.session && typeof db._.session.execute === "function") { + const { + tracerName = DEFAULT_TRACER_NAME, + dbSystem = DEFAULT_DB_SYSTEM, + dbName, + captureQueryText = true, + maxQueryTextLength = 1000, + peerName, + peerPort, + } = config ?? {}; + + const tracer = trace.getTracer(tracerName); + let instrumented = false; + + // First priority: Instrument the session directly + // This is where all queries actually go through + if ((db as any).session && !instrumented) { + const session = (db as any).session; + + // Check if session has prepareQuery method (used by select/insert/update/delete) + if (typeof session.prepareQuery === "function" && !session[INSTRUMENTED_FLAG]) { + const originalPrepareQuery = session.prepareQuery; + + session.prepareQuery = function(...args: any[]) { + const prepared = originalPrepareQuery.apply(this, args); + + // Wrap the prepared query's execute method + if (prepared && typeof prepared.execute === "function") { + const originalPreparedExecute = prepared.execute; + + prepared.execute = function(this: any, ...executeArgs: any[]) { + // Extract query information from the query object + const queryObj = args[0]; // The query object passed to prepareQuery + const queryText = queryObj?.sql || queryObj?.queryString || extractQueryText(queryObj); + const operation = queryText ? extractOperation(queryText) : undefined; + const spanName = operation + ? `drizzle.${operation.toLowerCase()}` + : "drizzle.query"; + + // Start span + const span = tracer.startSpan(spanName, { kind: SpanKind.CLIENT }); + span.setAttribute(SEMATTRS_DB_SYSTEM, dbSystem); + + if (operation) { + span.setAttribute(SEMATTRS_DB_OPERATION, operation); + } + + if (dbName) { + span.setAttribute(SEMATTRS_DB_NAME, dbName); + } + + if (captureQueryText && queryText !== undefined) { + const sanitized = sanitizeQueryText(queryText, maxQueryTextLength); + span.setAttribute(SEMATTRS_DB_STATEMENT, sanitized); + } + + if (peerName) { + span.setAttribute(SEMATTRS_NET_PEER_NAME, peerName); + } + + if (peerPort) { + span.setAttribute(SEMATTRS_NET_PEER_PORT, peerPort); + } + + const activeContext = trace.setSpan(context.active(), span); + + // Execute the prepared query + return context.with(activeContext, () => { + try { + const result = originalPreparedExecute.apply(this, executeArgs); + return Promise.resolve(result) + .then((value) => { + finalizeSpan(span); + return value; + }) + .catch((error) => { + finalizeSpan(span, error); + throw error; + }); + } catch (error) { + finalizeSpan(span, error); + throw error; + } + }); + }; + } + + return prepared; + }; + + session[INSTRUMENTED_FLAG] = true; + instrumented = true; + } + + // Also instrument direct query method if exists + if (typeof session.query === "function" && !session[INSTRUMENTED_FLAG + "_query"]) { + const originalQuery = session.query; + + session.query = function(this: any, queryString: string, params: any[]) { + const operation = queryString ? extractOperation(queryString) : undefined; + const spanName = operation + ? `drizzle.${operation.toLowerCase()}` + : "drizzle.query"; + + // Start span + const span = tracer.startSpan(spanName, { kind: SpanKind.CLIENT }); + span.setAttribute(SEMATTRS_DB_SYSTEM, dbSystem); + + if (operation) { + span.setAttribute(SEMATTRS_DB_OPERATION, operation); + } + + if (dbName) { + span.setAttribute(SEMATTRS_DB_NAME, dbName); + } + + if (captureQueryText && queryString !== undefined) { + const sanitized = sanitizeQueryText(queryString, maxQueryTextLength); + span.setAttribute(SEMATTRS_DB_STATEMENT, sanitized); + } + + if (peerName) { + span.setAttribute(SEMATTRS_NET_PEER_NAME, peerName); + } + + if (peerPort) { + span.setAttribute(SEMATTRS_NET_PEER_PORT, peerPort); + } + + const activeContext = trace.setSpan(context.active(), span); + + // Execute the query + return context.with(activeContext, () => { + try { + const result = originalQuery.apply(this, [queryString, params]); + return Promise.resolve(result) + .then((value) => { + finalizeSpan(span); + return value; + }) + .catch((error) => { + finalizeSpan(span, error); + throw error; + }); + } catch (error) { + finalizeSpan(span, error); + throw error; + } + }); + }; + + session[INSTRUMENTED_FLAG + "_query"] = true; + instrumented = true; + } + + // Instrument transaction method to ensure transaction sessions are also instrumented + if (typeof session.transaction === "function" && !session[INSTRUMENTED_FLAG + "_transaction"]) { + const originalTransaction = session.transaction; + + session.transaction = function(this: any, transactionCallback: any, ...restArgs: any[]) { + // Wrap the transaction callback to instrument the tx object + const wrappedCallback = async function(tx: any) { + // Instrument the transaction's session if it has one + if (tx && (tx.session || tx._?.session || tx)) { + const txSession = tx.session || tx._?.session || tx; + + // Instrument tx.execute if it exists + if (typeof tx.execute === "function" && !tx[INSTRUMENTED_FLAG + "_execute"]) { + const originalTxExecute = tx.execute; + + tx.execute = function(this: any, ...executeArgs: any[]) { + const queryText = extractQueryText(executeArgs[0]); + const operation = queryText ? extractOperation(queryText) : undefined; + const spanName = operation + ? `drizzle.${operation.toLowerCase()}` + : "drizzle.query"; + + // Start span + const span = tracer.startSpan(spanName, { kind: SpanKind.CLIENT }); + span.setAttribute(SEMATTRS_DB_SYSTEM, dbSystem); + span.setAttribute("db.transaction", true); + + if (operation) { + span.setAttribute(SEMATTRS_DB_OPERATION, operation); + } + + if (dbName) { + span.setAttribute(SEMATTRS_DB_NAME, dbName); + } + + if (captureQueryText && queryText !== undefined) { + const sanitized = sanitizeQueryText(queryText, maxQueryTextLength); + span.setAttribute(SEMATTRS_DB_STATEMENT, sanitized); + } + + if (peerName) { + span.setAttribute(SEMATTRS_NET_PEER_NAME, peerName); + } + + if (peerPort) { + span.setAttribute(SEMATTRS_NET_PEER_PORT, peerPort); + } + + const activeContext = trace.setSpan(context.active(), span); + + // Execute the query + return context.with(activeContext, () => { + try { + const result = originalTxExecute.apply(this, executeArgs); + return Promise.resolve(result) + .then((value) => { + finalizeSpan(span); + return value; + }) + .catch((error) => { + finalizeSpan(span, error); + throw error; + }); + } catch (error) { + finalizeSpan(span, error); + throw error; + } + }); + }; + + tx[INSTRUMENTED_FLAG + "_execute"] = true; + } + + // Also instrument txSession.prepareQuery if it exists + if (typeof txSession.prepareQuery === "function" && !txSession[INSTRUMENTED_FLAG + "_tx"]) { + const originalTxPrepareQuery = txSession.prepareQuery; + + txSession.prepareQuery = function(...prepareArgs: any[]) { + const prepared = originalTxPrepareQuery.apply(this, prepareArgs); + + // Wrap the prepared query's execute method + if (prepared && typeof prepared.execute === "function") { + const originalPreparedExecute = prepared.execute; + + prepared.execute = function(this: any, ...executeArgs: any[]) { + // Extract query information from the query object + const queryObj = prepareArgs[0]; // The query object passed to prepareQuery + const queryText = queryObj?.sql || queryObj?.queryString || extractQueryText(queryObj); + const operation = queryText ? extractOperation(queryText) : undefined; + const spanName = operation + ? `drizzle.${operation.toLowerCase()}` + : "drizzle.query"; + + // Start span + const span = tracer.startSpan(spanName, { kind: SpanKind.CLIENT }); + span.setAttribute(SEMATTRS_DB_SYSTEM, dbSystem); + span.setAttribute("db.transaction", true); + + if (operation) { + span.setAttribute(SEMATTRS_DB_OPERATION, operation); + } + + if (dbName) { + span.setAttribute(SEMATTRS_DB_NAME, dbName); + } + + if (captureQueryText && queryText !== undefined) { + const sanitized = sanitizeQueryText(queryText, maxQueryTextLength); + span.setAttribute(SEMATTRS_DB_STATEMENT, sanitized); + } + + if (peerName) { + span.setAttribute(SEMATTRS_NET_PEER_NAME, peerName); + } + + if (peerPort) { + span.setAttribute(SEMATTRS_NET_PEER_PORT, peerPort); + } + + const activeContext = trace.setSpan(context.active(), span); + + // Execute the prepared query + return context.with(activeContext, () => { + try { + const result = originalPreparedExecute.apply(this, executeArgs); + return Promise.resolve(result) + .then((value) => { + finalizeSpan(span); + return value; + }) + .catch((error) => { + finalizeSpan(span, error); + throw error; + }); + } catch (error) { + finalizeSpan(span, error); + throw error; + } + }); + }; + } + + return prepared; + }; + + txSession[INSTRUMENTED_FLAG + "_tx"] = true; + } + } + + // Call the original callback with the instrumented tx + return transactionCallback(tx); + }; + + // Call the original transaction with the wrapped callback + return originalTransaction.apply(this, [wrappedCallback, ...restArgs]); + }; + + session[INSTRUMENTED_FLAG + "_transaction"] = true; + instrumented = true; + } + } + + // Second priority: Try to instrument via $client + // This handles the underlying connection pool + if (db.$client && !instrumented) { + const client = db.$client; + // Check if client has query or execute function + if (typeof client.query === "function" || typeof client.execute === "function") { + instrumentDrizzle(client, config); + instrumented = true; + } + } + + // Third priority: Try to instrument via session.execute as fallback + if (db._ && db._.session && typeof db._.session.execute === "function" && !instrumented) { const session = db._.session; // Check if already instrumented @@ -423,8 +797,8 @@ export function instrumentDrizzleClient( } const instrumentedExecute: QueryFunction = function instrumented( - this: unknown, - ...args: unknown[] + this: any, + ...args: any[] ) { // Extract query information const queryText = extractQueryText(args[0]); @@ -482,6 +856,12 @@ export function instrumentDrizzleClient( session[INSTRUMENTED_FLAG] = true; session.execute = instrumentedExecute; + instrumented = true; + } + + // Mark the db as instrumented if we instrumented anything + if (instrumented) { + db[INSTRUMENTED_FLAG] = true; } return db; From 7b7dc83f7d8733d63a77f2a0d655513053f00bcd Mon Sep 17 00:00:00 2001 From: Alex Holovach Date: Sun, 5 Oct 2025 08:19:58 -0500 Subject: [PATCH 4/5] update docs --- packages/otel-drizzle/README.md | 177 +++++---- packages/otel-drizzle/src/index.ts | 563 ++++++++++++++++++----------- 2 files changed, 447 insertions(+), 293 deletions(-) diff --git a/packages/otel-drizzle/README.md b/packages/otel-drizzle/README.md index 1e8de12..aff11b8 100644 --- a/packages/otel-drizzle/README.md +++ b/packages/otel-drizzle/README.md @@ -42,118 +42,149 @@ Works with any observability platform that supports OpenTelemetry including: ## Usage -There are two ways to instrument Drizzle ORM with OpenTelemetry: +### Instrument Your Drizzle Database (Recommended) -### Option 1: Instrument the Connection Pool (Recommended) - -Wrap your database connection pool with `instrumentDrizzle()` before passing it to Drizzle: +Use `instrumentDrizzleClient()` to add tracing to your Drizzle database instance. This is the simplest and most straightforward approach: ```typescript -import { drizzle } from "drizzle-orm/node-postgres"; -import { Pool } from "pg"; -import { instrumentDrizzle } from "@kubiks/otel-drizzle"; +import { drizzle } from "drizzle-orm/postgres-js"; +import { instrumentDrizzleClient } from "@kubiks/otel-drizzle"; -const pool = new Pool({ connectionString: process.env.DATABASE_URL }); -const instrumentedPool = instrumentDrizzle(pool); -const db = drizzle(instrumentedPool); +// Create your Drizzle database instance as usual +const db = drizzle(process.env.DATABASE_URL!); + +// Add instrumentation with a single line +instrumentDrizzleClient(db); // That's it! All queries are now traced automatically const users = await db.select().from(usersTable); ``` -### Option 2: Instrument an Existing Drizzle Client +### Database-Specific Examples -If you already have a Drizzle database instance or don't have access to the underlying pool, use `instrumentDrizzleClient()`. This method instruments the database at the session level, capturing all query operations: +#### PostgreSQL ```typescript -// Works with postgres-js (Postgres.js) +// PostgreSQL with postgres.js (recommended for serverless) import { drizzle } from "drizzle-orm/postgres-js"; +import postgres from "postgres"; import { instrumentDrizzleClient } from "@kubiks/otel-drizzle"; -import * as schema from "./schema"; -const db = drizzle(process.env.DATABASE_URL!, { schema }); +// Using connection string directly +const db = drizzle(process.env.DATABASE_URL!); +instrumentDrizzleClient(db, { dbSystem: "postgresql" }); -// Instrument the existing database instance -instrumentDrizzleClient(db); - -// All queries are now traced automatically -const users = await db.select().from(schema.users); -// Direct execute calls are also traced -await db.execute("SELECT * FROM users"); -// Transactions are also traced -await db.transaction(async (tx) => { - await tx.insert(schema.users).values({ name: "John" }); -}); -``` - -### Optional Configuration - -Both instrumentation methods accept the same configuration options: - -```typescript -// Option 1: Instrument the pool -const pool = new Pool({ connectionString: process.env.DATABASE_URL }); -const instrumentedPool = instrumentDrizzle(pool, { - dbSystem: "postgresql", // Database type (default: 'postgresql') - dbName: "myapp", // Database name for spans - captureQueryText: true, // Include SQL in traces (default: true) - maxQueryTextLength: 1000, // Max SQL length (default: 1000) - peerName: "db.example.com", // Database server hostname - peerPort: 5432, // Database server port -}); -const db = drizzle(instrumentedPool); - -// Option 2: Instrument the Drizzle client -const db = drizzle(pool, { schema }); +// Or with a client instance +const queryClient = postgres(process.env.DATABASE_URL!); +const db = drizzle({ client: queryClient }); instrumentDrizzleClient(db, { dbSystem: "postgresql", dbName: "myapp", - captureQueryText: true, peerName: "db.example.com", peerPort: 5432, }); ``` -### Works with All Drizzle-Supported Databases - -This package automatically detects and instruments **all databases that Drizzle ORM supports**. It works by detecting whether your database driver uses a `query` or `execute` method and instrumenting it appropriately. This includes: - -- **PostgreSQL** (node-postgres, postgres.js, Neon, Vercel Postgres, etc.) -- **MySQL** (mysql2, PlanetScale, TiDB, etc.) -- **SQLite** (better-sqlite3, LibSQL/Turso, Cloudflare D1, etc.) -- **And any other Drizzle-supported database** - ```typescript -// PostgreSQL with postgres-js (Postgres.js) - use instrumentDrizzleClient -import { drizzle } from "drizzle-orm/postgres-js"; -const db = drizzle(process.env.DATABASE_URL!); -instrumentDrizzleClient(db); - -// PostgreSQL with node-postgres (pg) - use instrumentDrizzle on pool +// PostgreSQL with node-postgres (pg) import { drizzle } from "drizzle-orm/node-postgres"; import { Pool } from "pg"; -const pool = new Pool({ connectionString: process.env.DATABASE_URL }); -const db = drizzle(instrumentDrizzle(pool)); +import { instrumentDrizzleClient } from "@kubiks/otel-drizzle"; -// MySQL with mysql2 (uses 'execute' or 'query' method) +// Using connection string directly +const db = drizzle(process.env.DATABASE_URL!); +instrumentDrizzleClient(db, { dbSystem: "postgresql" }); + +// Or with a pool instance +const pool = new Pool({ connectionString: process.env.DATABASE_URL }); +const db = drizzle({ client: pool }); +instrumentDrizzleClient(db, { dbSystem: "postgresql" }); +``` + +#### MySQL + +```typescript +// MySQL with mysql2 import { drizzle } from "drizzle-orm/mysql2"; import mysql from "mysql2/promise"; -const connection = await mysql.createConnection(process.env.DATABASE_URL); -const db = drizzle(instrumentDrizzle(connection, { dbSystem: "mysql" })); +import { instrumentDrizzleClient } from "@kubiks/otel-drizzle"; +// Using connection string directly +const db = drizzle(process.env.DATABASE_URL!); +instrumentDrizzleClient(db, { dbSystem: "mysql" }); + +// Or with a connection instance +const connection = await mysql.createConnection({ + host: "localhost", + user: "root", + database: "mydb", + // ... other connection options +}); +const db = drizzle({ client: connection }); +instrumentDrizzleClient(db, { + dbSystem: "mysql", + dbName: "mydb", + peerName: "localhost", + peerPort: 3306, +}); +``` + +#### SQLite + +```typescript // SQLite with better-sqlite3 import { drizzle } from "drizzle-orm/better-sqlite3"; import Database from "better-sqlite3"; -const sqlite = new Database("database.db"); -const db = drizzle(instrumentDrizzle(sqlite, { dbSystem: "sqlite" })); +import { instrumentDrizzleClient } from "@kubiks/otel-drizzle"; -// LibSQL/Turso (uses 'execute' method) +// Using file path directly +const db = drizzle("sqlite.db"); +instrumentDrizzleClient(db, { dbSystem: "sqlite" }); + +// Or with a Database instance +const sqlite = new Database("sqlite.db"); +const db = drizzle({ client: sqlite }); +instrumentDrizzleClient(db, { dbSystem: "sqlite" }); +``` + +```typescript +// SQLite with LibSQL/Turso import { drizzle } from "drizzle-orm/libsql"; import { createClient } from "@libsql/client"; -const client = createClient({ url: "...", authToken: "..." }); -const db = drizzle(instrumentDrizzle(client, { dbSystem: "sqlite" })); +import { instrumentDrizzleClient } from "@kubiks/otel-drizzle"; + +// Using connection config directly +const db = drizzle({ + connection: { + url: process.env.DATABASE_URL!, + authToken: process.env.DATABASE_AUTH_TOKEN, + } +}); +instrumentDrizzleClient(db, { dbSystem: "sqlite" }); + +// Or with a client instance +const client = createClient({ + url: process.env.DATABASE_URL!, + authToken: process.env.DATABASE_AUTH_TOKEN, +}); +const db = drizzle({ client }); +instrumentDrizzleClient(db, { dbSystem: "sqlite" }); ``` +### Configuration Options + +```typescript +instrumentDrizzleClient(db, { + dbSystem: "postgresql", // Database type: 'postgresql' | 'mysql' | 'sqlite' (default: 'postgresql') + dbName: "myapp", // Database name for spans + captureQueryText: true, // Include SQL in traces (default: true) + maxQueryTextLength: 1000, // Max SQL length (default: 1000) + peerName: "db.example.com", // Database server hostname + peerPort: 5432, // Database server port +}); +``` + + ## What You Get Each database query automatically creates a span with rich telemetry data: diff --git a/packages/otel-drizzle/src/index.ts b/packages/otel-drizzle/src/index.ts index afe30a0..fa91c34 100644 --- a/packages/otel-drizzle/src/index.ts +++ b/packages/otel-drizzle/src/index.ts @@ -143,16 +143,15 @@ function finalizeSpan(span: Span, error?: unknown): void { } /** - * Instruments a database connection pool or client with OpenTelemetry tracing. + * Instruments a database connection pool/client with OpenTelemetry tracing. * - * This function wraps the connection's `query` or `execute` method to automatically create - * spans for each database operation. It automatically detects which method is available - * and instruments it appropriately for any database driver. - * The instrumentation is idempotent - calling it multiple times on the same - * connection will only instrument it once. + * This function wraps the connection's `query` and `execute` methods to create spans for each database + * operation. + * The instrumentation is idempotent - calling it multiple times on the same connection will only + * instrument it once. * * @typeParam TClient - The type of the database connection pool or client - * @param client - The database connection pool or client to instrument (e.g., pg Pool, mysql2 Connection, LibSQL Client) + * @param client - The database connection pool or client to instrument * @param config - Optional configuration for instrumentation behavior * @returns The instrumented pool/client (same instance, modified in place) * @@ -164,19 +163,13 @@ function finalizeSpan(span: Span, error?: unknown): void { * import { instrumentDrizzle } from '@kubiks/otel-drizzle'; * * const pool = new Pool({ connectionString: process.env.DATABASE_URL }); - * const instrumentedPool = instrumentDrizzle(pool); - * const db = drizzle(instrumentedPool); - * - * // With custom configuration * const instrumentedPool = instrumentDrizzle(pool, { * dbSystem: 'postgresql', * dbName: 'myapp', - * captureQueryText: true, - * maxQueryTextLength: 1000, * peerName: 'db.example.com', * peerPort: 5432, * }); - * const db = drizzle(instrumentedPool); + * const db = drizzle({ client: instrumentedPool }); * ``` * * @example @@ -186,8 +179,13 @@ function finalizeSpan(span: Span, error?: unknown): void { * import mysql from 'mysql2/promise'; * import { instrumentDrizzle } from '@kubiks/otel-drizzle'; * - * const connection = await mysql.createConnection(process.env.DATABASE_URL); - * const db = drizzle(instrumentDrizzle(connection, { dbSystem: 'mysql' })); + * const connection = await mysql.createConnection({ + * host: 'localhost', + * user: 'root', + * database: 'mydb', + * }); + * const instrumentedConnection = instrumentDrizzle(connection, { dbSystem: 'mysql' }); + * const db = drizzle({ client: instrumentedConnection }); * ``` * * @example @@ -197,19 +195,24 @@ function finalizeSpan(span: Span, error?: unknown): void { * import Database from 'better-sqlite3'; * import { instrumentDrizzle } from '@kubiks/otel-drizzle'; * - * const sqlite = new Database('database.db'); - * const db = drizzle(instrumentDrizzle(sqlite, { dbSystem: 'sqlite' })); + * const sqlite = new Database('sqlite.db'); + * const instrumentedSqlite = instrumentDrizzle(sqlite, { dbSystem: 'sqlite' }); + * const db = drizzle({ client: instrumentedSqlite }); * ``` * * @example * ```typescript - * // LibSQL/Turso (automatically detects 'execute' method) + * // LibSQL/Turso * import { drizzle } from 'drizzle-orm/libsql'; * import { createClient } from '@libsql/client'; * import { instrumentDrizzle } from '@kubiks/otel-drizzle'; * - * const client = createClient({ url: '...', authToken: '...' }); - * const db = drizzle(instrumentDrizzle(client, { dbSystem: 'sqlite' })); + * const client = createClient({ + * url: process.env.DATABASE_URL!, + * authToken: process.env.DATABASE_AUTH_TOKEN, + * }); + * const instrumentedClient = instrumentDrizzle(client, { dbSystem: 'sqlite' }); + * const db = drizzle({ client: instrumentedClient }); * ``` */ export function instrumentDrizzle( @@ -219,11 +222,11 @@ export function instrumentDrizzle( if (!client) { return client; } - + // Check if client has query or execute method const hasQuery = typeof client.query === "function"; const hasExecute = typeof client.execute === "function"; - + if (!hasQuery && !hasExecute) { return client; } @@ -243,7 +246,7 @@ export function instrumentDrizzle( } = config ?? {}; const tracer = trace.getTracer(tracerName); - + // Store the original method (query or execute) const methodName = hasQuery ? "query" : "execute"; const originalMethod = hasQuery ? client.query : client.execute; @@ -338,7 +341,7 @@ export function instrumentDrizzle( }; client[INSTRUMENTED_FLAG] = true; - + // Replace the original method with the instrumented one if (hasQuery) { client.query = instrumentedMethod; @@ -369,21 +372,10 @@ interface DrizzleDbLike { } /** - * Instruments an already created Drizzle database instance with OpenTelemetry tracing. + * Instruments a Drizzle database instance with OpenTelemetry tracing. * - * This function instruments the database at the session level, intercepting: - * - `session.prepareQuery` - Used by all query builders (select, insert, update, delete) - * - `session.query` - Used for direct SQL execution - * - `$client` methods - As a fallback for underlying connection - * - * This ensures all database operations are traced, whether you use: - * - Query builders: `db.select().from(table)` - * - Direct execution: `db.execute(sql)` - * - Transactions: `db.transaction()` - * - * This is useful when you have a Drizzle database instance created with - * `const db = drizzle(connectionString)` or `const db = drizzle(pool, { schema })` - * and want to add instrumentation to it without having direct access to the underlying pool. + * This function instruments the database at the session level, automatically tracing all database + * operations including query builders, direct SQL execution, and transactions. * * The instrumentation is idempotent - calling it multiple times on the same * database will only instrument it once. @@ -395,39 +387,109 @@ interface DrizzleDbLike { * * @example * ```typescript - * // When you have a pre-created Drizzle database instance - * import { drizzle } from 'drizzle-orm/node-postgres'; - * import { Pool } from 'pg'; + * // PostgreSQL with postgres.js + * import { drizzle } from 'drizzle-orm/postgres-js'; + * import postgres from 'postgres'; * import { instrumentDrizzleClient } from '@kubiks/otel-drizzle'; - * import * as schema from './schema'; * - * const pool = new Pool({ connectionString: process.env.DATABASE_URL }); - * const db = drizzle(pool, { schema }); + * // Using connection string + * const db = drizzle(process.env.DATABASE_URL!); + * instrumentDrizzleClient(db, { dbSystem: 'postgresql' }); * - * // Instrument the existing db instance - * instrumentDrizzleClient(db, { - * dbSystem: 'postgresql', - * dbName: 'myapp', - * captureQueryText: true, - * peerName: 'db.example.com', - * peerPort: 5432, - * }); - * - * // Now all queries through db are traced - * const users = await db.select().from(schema.users); + * // Or with a client instance + * const queryClient = postgres(process.env.DATABASE_URL!); + * const db = drizzle({ client: queryClient }); + * instrumentDrizzleClient(db, { dbSystem: 'postgresql' }); * ``` * * @example * ```typescript - * // Works with any Drizzle driver + * // PostgreSQL with node-postgres (pg) + * import { drizzle } from 'drizzle-orm/node-postgres'; + * import { Pool } from 'pg'; + * import { instrumentDrizzleClient } from '@kubiks/otel-drizzle'; + * + * // Using connection string + * const db = drizzle(process.env.DATABASE_URL!); + * instrumentDrizzleClient(db, { dbSystem: 'postgresql' }); + * + * // Or with a pool + * const pool = new Pool({ connectionString: process.env.DATABASE_URL }); + * const db = drizzle({ client: pool }); + * instrumentDrizzleClient(db, { + * dbSystem: 'postgresql', + * dbName: 'myapp', + * peerName: 'db.example.com', + * peerPort: 5432, + * }); + * ``` + * + * @example + * ```typescript + * // MySQL with mysql2 * import { drizzle } from 'drizzle-orm/mysql2'; * import mysql from 'mysql2/promise'; * import { instrumentDrizzleClient } from '@kubiks/otel-drizzle'; * - * const connection = await mysql.createConnection(process.env.DATABASE_URL); - * const db = drizzle(connection); - * + * // Using connection string + * const db = drizzle(process.env.DATABASE_URL!); * instrumentDrizzleClient(db, { dbSystem: 'mysql' }); + * + * // Or with a connection + * const connection = await mysql.createConnection({ + * host: 'localhost', + * user: 'root', + * database: 'mydb', + * }); + * const db = drizzle({ client: connection }); + * instrumentDrizzleClient(db, { + * dbSystem: 'mysql', + * dbName: 'mydb', + * peerName: 'localhost', + * peerPort: 3306, + * }); + * ``` + * + * @example + * ```typescript + * // SQLite with better-sqlite3 + * import { drizzle } from 'drizzle-orm/better-sqlite3'; + * import Database from 'better-sqlite3'; + * import { instrumentDrizzleClient } from '@kubiks/otel-drizzle'; + * + * // Using file path + * const db = drizzle('sqlite.db'); + * instrumentDrizzleClient(db, { dbSystem: 'sqlite' }); + * + * // Or with a Database instance + * const sqlite = new Database('sqlite.db'); + * const db = drizzle({ client: sqlite }); + * instrumentDrizzleClient(db, { dbSystem: 'sqlite' }); + * ``` + * + * @example + * ```typescript + * // SQLite with LibSQL/Turso + * import { drizzle } from 'drizzle-orm/libsql'; + * import { createClient } from '@libsql/client'; + * import { instrumentDrizzleClient } from '@kubiks/otel-drizzle'; + * + * // Using connection config + * const db = drizzle({ + * connection: { + * url: process.env.DATABASE_URL!, + * authToken: process.env.DATABASE_AUTH_TOKEN, + * } + * }); + * instrumentDrizzleClient(db, { dbSystem: 'sqlite' }); + * + * // Or with a client instance + * const client = createClient({ + * url: process.env.DATABASE_URL!, + * authToken: process.env.DATABASE_AUTH_TOKEN, + * }); + * const db = drizzle({ client }); + * instrumentDrizzleClient(db, { dbSystem: 'sqlite' }); * ``` */ export function instrumentDrizzleClient( @@ -460,23 +522,31 @@ export function instrumentDrizzleClient( // This is where all queries actually go through if ((db as any).session && !instrumented) { const session = (db as any).session; - + // Check if session has prepareQuery method (used by select/insert/update/delete) - if (typeof session.prepareQuery === "function" && !session[INSTRUMENTED_FLAG]) { + if ( + typeof session.prepareQuery === "function" && + !session[INSTRUMENTED_FLAG] + ) { const originalPrepareQuery = session.prepareQuery; - - session.prepareQuery = function(...args: any[]) { + + session.prepareQuery = function (...args: any[]) { const prepared = originalPrepareQuery.apply(this, args); - + // Wrap the prepared query's execute method if (prepared && typeof prepared.execute === "function") { const originalPreparedExecute = prepared.execute; - - prepared.execute = function(this: any, ...executeArgs: any[]) { + + prepared.execute = function (this: any, ...executeArgs: any[]) { // Extract query information from the query object const queryObj = args[0]; // The query object passed to prepareQuery - const queryText = queryObj?.sql || queryObj?.queryString || extractQueryText(queryObj); - const operation = queryText ? extractOperation(queryText) : undefined; + const queryText = + queryObj?.sql || + queryObj?.queryString || + extractQueryText(queryObj); + const operation = queryText + ? extractOperation(queryText) + : undefined; const spanName = operation ? `drizzle.${operation.toLowerCase()}` : "drizzle.query"; @@ -494,7 +564,10 @@ export function instrumentDrizzleClient( } if (captureQueryText && queryText !== undefined) { - const sanitized = sanitizeQueryText(queryText, maxQueryTextLength); + const sanitized = sanitizeQueryText( + queryText, + maxQueryTextLength, + ); span.setAttribute(SEMATTRS_DB_STATEMENT, sanitized); } @@ -528,20 +601,25 @@ export function instrumentDrizzleClient( }); }; } - + return prepared; }; - + session[INSTRUMENTED_FLAG] = true; instrumented = true; } // Also instrument direct query method if exists - if (typeof session.query === "function" && !session[INSTRUMENTED_FLAG + "_query"]) { + if ( + typeof session.query === "function" && + !session[INSTRUMENTED_FLAG + "_query"] + ) { const originalQuery = session.query; - - session.query = function(this: any, queryString: string, params: any[]) { - const operation = queryString ? extractOperation(queryString) : undefined; + + session.query = function (this: any, queryString: string, params: any[]) { + const operation = queryString + ? extractOperation(queryString) + : undefined; const spanName = operation ? `drizzle.${operation.toLowerCase()}` : "drizzle.query"; @@ -592,188 +670,233 @@ export function instrumentDrizzleClient( } }); }; - + session[INSTRUMENTED_FLAG + "_query"] = true; instrumented = true; } // Instrument transaction method to ensure transaction sessions are also instrumented - if (typeof session.transaction === "function" && !session[INSTRUMENTED_FLAG + "_transaction"]) { + if ( + typeof session.transaction === "function" && + !session[INSTRUMENTED_FLAG + "_transaction"] + ) { const originalTransaction = session.transaction; - - session.transaction = function(this: any, transactionCallback: any, ...restArgs: any[]) { + + session.transaction = function ( + this: any, + transactionCallback: any, + ...restArgs: any[] + ) { // Wrap the transaction callback to instrument the tx object - const wrappedCallback = async function(tx: any) { - // Instrument the transaction's session if it has one - if (tx && (tx.session || tx._?.session || tx)) { - const txSession = tx.session || tx._?.session || tx; - - // Instrument tx.execute if it exists - if (typeof tx.execute === "function" && !tx[INSTRUMENTED_FLAG + "_execute"]) { - const originalTxExecute = tx.execute; - - tx.execute = function(this: any, ...executeArgs: any[]) { - const queryText = extractQueryText(executeArgs[0]); - const operation = queryText ? extractOperation(queryText) : undefined; - const spanName = operation - ? `drizzle.${operation.toLowerCase()}` - : "drizzle.query"; + const wrappedCallback = async function (tx: any) { + // Instrument the transaction's session if it has one + if (tx && (tx.session || tx._?.session || tx)) { + const txSession = tx.session || tx._?.session || tx; - // Start span - const span = tracer.startSpan(spanName, { kind: SpanKind.CLIENT }); - span.setAttribute(SEMATTRS_DB_SYSTEM, dbSystem); - span.setAttribute("db.transaction", true); + // Instrument tx.execute if it exists + if ( + typeof tx.execute === "function" && + !tx[INSTRUMENTED_FLAG + "_execute"] + ) { + const originalTxExecute = tx.execute; - if (operation) { - span.setAttribute(SEMATTRS_DB_OPERATION, operation); - } + tx.execute = function (this: any, ...executeArgs: any[]) { + const queryText = extractQueryText(executeArgs[0]); + const operation = queryText + ? extractOperation(queryText) + : undefined; + const spanName = operation + ? `drizzle.${operation.toLowerCase()}` + : "drizzle.query"; - if (dbName) { - span.setAttribute(SEMATTRS_DB_NAME, dbName); - } + // Start span + const span = tracer.startSpan(spanName, { + kind: SpanKind.CLIENT, + }); + span.setAttribute(SEMATTRS_DB_SYSTEM, dbSystem); + span.setAttribute("db.transaction", true); - if (captureQueryText && queryText !== undefined) { - const sanitized = sanitizeQueryText(queryText, maxQueryTextLength); - span.setAttribute(SEMATTRS_DB_STATEMENT, sanitized); - } + if (operation) { + span.setAttribute(SEMATTRS_DB_OPERATION, operation); + } - if (peerName) { - span.setAttribute(SEMATTRS_NET_PEER_NAME, peerName); - } + if (dbName) { + span.setAttribute(SEMATTRS_DB_NAME, dbName); + } - if (peerPort) { - span.setAttribute(SEMATTRS_NET_PEER_PORT, peerPort); - } + if (captureQueryText && queryText !== undefined) { + const sanitized = sanitizeQueryText( + queryText, + maxQueryTextLength, + ); + span.setAttribute(SEMATTRS_DB_STATEMENT, sanitized); + } - const activeContext = trace.setSpan(context.active(), span); + if (peerName) { + span.setAttribute(SEMATTRS_NET_PEER_NAME, peerName); + } - // Execute the query - return context.with(activeContext, () => { - try { - const result = originalTxExecute.apply(this, executeArgs); - return Promise.resolve(result) - .then((value) => { - finalizeSpan(span); - return value; - }) - .catch((error) => { - finalizeSpan(span, error); - throw error; - }); - } catch (error) { - finalizeSpan(span, error); - throw error; - } - }); - }; - - tx[INSTRUMENTED_FLAG + "_execute"] = true; - } - - // Also instrument txSession.prepareQuery if it exists - if (typeof txSession.prepareQuery === "function" && !txSession[INSTRUMENTED_FLAG + "_tx"]) { - const originalTxPrepareQuery = txSession.prepareQuery; - - txSession.prepareQuery = function(...prepareArgs: any[]) { - const prepared = originalTxPrepareQuery.apply(this, prepareArgs); - - // Wrap the prepared query's execute method - if (prepared && typeof prepared.execute === "function") { - const originalPreparedExecute = prepared.execute; - - prepared.execute = function(this: any, ...executeArgs: any[]) { - // Extract query information from the query object - const queryObj = prepareArgs[0]; // The query object passed to prepareQuery - const queryText = queryObj?.sql || queryObj?.queryString || extractQueryText(queryObj); - const operation = queryText ? extractOperation(queryText) : undefined; - const spanName = operation - ? `drizzle.${operation.toLowerCase()}` - : "drizzle.query"; + if (peerPort) { + span.setAttribute(SEMATTRS_NET_PEER_PORT, peerPort); + } - // Start span - const span = tracer.startSpan(spanName, { kind: SpanKind.CLIENT }); - span.setAttribute(SEMATTRS_DB_SYSTEM, dbSystem); - span.setAttribute("db.transaction", true); + const activeContext = trace.setSpan(context.active(), span); - if (operation) { - span.setAttribute(SEMATTRS_DB_OPERATION, operation); - } - - if (dbName) { - span.setAttribute(SEMATTRS_DB_NAME, dbName); - } - - if (captureQueryText && queryText !== undefined) { - const sanitized = sanitizeQueryText(queryText, maxQueryTextLength); - span.setAttribute(SEMATTRS_DB_STATEMENT, sanitized); - } - - if (peerName) { - span.setAttribute(SEMATTRS_NET_PEER_NAME, peerName); - } - - if (peerPort) { - span.setAttribute(SEMATTRS_NET_PEER_PORT, peerPort); - } - - const activeContext = trace.setSpan(context.active(), span); - - // Execute the prepared query - return context.with(activeContext, () => { - try { - const result = originalPreparedExecute.apply(this, executeArgs); - return Promise.resolve(result) - .then((value) => { - finalizeSpan(span); - return value; - }) - .catch((error) => { - finalizeSpan(span, error); - throw error; - }); - } catch (error) { - finalizeSpan(span, error); - throw error; - } + // Execute the query + return context.with(activeContext, () => { + try { + const result = originalTxExecute.apply(this, executeArgs); + return Promise.resolve(result) + .then((value) => { + finalizeSpan(span); + return value; + }) + .catch((error) => { + finalizeSpan(span, error); + throw error; }); - }; + } catch (error) { + finalizeSpan(span, error); + throw error; } - - return prepared; - }; - - txSession[INSTRUMENTED_FLAG + "_tx"] = true; - } + }); + }; + + tx[INSTRUMENTED_FLAG + "_execute"] = true; } - - // Call the original callback with the instrumented tx - return transactionCallback(tx); - }; - + + // Also instrument txSession.prepareQuery if it exists + if ( + typeof txSession.prepareQuery === "function" && + !txSession[INSTRUMENTED_FLAG + "_tx"] + ) { + const originalTxPrepareQuery = txSession.prepareQuery; + + txSession.prepareQuery = function (...prepareArgs: any[]) { + const prepared = originalTxPrepareQuery.apply( + this, + prepareArgs, + ); + + // Wrap the prepared query's execute method + if (prepared && typeof prepared.execute === "function") { + const originalPreparedExecute = prepared.execute; + + prepared.execute = function ( + this: any, + ...executeArgs: any[] + ) { + // Extract query information from the query object + const queryObj = prepareArgs[0]; // The query object passed to prepareQuery + const queryText = + queryObj?.sql || + queryObj?.queryString || + extractQueryText(queryObj); + const operation = queryText + ? extractOperation(queryText) + : undefined; + const spanName = operation + ? `drizzle.${operation.toLowerCase()}` + : "drizzle.query"; + + // Start span + const span = tracer.startSpan(spanName, { + kind: SpanKind.CLIENT, + }); + span.setAttribute(SEMATTRS_DB_SYSTEM, dbSystem); + span.setAttribute("db.transaction", true); + + if (operation) { + span.setAttribute(SEMATTRS_DB_OPERATION, operation); + } + + if (dbName) { + span.setAttribute(SEMATTRS_DB_NAME, dbName); + } + + if (captureQueryText && queryText !== undefined) { + const sanitized = sanitizeQueryText( + queryText, + maxQueryTextLength, + ); + span.setAttribute(SEMATTRS_DB_STATEMENT, sanitized); + } + + if (peerName) { + span.setAttribute(SEMATTRS_NET_PEER_NAME, peerName); + } + + if (peerPort) { + span.setAttribute(SEMATTRS_NET_PEER_PORT, peerPort); + } + + const activeContext = trace.setSpan(context.active(), span); + + // Execute the prepared query + return context.with(activeContext, () => { + try { + const result = originalPreparedExecute.apply( + this, + executeArgs, + ); + return Promise.resolve(result) + .then((value) => { + finalizeSpan(span); + return value; + }) + .catch((error) => { + finalizeSpan(span, error); + throw error; + }); + } catch (error) { + finalizeSpan(span, error); + throw error; + } + }); + }; + } + + return prepared; + }; + + txSession[INSTRUMENTED_FLAG + "_tx"] = true; + } + } + + // Call the original callback with the instrumented tx + return transactionCallback(tx); + }; + // Call the original transaction with the wrapped callback return originalTransaction.apply(this, [wrappedCallback, ...restArgs]); }; - + session[INSTRUMENTED_FLAG + "_transaction"] = true; instrumented = true; } } - // Second priority: Try to instrument via $client - // This handles the underlying connection pool if (db.$client && !instrumented) { const client = db.$client; // Check if client has query or execute function - if (typeof client.query === "function" || typeof client.execute === "function") { + if ( + typeof client.query === "function" || + typeof client.execute === "function" + ) { instrumentDrizzle(client, config); instrumented = true; } } // Third priority: Try to instrument via session.execute as fallback - if (db._ && db._.session && typeof db._.session.execute === "function" && !instrumented) { + if ( + db._ && + db._.session && + typeof db._.session.execute === "function" && + !instrumented + ) { const session = db._.session; - + // Check if already instrumented if (session[INSTRUMENTED_FLAG]) { return db; From 50f189aac1629c6197668b2da4d98cd259cd5db1 Mon Sep 17 00:00:00 2001 From: Alex Holovach Date: Sun, 5 Oct 2025 08:25:25 -0500 Subject: [PATCH 5/5] changeset --- .changeset/quick-walls-brake.md | 5 +++++ packages/otel-drizzle/package.json | 2 +- 2 files changed, 6 insertions(+), 1 deletion(-) create mode 100644 .changeset/quick-walls-brake.md diff --git a/.changeset/quick-walls-brake.md b/.changeset/quick-walls-brake.md new file mode 100644 index 0000000..6b68f8b --- /dev/null +++ b/.changeset/quick-walls-brake.md @@ -0,0 +1,5 @@ +--- +"@kubiks/otel-drizzle": minor +--- + +Added instrumentDrizzleClient func to instrument any drizzle instance diff --git a/packages/otel-drizzle/package.json b/packages/otel-drizzle/package.json index 86273fa..abadd31 100644 --- a/packages/otel-drizzle/package.json +++ b/packages/otel-drizzle/package.json @@ -1,6 +1,6 @@ { "name": "@kubiks/otel-drizzle", - "version": "2.0.9", + "version": "2.0.3", "private": false, "publishConfig": { "access": "public"