Implement better query helpers

porsager · porsager · commit fcbdcab539ba · 2019-12-09T08:16:07.000+01:00
diff --git a/lib/index.js b/lib/index.js
@@ -13,9 +13,21 @@ import {
   toCamel,
   toKebab,
   errors,
+  escape,
   types
 } from './types.js'
 
+const notPromise = {
+  P: {},
+  finally: notTagged,
+  then: notTagged,
+  catch: notTagged
+}
+
+function notTagged() {
+  throw errors.generic({ message: 'Query not called as a tagged template literal', code: 'NOT_TAGGED_CALL' })
+}
+
 Object.assign(Postgres, {
   toPascal,
   toCamel,
@@ -37,6 +49,8 @@ export default function Postgres(url, options) {
       , listeners = {}
       , typeArrayMap = {}
       , files = {}
+      , isInsert = /(^|\))\s*insert\s*into\s/i
+      , isSelect = /(^|\))\s*select\s/i
 
   function postgres(xs, ...args) {
     return query({}, getConnection(), xs, args)
@@ -138,7 +152,7 @@ export default function Postgres(url, options) {
 
   function query(query, connection, xs, args) {
     if (!query.raw && (!Array.isArray(xs) || !Array.isArray(xs.raw)))
-      throw errors.generic({ message: 'Query not called as a tagged template literal', code: 'NOT_TAGGED_CALL' })
+      return nested(xs, args)
 
     const promise = new Promise((resolve, reject) => {
       query.resolve = resolve
@@ -155,6 +169,13 @@ export default function Postgres(url, options) {
     return promise
   }
 
+  function nested(first, rest) {
+    const o = Object.create(notPromise)
+    o.first = first
+    o.rest = rest
+    return o
+  }
+
   function send(connection, query, xs, args) {
     connection
       ? connection.send(query, query.raw ? parseRaw(xs, args) : parse(xs, args))
@@ -206,22 +227,6 @@ export default function Postgres(url, options) {
     }
   }
 
-  function rows(rows, ...args) {
-    return {
-      rows: typeof args[0] === 'string'
-        ? rows.map(x => Array.isArray(x) ? x : args.map(a => x[a])) // pluck
-        : typeof args[0] === 'function'
-          ? rows.map(x => args[0](x)) // map
-          : rows
-    }
-  }
-
-  function row(row, ...args) {
-    return {
-      row: args.map(a => row[a])
-    }
-  }
-
   function array(value) {
     return {
       type: inferType(value) || 25,
@@ -267,9 +272,7 @@ export default function Postgres(url, options) {
       unsafe,
       array,
       file,
-      json,
-      rows,
-      row
+      json
     })
 
     function notify(channel, payload) {
@@ -380,13 +383,9 @@ export default function Postgres(url, options) {
 
     for (let i = 1; i < xs.length; i++) {
       arg = args[i - 1]
-      str += (!arg
-        ? parseValue(arg, xargs, types)
-        : arg.rows
-          ? parseRows(arg.rows, xargs, types)
-          : arg.row
-            ? parseRow(arg.row, xargs, types)
-            : parseValue(arg, xargs, types)
+      str += (arg && arg.P === notPromise.P
+        ? parseHelper(str, arg, xargs, types)
+        : parseValue(arg, xargs, types)
       ) + xs[i]
     }
 
@@ -397,19 +396,66 @@ export default function Postgres(url, options) {
     }
   }
 
-  function parseRows(rows, xargs, types) {
-    xargs.dynamic = true
-    return rows.map(row => parseRow(row, xargs, types)).join(',')
+  function parseHelper(str, { first, rest }, xargs, types) {
+    if (first !== null && typeof first === 'object' && typeof first[0] !== 'string') {
+      if (isInsert.test(str))
+        return insertHelper(first, rest, xargs, types)
+      else if(isSelect.test(str))
+        return selectHelper(first, rest, xargs, types)
+      else if(!Array.isArray(first))
+        return equalsHelper(first, rest, xargs, types)
+    }
+
+    return escapeHelper(Array.isArray(first) ? first : [first].concat(rest))
   }
 
-  function parseRow(row, xargs, types) {
-    return '(' + row.map(x => parseValue(x, xargs, types)).join(',') + ')'
+  function selectHelper(first, columns, xargs, types) {
+    return Object.entries(first).reduce((acc, [k, v]) =>
+      acc + (!columns.length || columns.indexOf(k) > -1
+        ? (acc ? ', ' : '') + parseValue(v, xargs, types) + ' as ' + escape(k)
+        : ''
+      )
+    , '')
+  }
+
+  function insertHelper(first, columns, xargs, types) {
+    first = Array.isArray(first) ? first : [first]
+    return '(' + escapeHelper(Object.keys(first[0])) + ') values ' +
+      first.map(row =>
+        '(' + Object.entries(row).reduce((acc, [k, v]) =>
+          acc + (!columns.length || columns.indexOf(k) > -1
+            ? (acc ? ', ' : '') + parseValue(v, xargs, types)
+            : ''
+          )
+        , '') + ')'
+      ).join(', ')
+  }
+
+  function equalsHelper(first, columns, xargs, types) {
+    return Object.entries(first).reduce((acc, [k, v]) =>
+      acc + (!columns.length || columns.indexOf(k) > -1
+        ? (acc ? ', ' : '') + escape(k) + ' = ' + parseValue(v, xargs, types)
+        : ''
+      )
+    , '')
+  }
+
+  function escapeHelper(xs) {
+    return xs.reduce((acc, x) => acc + (acc ? ', ' : '') + escape(x), '')
   }
 
   function parseValue(x, xargs, types) {
+    return Array.isArray(x)
+      ? x.reduce((acc, x) => addValue(x, xargs, types), '')
+      : addValue(x, xargs, types)
+  }
+
+  function addValue(x, xargs, types) {
     const type = getType(x)
-    types.push(type.type)
-    return '$' + xargs.push(type)
+        , i = types.push(type-type)
+
+    xargs.push(type)
+    return '$' + i
   }
 
   function getType(x) {
diff --git a/lib/types.js b/lib/types.js
@@ -65,6 +65,28 @@ const type = {
   boolean: 16
 }
 
+export function escape(str) {
+  let result = ''
+  let q = str[0] < 10 || str[0] === '$'
+  let last = 0
+  let char
+  let code
+
+  for (let i = 0; i < str.length; i++) {
+    const char = str[i]
+    const code = char.charCodeAt(0)
+    if (str[i] === '"') {
+      q = true
+      result += str.slice(last, i) + '"'
+      last = i
+    } else if (code === 96 || (code !== 36 && code <= 47) || (code >= 58 && code <= 64) || (code >= 91 && code <= 94) || (code >= 123 && code <= 128)) {
+      q = true
+    }
+  }
+
+  return (q ? '"' : '') + (q ? result + str.slice(last, str.length) : str) + (q ? '"' : '')
+}
+
 export function inferType(x) {
   return x.type || (x instanceof Date
     ? 1184
