From 835598e1c1d96800b8e00cf1cb01868c4caa331f Mon Sep 17 00:00:00 2001
From: Qian Li <qian.li@dbos.dev>
Date: Wed, 28 Aug 2024 13:36:32 -0700
Subject: [PATCH] Add Drizzle support

---
 dbos-rules.test.ts | 12 +++++++++++-
 dbos-rules.ts      |  2 +-
 package-lock.json  |  4 ++--
 package.json       |  2 +-
 4 files changed, 15 insertions(+), 5 deletions(-)

diff --git a/dbos-rules.test.ts b/dbos-rules.test.ts
index 7fb2da3..5ac3ef7 100644
--- a/dbos-rules.test.ts
+++ b/dbos-rules.test.ts
@@ -90,6 +90,10 @@ function makeSqlInjectionCode(code: string, sqlClient: string): string {
       query(query: string, ...values: any[]) {}
     }
 
+    class PgDatabase {
+      execute(query: string, ...values: any[]) {}
+    }
+
     function Transaction(target?: any, key?: any, descriptor?: any): any {
       return descriptor;
     }
@@ -244,7 +248,13 @@ const testSet: TestSet = [
       // Success test #9 (testing unsubstituted template literals)
       makeSqlInjectionSuccessTest(`
         ctxt.client.raw(\`foo\`);
-      `)
+      `),
+
+      // Success test #10 (testing Drizzle support)
+      makeSqlInjectionSuccessTest(`
+        ctxt.client.execute("foo");`,
+        "PgDatabase"
+      )
     ],
 
     [
diff --git a/dbos-rules.ts b/dbos-rules.ts
index 70000e3..5648aa3 100644
--- a/dbos-rules.ts
+++ b/dbos-rules.ts
@@ -57,7 +57,7 @@ const ormClientInfoForRawSqlQueries: Map<string, string[]> = new Map([
   ["PrismaClient", ["$queryRawUnsafe", "$executeRawUnsafe"]], // For Prisma
   ["EntityManager", ["query"]], // For TypeORM
   ["PoolClient", ["query"]], // This is supported in `dbos-transact` (see `user_database.ts`, but not sure what ORM this corresponds to)
-  // ["PgDatabase", []], // For Drizzle (TODO: add full support for this)
+  ["PgDatabase", []], // For Drizzle (Currently we don't detect raw SQL calls for Drizzle)
 ]);
 
 const assignmentTokenKinds = new Set([
diff --git a/package-lock.json b/package-lock.json
index b0ddfd7..b21fa20 100644
--- a/package-lock.json
+++ b/package-lock.json
@@ -1,12 +1,12 @@
 {
   "name": "@dbos-inc/eslint-plugin",
-  "version": "3.3.3",
+  "version": "3.3.4",
   "lockfileVersion": 3,
   "requires": true,
   "packages": {
     "": {
       "name": "@dbos-inc/eslint-plugin",
-      "version": "3.3.3",
+      "version": "3.3.4",
       "license": "MIT",
       "dependencies": {
         "@types/node": "^20.14.12",
diff --git a/package.json b/package.json
index 7ba6ebd..169aa9f 100644
--- a/package.json
+++ b/package.json
@@ -1,6 +1,6 @@
 {
   "name": "@dbos-inc/eslint-plugin",
-  "version": "3.3.3",
+  "version": "3.3.4",
   "description": "eslint plugin for DBOS SDK",
   "license": "MIT",
   "repository": {