feat(prisma): handle raw queries

lumigo-io · Dec 26, 2023 · 2a25844 · 2a25844
1 parent cda10ab
commit 2a25844
Show file tree

Hide file tree

Showing 4 changed files with 36 additions and 17 deletions.
diff --git a/src/hooks/prisma.test.js b/src/hooks/prisma.test.js
@@ -18,7 +18,7 @@ describe('Prisma', () => {
       prismaClientLibrary = require('@prisma/client')
     });
 
-    test('no $extends hook', async () => {
+    test('does not fail on no $extends hook', async () => {
       // Override constructor
       prismaClientLibrary.PrismaClient = function () {
         return {
@@ -31,7 +31,7 @@ describe('Prisma', () => {
       expect(() => new prismaClientLibrary.PrismaClient()).not.toThrow()
     });
 
-    test('no PrismaClient class', async () => {
+    test('does not fail when no PrismaClient is missing', async () => {
       delete prismaClientLibrary.PrismaClient
 
       expect(() => hookPrisma(prismaClientLibrary)).not.toThrow()
@@ -56,7 +56,7 @@ describe('Prisma', () => {
 
     afterAll(() => client && client.$disconnect());
 
-    test('simple operation', async () => {
+    test('produces simple operation span', async () => {
       const count = await client.user.count();
 
       const spans = SpansContainer.getSpans();
@@ -66,7 +66,7 @@ describe('Prisma', () => {
         .withId(spans[0].id)
         .withStarted(spans[0].started)
         .withEnded(spans[0].ended)
-        .withModelName('User')
+        .withModel('User')
         .withOperation('count')
         .withQueryArgs("{}")
         .withResult(count.toString())
@@ -75,7 +75,7 @@ describe('Prisma', () => {
       expect(spans[0]).toEqual(expectedSpan);
     });
 
-    test('operation with arguments', async () => {
+    test('produces span for operations with arguments', async () => {
       const email = getRandomId()
 
       const user = await client.user.create({ data: { name: 'John Doe', email } });
@@ -90,7 +90,7 @@ describe('Prisma', () => {
         .withId(createUserSpan.id)
         .withStarted(createUserSpan.started)
         .withEnded(createUserSpan.ended)
-        .withModelName('User')
+        .withModel('User')
         .withOperation('create')
         .withQueryArgs(`{\"data\":{\"name\":\"John Doe\",\"email\":\"${email}\"}}`)
         .withResult(`{\"id\":${user.id},\"email\":\"${email}\",\"name\":\"John Doe\"}`)
@@ -100,7 +100,7 @@ describe('Prisma', () => {
         .withId(findUserSpan.id)
         .withStarted(findUserSpan.started)
         .withEnded(findUserSpan.ended)
-        .withModelName('User')
+        .withModel('User')
         .withOperation('findFirst')
         .withQueryArgs(`{\"where\":{\"id\":${user.id}}}`)
         .withResult(`{\"id\":${user.id},\"email\":\"${email}\",\"name\":\"John Doe\"}`)
@@ -111,7 +111,7 @@ describe('Prisma', () => {
       expect(findUserSpan).toEqual(expectedFindSpan);
     });
 
-    test("operation with error", async () => {
+    test("produces an error span on erroneous operation", async () => {
       const email = getRandomId();
 
       const user = await client.user.create({ data: { name: 'John Doe', email } });
@@ -126,7 +126,7 @@ describe('Prisma', () => {
         .withId(successSpan.id)
         .withStarted(successSpan.started)
         .withEnded(successSpan.ended)
-        .withModelName('User')
+        .withModel('User')
         .withOperation('create')
         .withQueryArgs(`{\"data\":{\"name\":\"John Doe\",\"email\":\"${email}\"}}`)
         .withResult(`{\"id\":${user.id},\"email\":\"${email}\",\"name\":\"John Doe\"}`)
@@ -136,7 +136,7 @@ describe('Prisma', () => {
         .withId(errorSpan.id)
         .withStarted(errorSpan.started)
         .withEnded(errorSpan.ended)
-        .withModelName('User')
+        .withModel('User')
         .withOperation('create')
         .withQueryArgs(`{\"data\":{\"name\":\"John Doe\",\"email\":\"${email}\"}}`)
         .withError(expect.stringContaining("Unique constraint failed on the fields"))
@@ -147,7 +147,7 @@ describe('Prisma', () => {
       expect(errorSpan).toEqual(expectedErrorSpan);
     })
 
-    test("allows other user-extensions", async () => {
+    test("allows other user-extensions to be run", async () => {
       let extensionExecuted = false
 
       const clientWithUserExtension = client.$extends({
@@ -164,6 +164,25 @@ describe('Prisma', () => {
       expect(extensionExecuted).toBe(true)
     })
 
-    test.todo("raw query")
+    test("produces a raw query span", async () => {
+      const userId = getRandomId();
+      await client.$queryRaw`SELECT * FROM User WHERE id = ${userId};`
+
+      const spans = SpansContainer.getSpans();
+      expect(spans).toHaveLength(1);
+
+      const [rawQuerySpan] = spans;
+
+      const expectedRawQuerySpan = new PrismaSpanBuilder()
+        .withId(rawQuerySpan.id)
+        .withStarted(rawQuerySpan.started)
+        .withEnded(rawQuerySpan.ended)
+        .withOperation('$queryRaw')
+        .withQueryArgs(`{\"values\":[\"${userId}\"],\"strings\":[\"SELECT * FROM User WHERE id = \",\";\"]}`)
+        .withResult(`[]`)
+        .build();
+
+      expect(rawQuerySpan).toEqual(expectedRawQuerySpan);
+    })
   })
 })
diff --git a/src/hooks/prisma.ts b/src/hooks/prisma.ts
@@ -18,8 +18,8 @@ async function queryExtension({ query, args, model, operation }) {
       started: Date.now(),
     },
     {
-      modelName: model,
-      operation: operation,
+      model,
+      operation,
       queryArgs: args,
     }
   );

diff --git a/src/spans/prismaSpan.ts b/src/spans/prismaSpan.ts
@@ -13,7 +13,7 @@ export const createPrismaSpan = (
   return {
     ...baseSpan,
     started: requestMetadata.started,
-    modelName: prismaFields.modelName,
+    model: prismaFields.model,
     operation: prismaFields.operation,
     queryArgs: prismaFields.queryArgs ? payloadStringify(prismaFields.queryArgs) : '',
   };

diff --git a/testUtils/prismaSpanBuilder.js b/testUtils/prismaSpanBuilder.js
@@ -41,8 +41,8 @@ export class PrismaSpanBuilder {
     return this;
   }
 
-  withModelName = (modelName) => {
-    this._span.modelName = modelName;
+  withModel = (model) => {
+    this._span.model = model;
     return this;
   }