Fix parser for rows spanning several chunks / compatibility with pg-copy-stream 3.0+

Author: jeromew

lib/parser.js

@@ -11,6 +11,14 @@ module.exports = function (txt, options) {
 const { Transform } = require('stream')
 const BP = require('bufferput')
 const { parse } = require('./pg_types')
+const BufferList = require('bl/BufferList')
+
+const PG_HEADER = 0
+const PG_ROW_START = 1
+const PG_FIELD_START = 2
+const PG_FIELD_DATA = 3
+const PG_FIELD_END = 4
+const PG_TRAILER = 5
 
 class CopyStream extends Transform {
   constructor(options) {
@@ -24,75 +32,96 @@ class CopyStream extends Transform {
       .word32be(0)
       .buffer()
 
-    this.COPYTrailer = Buffer.from([0xff, 0xff])
-
-    this._headerReceived = false
-    this._trailerReceived = false
-    this._remainder = false
+    this._buffer = new BufferList()
+    this._state = PG_HEADER
+    this._row = null
+    this._fieldCount = null
+    this._fieldIndex = null
+    this._fieldLength = null
+    this._fieldBuffer = null
 
     this.mapping = options.mapping || false
   }
 
   _transform(chunk, enc, cb) {
-    if (this._remainder && chunk) {
-      chunk = Buffer.concat([this._remainder, chunk])
-    }
+    this._buffer.append(chunk)
+    while (this._buffer.length > 0) {
+      if (PG_HEADER === this._state) {
+        if (this._buffer.length < this.COPYHeaderFull.length) break
+        if (!this.COPYHeaderFull.equals(this._buffer.slice(0, this.COPYHeaderFull.length))) {
+          return cb(new Error('COPY BINARY Header mismatch'))
+        }
+        this._buffer.consume(this.COPYHeaderFull.length)
+        this._state = PG_ROW_START
+      }
 
-    let offset = 0
-    if (!this._headerReceived && chunk.length >= this.COPYHeaderFull.length) {
-      if (this.COPYHeaderFull.equals(chunk.slice(0, this.COPYHeaderFull.length))) {
-        this._headerReceived = true
-        offset += this.COPYHeaderFull.length
+      if (PG_ROW_START === this._state) {
+        if (this._buffer.length < 2) break
+        this._fieldCount = this._buffer.readUInt16BE(0)
+        this._buffer.consume(2)
+        const UInt16_0xffff = 65535
+        if (this._fieldCount === UInt16_0xffff) {
+          this._state = PG_TRAILER
+        } else {
+          this._row = this.mapping ? {} : []
+          this._state = PG_FIELD_START
+          this._fieldIndex = -1
+        }
       }
-    }
 
-    // Copy-out mode (data transfer from the server) is initiated when the backend executes a COPY TO STDOUT SQL statement.
-    // The backend sends a CopyOutResponse message to the frontend, followed by zero or more CopyData messages (always one per row)
-    const UInt16Len = 2
-    while (this._headerReceived && chunk.length - offset >= UInt16Len) {
-      const fieldCount = chunk.readUInt16BE(offset)
-      offset += 2
-      const UInt32Len = 4
-      const UInt16_0xff = 65535
-      const UInt32_0xffffffff = 4294967295
-      if (fieldCount === UInt16_0xff) {
-        this._trailerReceived = true
+      if (PG_TRAILER === this._state) {
         this.push(null)
+        this._row = null
+        this._fieldBuffer = null
         return cb()
       }
-      const fields = this.mapping ? {} : []
-      for (let i = 0; i < fieldCount; i++) {
-        let v
-        const fieldLen = chunk.readUInt32BE(offset)
-        offset += UInt32Len
-        if (fieldLen === UInt32_0xffffffff) {
-          v = null
+
+      if (PG_FIELD_START === this._state) {
+        if (this._buffer.length < 4) break
+        this._fieldIndex++
+        this._fieldLength = this._buffer.readUInt32BE(0)
+        this._buffer.consume(4)
+        const UInt32_0xffffffff = 4294967295 /* Magic value for NULL */
+        if (this._fieldLength === UInt32_0xffffffff) {
+          this._fieldBuffer = null
+          this._fieldLength = 0
+          this._state = PG_FIELD_END
         } else {
-          v = chunk.slice(offset, offset + fieldLen)
-          if (this.mapping) {
-            v = parse(v, this.mapping[i].type)
-          }
-          offset += fieldLen
+          this._fieldBuffer = new BufferList()
+          this._state = PG_FIELD_DATA
+        }
+      }
+
+      if (PG_FIELD_DATA === this._state) {
+        if (this._buffer.length === 0) break
+        const bl = this._buffer.shallowSlice(0, this._fieldLength)
+        this._fieldBuffer.append(bl)
+        this._fieldLength -= bl.length
+        this._buffer.consume(bl.length)
+        if (this._fieldLength === 0) {
+          this._state = PG_FIELD_END
+        }
+      }
+
+      if (PG_FIELD_END === this._state) {
+        if (this._fieldBuffer && this.mapping) {
+          this._fieldBuffer = parse(this._fieldBuffer.slice(), this.mapping[this._fieldIndex].type)
         }
         if (this.mapping) {
-          fields[this.mapping[i].key] = v
+          this._row[this.mapping[this._fieldIndex].key] = this._fieldBuffer
         } else {
-          fields.push(v)
+          this._row.push(this._fieldBuffer)
         }
-      }
-      this.push(fields)
-    }
 
-    if (chunk.length - offset) {
-      const slice = chunk.slice(offset)
-      this._remainder = slice
-    } else {
-      this._remainder = false
+        this._state = PG_FIELD_START
+
+        if (this._fieldIndex === this._fieldCount - 1) {
+          this.push(this._row)
+          this._state = PG_ROW_START
+        }
+      }
     }
-    cb()
-  }
 
-  _flush(cb) {
     cb()
   }
 }