From 2cc9bed1580d12f6c1ccb58dfa79c8d07048f978 Mon Sep 17 00:00:00 2001
From: Simon Walker <s.r.walker101@googlemail.com>
Date: Fri, 16 Feb 2024 16:53:46 +0000
Subject: [PATCH] Support psql insert returning

---
 __tests__/__snapshots__/resolvers.test.js.snap | 2 +-
 __tests__/resolvers.test.js                    | 2 +-
 rds/index.js                                   | 6 +++++-
 3 files changed, 7 insertions(+), 3 deletions(-)

diff --git a/__tests__/__snapshots__/resolvers.test.js.snap b/__tests__/__snapshots__/resolvers.test.js.snap
index c50fed4..dbfaf70 100644
--- a/__tests__/__snapshots__/resolvers.test.js.snap
+++ b/__tests__/__snapshots__/resolvers.test.js.snap
@@ -111,7 +111,7 @@ exports[`rds resolvers mysql update 1`] = `
 exports[`rds resolvers postgresql insert 1`] = `
 {
   "statements": [
-    "INSERT INTO "persons" ("name") VALUES (:P0)",
+    "INSERT INTO "persons" ("name") VALUES (:P0) RETURNING *",
   ],
   "variableMap": {
     ":P0": "test",
diff --git a/__tests__/resolvers.test.js b/__tests__/resolvers.test.js
index c543214..8f9f7cf 100644
--- a/__tests__/resolvers.test.js
+++ b/__tests__/resolvers.test.js
@@ -427,7 +427,7 @@ describe("rds resolvers", () => {
       const code = `
         export function request(ctx) {
         const { input: values } = ctx.args;
-        const insertStatement = rds.insert({ table: 'persons', values });
+        const insertStatement = rds.insert({ table: 'persons', values, returning: "*" });
         
         return rds.createPgStatement(insertStatement)
         }
diff --git a/rds/index.js b/rds/index.js
index ea47e51..fdf1c33 100644
--- a/rds/index.js
+++ b/rds/index.js
@@ -129,7 +129,7 @@ class StatementBuilder {
           break;
         }
         case "INSERT": {
-          const { table, values } = properties;
+          const { table, values, returning } = properties;
           const tableName = this.getTableName(table);
 
           let query = `INSERT INTO ${tableName}`;
@@ -143,6 +143,10 @@ class StatementBuilder {
           }
           query = `${query} (${columnTextItems.join(', ')}) VALUES (${valuesTextItems.join(', ')})`;
 
+          if (returning) {
+            query = `${query} RETURNING ${returning}`;
+          }
+
           this.result.statements.push(query);
           break;
         }