code-gen: support setting primary key in generated insert queries (#491)

compasjs · Nov 14, 2020 · f336ead · f336ead
1 parent ff2fff1
commit f336ead
Show file tree

Hide file tree

Showing 3 changed files with 50 additions and 15 deletions.
diff --git a/packages/code-gen/src/generator/sql/query-basics.js b/packages/code-gen/src/generator/sql/query-basics.js
@@ -218,17 +218,19 @@ function insertQuery(context, imports, type) {
     /**
      * @param {Postgres} sql
      * @param {${type.partial.insertType}|(${type.partial.insertType}[])} insert
+     * @param {{ withPrimaryKey: boolean }=} options
      * @returns {Promise<${type.uniqueName}[]>}
      */
-    export function ${type.name}Insert(sql, insert) {
+    export function ${type.name}Insert(sql, insert, options = {}) {
       if (insert === undefined || insert.length === 0) {
         return []; 
       }
+      options.withPrimaryKey = options.withPrimaryKey ?? false;
 
       return query\`
         INSERT INTO "${type.name}" ($\{${type.name}Fields(
-        "", { excludePrimaryKey: true })})
-        VALUES $\{${type.name}InsertValues(insert)}
+        "", { excludePrimaryKey: !options.withPrimaryKey })})
+        VALUES $\{${type.name}InsertValues(insert, { includePrimaryKey: options.withPrimaryKey })}
         RETURNING $\{${type.name}Fields("")}
       \`.exec(sql);
     }

diff --git a/packages/code-gen/test/sql.test.js b/packages/code-gen/test/sql.test.js
@@ -194,6 +194,17 @@ test("code-gen/e2e/sql", async (t) => {
     }
   });
 
+  t.test("Insert with primary key", async (t) => {
+    const id = uuid();
+    const [category] = await client.queries.categoryInsert(
+      sql,
+      { id, label: "TestPK" },
+      { withPrimaryKey: true },
+    );
+
+    t.equal(category.id, id);
+  });
+
   t.test("deletedAt in the future", async (t) => {
     const future = new Date();
     future.setUTCDate(future.getUTCDate() + 1);

diff --git a/packages/store/src/generated/query-basics.js b/packages/store/src/generated/query-basics.js