From 61e6c3f23f79a8df4a5b3e37b24753c520e4ceed Mon Sep 17 00:00:00 2001
From: Jonathan Ginn <jonathan@ginn.work>
Date: Fri, 24 Jan 2025 15:11:49 -0500
Subject: [PATCH] feat: support `RETURNING` clause (#139)

---
 src/database.ts           | 12 +++++++++---
 test/connectors/_tests.ts | 33 ++++++++++++++++++++++-----------
 2 files changed, 31 insertions(+), 14 deletions(-)

diff --git a/src/database.ts b/src/database.ts
index 63e020c..8fb195e 100644
--- a/src/database.ts
+++ b/src/database.ts
@@ -1,7 +1,9 @@
 import { sqlTemplate } from "./template";
-import type { Connector, Database } from "./types";
+import type { Connector, Database, SQLDialect } from "./types";
 
-const SQL_WITH_RES_RE = /^select/i;
+const SQL_SELECT_RE = /^select/i;
+const SQL_RETURNING_RE = /[\s]returning[\s]/i;
+const DIALECTS_WITH_RET: Set<SQLDialect> = new Set(["postgresql", "sqlite"]);
 
 /**
  * Creates and returns a database interface using the specified connector.
@@ -33,7 +35,11 @@ export function createDatabase<TConnector extends Connector = Connector>(
 
     sql: async (strings, ...values) => {
       const [sql, params] = sqlTemplate(strings, ...values);
-      if (SQL_WITH_RES_RE.test(sql)) {
+      if (
+        SQL_SELECT_RE.test(sql) /* select */ ||
+        // prettier-ignore
+        (DIALECTS_WITH_RET.has(connector.dialect) && SQL_RETURNING_RE.test(sql)) /* returning */
+      ) {
         const rows = await connector.prepare(sql).all(...params);
         return {
           rows,
diff --git a/test/connectors/_tests.ts b/test/connectors/_tests.ts
index 4181006..230d547 100644
--- a/test/connectors/_tests.ts
+++ b/test/connectors/_tests.ts
@@ -8,6 +8,16 @@ export function testConnector<TConnector extends Connector = Connector>(opts: {
   });
 
   const userId = "1001";
+  const userSnapshot = `
+    [
+      {
+        "email": "",
+        "firstName": "John",
+        "id": "1001",
+        "lastName": "Doe",
+      },
+    ]
+  `;
 
   it("instance matches", async () => {
     const instance = await db.getInstance();
@@ -34,20 +44,21 @@ export function testConnector<TConnector extends Connector = Connector>(opts: {
   });
 
   it("insert", async () => {
-    await db.sql`INSERT INTO users VALUES (${userId}, 'John', 'Doe', '')`;
+    switch(opts.dialect) {
+      case "mysql": {
+        await db.sql`INSERT INTO users VALUES (${userId}, 'John', 'Doe', '')`;
+        break;
+      }
+      default: {
+        const { rows } = await db.sql`INSERT INTO users VALUES (${userId}, 'John', 'Doe', '') RETURNING *`;
+        expect(rows).toMatchInlineSnapshot(userSnapshot);
+        break;
+      }
+    }
   });
 
   it("select", async () => {
     const { rows } = await db.sql`SELECT * FROM users WHERE id = ${userId}`;
-    expect(rows).toMatchInlineSnapshot(`
-      [
-        {
-          "email": "",
-          "firstName": "John",
-          "id": "1001",
-          "lastName": "Doe",
-        },
-      ]
-    `);
+    expect(rows).toMatchInlineSnapshot(userSnapshot);
   });
 }