Add support for dynamic columns with returning - fixes porsager#317

porsager · porsager · commit 822fb2158ebb · 2022-04-14T22:21:30.000+02:00
diff --git a/src/types.js b/src/types.js
@@ -110,9 +110,28 @@ function values(first, rest, parameters, types, transform) {
   return valuesBuilder(multi ? first : [first], parameters, types, transform, columns)
 }
 
+function select(first, rest, parameters, types, transform) {
+  typeof first === 'string' && (first = [first].concat(rest))
+  if (Array.isArray(first))
+    return first.map(x => escapeIdentifier(transform.column.to ? transform.column.to(x) : x)).join(',')
+
+  let value
+  const columns = rest.length ? rest.flat() : Object.keys(first)
+  return columns.map(x => {
+    value = first[x]
+    return (
+      value instanceof Query ? value.strings[0] :
+      value instanceof Identifier ? value.value :
+      handleValue(value, parameters, types)
+    ) + ' as ' + escapeIdentifier(transform.column.to ? transform.column.to(x) : x)
+  }).join(',')
+}
+
 const builders = Object.entries({
   values,
   in: values,
+  select,
+  returning: select,
 
   update(first, rest, parameters, types, transform) {
     return (rest.length ? rest.flat() : Object.keys(first)).map(x =>
@@ -121,23 +140,6 @@ const builders = Object.entries({
     )
   },
 
-  select(first, rest, parameters, types, transform) {
-    typeof first === 'string' && (first = [first].concat(rest))
-    if (Array.isArray(first))
-      return first.map(x => escapeIdentifier(transform.column.to ? transform.column.to(x) : x)).join(',')
-
-    let value
-    const columns = rest.length ? rest.flat() : Object.keys(first)
-    return columns.map(x => {
-      value = first[x]
-      return (
-        value instanceof Query ? value.strings[0] :
-        value instanceof Identifier ? value.value :
-        handleValue(value, parameters, types)
-      ) + ' as ' + escapeIdentifier(transform.column.to ? transform.column.to(x) : x)
-    }).join(',')
-  },
-
   insert(first, rest, parameters, types, transform) {
     const columns = rest.length ? rest.flat() : Object.keys(Array.isArray(first) ? first[0] : first)
     return '(' + columns.map(x =>
diff --git a/tests/index.js b/tests/index.js
@@ -961,6 +961,15 @@ t('dynamic select array', async() => {
   return ['yay', (await sql`select ${ sql(['a', 'b']) } from test`)[0].b, await sql`drop table test`]
 })
 
+t('dynamic returning array', async() => {
+  await sql`create table test (a int, b text)`
+  return [
+    'yay',
+    (await sql`insert into test (a, b) values (42, 'yay') returning ${ sql(['a', 'b']) }`)[0].b,
+    await sql`drop table test`
+  ]
+})
+
 t('dynamic select args', async() => {
   await sql`create table test (a int, b text)`
   await sql`insert into test (a, b) values (42, 'yay')`
