From 5834c749c2ef3f5619566d7104670a51146e4dca Mon Sep 17 00:00:00 2001
From: Tomer <42910634+tomer-friedman@users.noreply.github.com>
Date: Tue, 13 Jun 2023 16:15:38 +0300
Subject: [PATCH] feat: sql compare (#76)

---
 docs/jest-otel/syntax/db-pg.mdx               |  1 +
 .../service/to-query-postgresql.test.ts       |  1 +
 .../src/matchers/utils/filters.ts             |  2 +-
 .../src/resources/postgresql-query.ts         | 28 +++++++++++++++++--
 4 files changed, 29 insertions(+), 3 deletions(-)

diff --git a/docs/jest-otel/syntax/db-pg.mdx b/docs/jest-otel/syntax/db-pg.mdx
index 050524c..2d4603a 100644
--- a/docs/jest-otel/syntax/db-pg.mdx
+++ b/docs/jest-otel/syntax/db-pg.mdx
@@ -30,6 +30,7 @@ So, a complete assertion can look like:
 expectTrace(traceloop.serviceByName('orders-service'))
   .toQueryPostgreSQL()
   .withDatabaseName('postgres')
+  .withOperationAndTable('INSERT', 'orders')
   .withStatement(
     /INSERT INTO orders \(id, price_in_cents\) VALUES \('[0-9A-Fa-f]{8}-[0-9A-Fa-f]{4}-[0-9A-Fa-f]{4}-[0-9A-Fa-f]{4}-[0-9A-Fa-f]{12}', [0-9]+\)/,
     { compareType: COMPARE_TYPE.REGEX },
diff --git a/packages/expect-opentelemetry/src/matchers/service/to-query-postgresql.test.ts b/packages/expect-opentelemetry/src/matchers/service/to-query-postgresql.test.ts
index 33f26f6..8d8d23e 100644
--- a/packages/expect-opentelemetry/src/matchers/service/to-query-postgresql.test.ts
+++ b/packages/expect-opentelemetry/src/matchers/service/to-query-postgresql.test.ts
@@ -24,6 +24,7 @@ describe('postgresql query', () => {
     expectTrace(traceloop.serviceByName('orders-service'))
       .toQueryPostgreSQL()
       .withDatabaseName('postgres', { compareType: COMPARE_TYPE.EQUALS })
+      .withOperationAndTable('INSERT', 'orders')
       .withStatement(
         /INSERT INTO orders \(id, price_in_cents\) VALUES \('[0-9A-Fa-f]{8}-[0-9A-Fa-f]{4}-[0-9A-Fa-f]{4}-[0-9A-Fa-f]{4}-[0-9A-Fa-f]{12}', [0-9]+\)/,
         { compareType: COMPARE_TYPE.REGEX },
diff --git a/packages/expect-opentelemetry/src/matchers/utils/filters.ts b/packages/expect-opentelemetry/src/matchers/utils/filters.ts
index 34b66c5..c17ad4b 100644
--- a/packages/expect-opentelemetry/src/matchers/utils/filters.ts
+++ b/packages/expect-opentelemetry/src/matchers/utils/filters.ts
@@ -7,7 +7,7 @@ export const filterByAttributeKey = (
   attName: string,
 ) =>
   spans.filter((span) => {
-    span.attributes?.find((attribute) => {
+    return span.attributes?.find((attribute) => {
       return attribute.key === attName;
     });
   });
diff --git a/packages/expect-opentelemetry/src/resources/postgresql-query.ts b/packages/expect-opentelemetry/src/resources/postgresql-query.ts
index cfa4116..438e273 100644
--- a/packages/expect-opentelemetry/src/resources/postgresql-query.ts
+++ b/packages/expect-opentelemetry/src/resources/postgresql-query.ts
@@ -11,7 +11,7 @@ export class PostgreSQLQuery {
     private readonly serviceName: string,
   ) {}
 
-  withDatabaseName(name: string | RegExp, options: CompareOptions) {
+  withDatabaseName(name: string | RegExp, options?: CompareOptions) {
     const filteredSpans = filterByAttributeStringValue(
       this.spans,
       SemanticAttributes.DB_NAME,
@@ -28,7 +28,7 @@ export class PostgreSQLQuery {
     return new PostgreSQLQuery(filteredSpans, this.serviceName);
   }
 
-  withStatement(statement: string | RegExp, options: CompareOptions) {
+  withStatement(statement: string | RegExp, options?: CompareOptions) {
     const filteredSpans = filterByAttributeStringValue(
       this.spans,
       SemanticAttributes.DB_STATEMENT,
@@ -44,4 +44,28 @@ export class PostgreSQLQuery {
 
     return new PostgreSQLQuery(filteredSpans, this.serviceName);
   }
+
+  withOperationAndTable(operation: string, table: string) {
+    const regex = new RegExp(`${operation}.*${table}`, 'i'); // case insensitive
+
+    const filteredSpans = this.spans.filter((span) => {
+      const statement = span.attributes?.find(
+        (attribute) => attribute.key === SemanticAttributes.DB_STATEMENT,
+      )?.value?.stringValue;
+
+      if (!statement) {
+        return false;
+      }
+
+      return regex.test(statement);
+    });
+
+    if (filteredSpans.length === 0) {
+      throw new Error(
+        `No query by ${this.serviceName} to postgresql with operation ${operation} and table ${table} was found`,
+      );
+    }
+
+    return new PostgreSQLQuery(filteredSpans, this.serviceName);
+  }
 }