11import { SemanticAttributes } from '@opentelemetry/semantic-conventions' ;
22import { opentelemetry } from '@traceloop/otel-proto' ;
3+ import { Parser } from 'node-sql-parser' ;
34import {
45 CompareOptions ,
56 filterByAttributeStringValue ,
@@ -9,6 +10,7 @@ export class PostgreSQLQuery {
910 constructor (
1011 readonly spans : opentelemetry . proto . trace . v1 . ISpan [ ] ,
1112 private readonly serviceName : string ,
13+ private parser = new Parser ( ) ,
1214 ) { }
1315
1416 withDatabaseName ( name : string | RegExp , options ?: CompareOptions ) {
@@ -45,9 +47,33 @@ export class PostgreSQLQuery {
4547 return new PostgreSQLQuery ( filteredSpans , this . serviceName ) ;
4648 }
4749
48- withOperationAndTable ( operation : string , table : string ) {
49- const regex = new RegExp ( `${ operation } .*${ table } ` , 'i' ) ; // case insensitive
50+ withOperations ( ...operations : string [ ] ) {
51+ const filteredSpans = this . spans . filter ( ( span ) => {
52+ const statement = span . attributes ?. find (
53+ ( attribute ) => attribute . key === SemanticAttributes . DB_STATEMENT ,
54+ ) ?. value ?. stringValue ;
55+
56+ if ( ! statement ) {
57+ return false ;
58+ }
59+
60+ const lowerCaseStatement = statement . toLowerCase ( ) ;
61+
62+ return operations . every ( ( operation ) =>
63+ lowerCaseStatement . includes ( operation . toLowerCase ( ) ) ,
64+ ) ;
65+ } ) ;
5066
67+ if ( filteredSpans . length === 0 ) {
68+ throw new Error (
69+ `No query by ${ this . serviceName } to postgresql with operations ${ operations } was found` ,
70+ ) ;
71+ }
72+
73+ return new PostgreSQLQuery ( filteredSpans , this . serviceName ) ;
74+ }
75+
76+ withTables ( ...tables : string [ ] ) {
5177 const filteredSpans = this . spans . filter ( ( span ) => {
5278 const statement = span . attributes ?. find (
5379 ( attribute ) => attribute . key === SemanticAttributes . DB_STATEMENT ,
@@ -57,15 +83,25 @@ export class PostgreSQLQuery {
5783 return false ;
5884 }
5985
60- return regex . test ( statement ) ;
86+ const allTablesInStatement = this . parser
87+ . tableList ( prepareQuery ( statement ) , { database : 'PostgresQL' } )
88+ . map ( ( table ) => table . split ( '::' ) [ 2 ] . trim ( ) ) ;
89+
90+ return tables . every ( ( table ) =>
91+ allTablesInStatement . includes ( table . toLowerCase ( ) ) ,
92+ ) ;
6193 } ) ;
6294
6395 if ( filteredSpans . length === 0 ) {
6496 throw new Error (
65- `No query by ${ this . serviceName } to postgresql with operation ${ operation } and table ${ table } was found` ,
97+ `No query by ${ this . serviceName } to postgresql with tables ${ tables } was found` ,
6698 ) ;
6799 }
68100
69101 return new PostgreSQLQuery ( filteredSpans , this . serviceName ) ;
70102 }
71103}
104+
105+ const prepareQuery = (
106+ query : string , // remove double quotes and replace %s with 111
107+ ) => query . replace ( / " / g, '' ) . replace ( / % s / g, '111' ) . toLocaleLowerCase ( ) ;
0 commit comments