Fixed issue with chunked rowsets when data packets are coalesced

Author: gionatamettifogo

Dockerfile

@@ -1,3 +1,4 @@
+# https://bun.sh/guides/ecosystem/docker
 # docker build -t sqlitcloud-gateway .
 # docker run -p 8090:8090 -p 4000:4000 --name sqlitecloud-gateway sqlitecloud-gateway       
 

package-lock.json

@@ -61,6 +61,7 @@
     "@typescript-eslint/eslint-plugin": "^4.22.0",
     "@typescript-eslint/parser": "^4.22.0",
     "dotenv": "^16.4.1",
+    "dotenv-cli": "^7.3.0",
     "eslint": "^7.32.0",
     "eslint-config-prettier": "^8.10.0",
     "eslint-plugin-node": "^11.1.0",

package.json

@@ -91,7 +91,7 @@ <h2>Messages:</h2>
         const sql = messageInput.value
         const startTime = Date.now()
         socket.emit('v1/sql', { sql }, response => {
-          console.debug(`sql - sql: ${sql}, elapsed: ${Date.now() - startTime}`, response)
+          console.debug(`sql - sql: ${sql} (${Date.now() - startTime}ms)`, response)
           appendMessage(`sql | ${JSON.stringify(response)}`)
         })
       })
@@ -100,11 +100,10 @@ <h2>Messages:</h2>
         if (!socket) {
           socket = setupSocket()
         }
-        socket.emit('v1/info', { }, response => {
-          appendMessage(`info | rss: ${(response.data.rss/1024/1024).toFixed(2)}mb, ${JSON.stringify(response)}`)
+        socket.emit('v1/info', {}, response => {
+          appendMessage(`info | rss: ${(response.data.rss / 1024 / 1024).toFixed(2)}mb, ${JSON.stringify(response)}`)
         })
       })
-
-</script>
+    </script>
   </body>
 </html>

public/index.html

@@ -4,7 +4,17 @@
 
 import { SQLiteCloudConfig, SQLiteCloudError, ErrorCallback, ResultsCallback } from './types'
 import { SQLiteCloudConnection } from './connection'
-import { formatCommand, hasCommandLength, parseCommandLength, popData, decompressBuffer, CMD_COMPRESSED, CMD_ROWSET_CHUNK } from './protocol'
+import {
+  formatCommand,
+  hasCommandLength,
+  parseCommandLength,
+  popData,
+  decompressBuffer,
+  CMD_COMPRESSED,
+  CMD_ROWSET_CHUNK,
+  bufferEndsWith,
+  ROWSET_CHUNKS_END
+} from './protocol'
 import { getInitializationCommands, anonimizeError, anonimizeCommand } from './utilities'
 import { parseRowsetChunks } from './protocol'
 
@@ -181,19 +191,15 @@ export class SQLiteCloudTlsConnection extends SQLiteCloudConnection {
               const { data } = popData(buffer)
               finish(null, data)
             } else {
-              // @ts-expect-error
               // check if rowset received the ending chunk
-              if (data.subarray(data.indexOf(' ') + 1, data.length).toString() === '0 0 0 ') {
+              if (bufferEndsWith(buffer, ROWSET_CHUNKS_END)) {
+                rowsetChunks.push(buffer)
                 const parsedData = parseRowsetChunks(rowsetChunks)
                 finish?.call(this, null, parsedData)
               } else {
                 // no ending string? ask server for another chunk
                 rowsetChunks.push(buffer)
                 buffer = Buffer.alloc(0)
-
-                // no longer need to ack the server
-                // const okCommand = formatCommand('OK')
-                // this.socket?.write(okCommand)
               }
             }
           }

src/drivers/connection-tls.ts

@@ -4,6 +4,7 @@
 
 import { SQLiteCloudError, type SQLCloudRowsetMetadata, type SQLiteCloudDataTypes } from './types'
 import { SQLiteCloudRowset } from './rowset'
+import { debug } from 'console'
 
 const lz4 = require('lz4js')
 
@@ -28,6 +29,10 @@ export const CMD_ARRAY = '='
 // const CMD_PUBSUB = '|'
 // const CMD_RECONNECT = '@'
 
+// To mark the end of the Rowset, the special string /LEN 0 0 0   is sent (LEN is always 6 in this case)
+// https://github.com/sqlitecloud/sdk/blob/master/PROTOCOL.md#scsp-rowset-chunk
+export const ROWSET_CHUNKS_END = '/6 0 0 0 '
+
 //
 // utility functions
 //
@@ -190,22 +195,34 @@ function parseRowset(buffer: Buffer, spaceIndex: number): SQLiteCloudRowset {
   return new SQLiteCloudRowset(metadata, data)
 }
 
+export function bufferStartsWith(buffer: Buffer, prefix: string): boolean {
+  return buffer.length >= prefix.length && buffer.subarray(0, prefix.length).toString('utf8') === prefix
+}
+
+export function bufferEndsWith(buffer: Buffer, suffix: string): boolean {
+  return buffer.length >= suffix.length && buffer.subarray(buffer.length - suffix.length, buffer.length).toString('utf8') === suffix
+}
+
 /**
  * Parse a chunk of a chunked rowset command, eg:
  * *LEN 0:VERS NROWS NCOLS DATA
+ * @see https://github.com/sqlitecloud/sdk/blob/master/PROTOCOL.md#scsp-rowset-chunk
  */
 export function parseRowsetChunks(buffers: Buffer[]): SQLiteCloudRowset {
+  let buffer = Buffer.concat(buffers)
+  if (!bufferStartsWith(buffer, CMD_ROWSET_CHUNK) || !bufferEndsWith(buffer, ROWSET_CHUNKS_END)) {
+    throw new Error('SQLiteCloudConnection.parseRowsetChunks - invalid chunks buffer')
+  }
+
   let metadata: SQLCloudRowsetMetadata = { version: 1, numberOfColumns: 0, numberOfRows: 0, columns: [] }
   const data: any[] = []
 
-  for (let i = 0; i < buffers.length; i++) {
-    let buffer: Buffer = buffers[i]
-
-    // validate and skip data type
-    const dataType = buffer.subarray(0, 1).toString()
-    console.assert(dataType === CMD_ROWSET_CHUNK)
-    buffer = buffer.subarray(buffer.indexOf(' ') + 1)
+  // validate and skip data type
+  const dataType = buffer.subarray(0, 1).toString()
+  console.assert(dataType === CMD_ROWSET_CHUNK)
+  buffer = buffer.subarray(buffer.indexOf(' ') + 1)
 
+  while (true) {
     // chunk header, eg: 0:VERS NROWS NCOLS
     const { index: chunkIndex, metadata: chunkMetadata, fwdBuffer } = parseRowsetHeader(buffer)
     buffer = fwdBuffer
@@ -224,10 +241,17 @@ export function parseRowsetChunks(buffers: Buffer[]): SQLiteCloudRowset {
       data.push(itemData)
       buffer = fwdBuffer
     }
+
+    // no more chunks?
+    if (bufferStartsWith(buffer, ROWSET_CHUNKS_END)) {
+      break
+    }
   }
 
-  console.assert(data && data.length === metadata.numberOfRows * metadata.numberOfColumns, 'SQLiteCloudConnection.parseRowsetChunks - invalid rowset data')
-  return new SQLiteCloudRowset(metadata, data)
+  console.assert(data && data.length === metadata.numberOfRows * metadata.numberOfColumns, 'parseRowsetChunks - invalid rowset data')
+  const rowset = new SQLiteCloudRowset(metadata, data)
+  // console.debug(`parseRowsetChunks - ${rowset.numberOfRows} rows, ${rowset.numberOfColumns} columns`)
+  return rowset
 }
 
 /** Pop one or more space separated integers from beginning of buffer, move buffer forward */

src/drivers/protocol.ts

@@ -8,8 +8,9 @@
 import { SQLiteCloudError } from '../drivers/types'
 import { SQLiteCloudBunConnection } from './connection-bun'
 
-const CHINOOK_DATABASE_URL = process.env['CHINOOK_DATABASE_URL'] as string
+let CHINOOK_DATABASE_URL = process.env['CHINOOK_DATABASE_URL'] as string
 console.assert(CHINOOK_DATABASE_URL, 'CHINOOK_DATABASE_URL is required')
+CHINOOK_DATABASE_URL = 'sqlitecloud://admin:uN3ARhdcKQ@oggdnp3zm.sqlite.cloud:8860/chinook.db?verbose=true'
 
 import { expect, test, describe, beforeEach, afterEach } from 'bun:test'
 
@@ -177,7 +178,7 @@ describe('SQLiteCloudBunConnection', () => {
       expect(results.columnsNames).toEqual(['key', 'value'])
     })
 
-    test('should test chunked rowset', async () => {
+    test.only('should test chunked rowset', async () => {
       const results = await sendCommands(chinookConnection, 'TEST ROWSET_CHUNK')
       expect(results.numberOfRows).toBe(147)
       expect(results.numberOfColumns).toBe(1)

src/gateway/connection-bun.test.ts

@@ -13,7 +13,9 @@ import {
   decompressBuffer,
   parseRowsetChunks,
   CMD_COMPRESSED,
-  CMD_ROWSET_CHUNK
+  CMD_ROWSET_CHUNK,
+  bufferEndsWith,
+  ROWSET_CHUNKS_END
 } from '../drivers/protocol'
 import type { Socket } from 'bun'
 
@@ -125,9 +127,8 @@ export class SQLiteCloudBunConnection extends SQLiteCloudConnection {
 
     // reset buffer and rowset chunks, define response callback
     this.buffer = Buffer.alloc(0)
-    this.rowsetChunks = []
-    this.processCallback = callback
     this.startedOn = new Date()
+    this.processCallback = callback
 
     // compose commands following SCPC protocol
     const formattedCommands = formatCommand(commands)
@@ -144,9 +145,8 @@ export class SQLiteCloudBunConnection extends SQLiteCloudConnection {
   // onData is called when data is received, it will process the data until all data is retrieved for a response
   // when response is complete or there's an error, finish is called to call the results callback set by processCommands...
 
-  // buffer to store incoming data
+  // buffer to accumulate incoming data until an whole command is received and can be parsed
   private buffer: Buffer = Buffer.alloc(0)
-  private rowsetChunks: Buffer[] = []
   private startedOn: Date = new Date()
 
   // callback to be called when a command is finished processing
@@ -155,14 +155,16 @@ export class SQLiteCloudBunConnection extends SQLiteCloudConnection {
   /** Handles data received in response to an outbound command sent by processCommands */
   private processCommandsData(socket: Socket<any>, data: Buffer) {
     try {
-      // on first ondata event, dataType is read from data, on subsequent ondata event, is read from buffer that is the concatanations of data received on each ondata event
-      let dataType = this.buffer.length === 0 ? data.subarray(0, 1).toString() : this.buffer.subarray(0, 1).toString('utf8')
-      this.buffer = Buffer.concat([this.buffer, data])
-      const commandLength = hasCommandLength(dataType)
+      // append data to buffer as it arrives
+      if (data.length && data.length > 0) {
+        this.buffer = Buffer.concat([this.buffer, data])
+      }
 
-      if (commandLength) {
+      let dataType = this.buffer?.subarray(0, 1).toString()
+      if (hasCommandLength(dataType)) {
         const commandLength = parseCommandLength(this.buffer)
         const hasReceivedEntireCommand = this.buffer.length - this.buffer.indexOf(' ') - 1 >= commandLength ? true : false
+        console.debug(`dataType: ${dataType} buffer.length: ${this.buffer.length}, commandLenght: ${commandLength}`)
         if (hasReceivedEntireCommand) {
           if (this.config?.verbose) {
             let bufferString = this.buffer.toString('utf8')
@@ -182,18 +184,10 @@ export class SQLiteCloudBunConnection extends SQLiteCloudConnection {
             const { data } = popData(this.buffer)
             this.processCommandsFinish?.call(this, null, data)
           } else {
-            // check if rowset received the ending chunk
-            if (data.subarray(data.indexOf(' ') + 1, data.length).toString() === '0 0 0 ') {
-              const parsedData = parseRowsetChunks(this.rowsetChunks)
+            // check if rowset received the ending chunk in which case it can be unpacked
+            if (bufferEndsWith(this.buffer, ROWSET_CHUNKS_END)) {
+              const parsedData = parseRowsetChunks([this.buffer])
               this.processCommandsFinish?.call(this, null, parsedData)
-            } else {
-              // no ending string? ask server for another chunk
-              this.rowsetChunks.push(this.buffer)
-              this.buffer = Buffer.alloc(0)
-
-              // no longer need to ack the server
-              // const okCommand = formatCommand('OK')
-              // this.socket?.write(okCommand)
             }
           }
         }
@@ -221,6 +215,7 @@ export class SQLiteCloudBunConnection extends SQLiteCloudConnection {
       // console.debug('BunTransport.finish - result', result)
     }
     if (this.processCallback) {
+      // console.error(`SQLiteCloudBunConnection.processCommandsFinish - error:${error}, result: ${result}`, error, result)
       this.processCallback(error, result)
     }
   }

src/gateway/connection-bun.ts

@@ -89,7 +89,9 @@ describe('connection-tls', () => {
     })
 
     it('should connect with insecure connection string', done => {
-      if (INSECURE_DATABASE_URL) {
+      if (!INSECURE_DATABASE_URL) {
+        done()
+      } else {
         expect(INSECURE_DATABASE_URL).toBeDefined()
         const conn = new SQLiteCloudTlsConnection(INSECURE_DATABASE_URL, error => {
           expect(error).toBeNull()
@@ -295,29 +297,7 @@ describe('connection-tls', () => {
           expect(results[1]['key']).toBe('INDEXED')
           expect(results[2]['key']).toBe('INDEX')
           expect(results[3]['key']).toBe('DESC')
-
-          database.close()
-          done()
-        })
-      },
-      LONG_TIMEOUT
-    )
-
-    it(
-      'should test chunked rowset via ',
-      done => {
-        // this operation sends 150 packets, so we need to increase the timeout
-        const database = getChinookTlsConnection(undefined, { timeout: 60 * 1000 })
-        database.sendCommands('TEST ROWSET_CHUNK', (error, results) => {
-          expect(error).toBeNull()
-          expect(results.numberOfRows).toBe(147)
-          expect(results.numberOfColumns).toBe(1)
-          expect(results.columnsNames).toEqual(['key'])
-
-          expect(results[0]['key']).toBe('REINDEX')
-          expect(results[1]['key']).toBe('INDEXED')
-          expect(results[2]['key']).toBe('INDEX')
-          expect(results[3]['key']).toBe('DESC')
+          expect(results[146]['key']).toBe('PRIMARY')
 
           database.close()
           done()

test/connection-tls.test.ts

@@ -0,0 +1,31 @@
+//
+// protocol.test.ts
+//
+
+import { parseRowsetChunks } from '../src/drivers/protocol'
+
+// response sent by the server when we TEST ROWSET_CHUNK
+const CHUNKED_RESPONSE = Buffer.from(
+  '/24 1:1 1 1 +3 key+7 REINDEX/18 2:1 1 1 +7 INDEXED/16 3:1 1 1 +5 INDEX/15 4:1 1 1 +4 DESC/17 5:1 1 1 +6 ESCAPE/15 6:1 1 1 +4 EACH/16 7:1 1 1 +5 CHECK/14 8:1 1 1 +3 KEY/17 9:1 1 1 +6 BEFORE/19 10:1 1 1 +7 FOREIGN/15 11:1 1 1 +3 FOR/18 12:1 1 1 +6 IGNORE/18 13:1 1 1 +6 REGEXP/19 14:1 1 1 +7 EXPLAIN/19 15:1 1 1 +7 INSTEAD/15 16:1 1 1 +3 ADD/20 17:1 1 1 +8 DATABASE/14 18:1 1 1 +2 AS/18 19:1 1 1 +6 SELECT/17 20:1 1 1 +5 TABLE/16 21:1 1 1 +4 LEFT/16 22:1 1 1 +4 THEN/15 23:1 1 1 +3 END/23 24:1 1 1 +10 DEFERRABLE/16 25:1 1 1 +4 ELSE/19 26:1 1 1 +7 EXCLUDE/18 27:1 1 1 +6 DELETE/21 28:1 1 1 +9 TEMPORARY/16 29:1 1 1 +4 TEMP/14 30:1 1 1 +2 OR/18 31:1 1 1 +6 ISNULL/17 32:1 1 1 +5 NULLS/21 33:1 1 1 +9 SAVEPOINT/21 34:1 1 1 +9 INTERSECT/16 35:1 1 1 +4 TIES/19 36:1 1 1 +7 NOTNULL/15 37:1 1 1 +3 NOT/14 38:1 1 1 +2 NO/16 39:1 1 1 +4 NULL/16 40:1 1 1 +4 LIKE/18 41:1 1 1 +6 EXCEPT/24 42:1 1 1 +11 TRANSACTION/18 43:1 1 1 +6 ACTION/14 44:1 1 1 +2 ON/19 45:1 1 1 +7 NATURAL/17 46:1 1 1 +5 ALTER/17 47:1 1 1 +5 RAISE/21 48:1 1 1 +9 EXCLUSIVE/18 49:1 1 1 +6 EXISTS/23 50:1 1 1 +10 CONSTRAINT/16 51:1 1 1 +4 INTO/18 52:1 1 1 +6 OFFSET/14 53:1 1 1 +2 OF/15 54:1 1 1 +3 SET/19 55:1 1 1 +7 TRIGGER/17 56:1 1 1 +5 RANGE/21 57:1 1 1 +9 GENERATED/18 58:1 1 1 +6 DETACH/18 59:1 1 1 +6 HAVING/16 60:1 1 1 +4 GLOB/17 61:1 1 1 +5 BEGIN/17 62:1 1 1 +5 INNER/23 63:1 1 1 +10 REFERENCES/18 64:1 1 1 +6 UNIQUE/17 65:1 1 1 +5 QUERY/19 66:1 1 1 +7 WITHOUT/16 67:1 1 1 +4 WITH/17 68:1 1 1 +5 OUTER/19 69:1 1 1 +7 RELEASE/18 70:1 1 1 +6 ATTACH/19 71:1 1 1 +7 BETWEEN/19 72:1 1 1 +7 NOTHING/18 73:1 1 1 +6 GROUPS/17 74:1 1 1 +5 GROUP/19 75:1 1 1 +7 CASCADE/15 76:1 1 1 +3 ASC/19 77:1 1 1 +7 DEFAULT/16 78:1 1 1 +4 CASE/19 79:1 1 1 +7 COLLATE/18 80:1 1 1 +6 CREATE/25 81:1 1 1 +12 CURRENT_DATE/21 82:1 1 1 +9 IMMEDIATE/16 83:1 1 1 +4 JOIN/18 84:1 1 1 +6 INSERT/17 85:1 1 1 +5 MATCH/16 86:1 1 1 +4 PLAN/19 87:1 1 1 +7 ANALYZE/18 88:1 1 1 +6 PRAGMA/25 89:1 1 1 +12 MATERIALIZED/20 90:1 1 1 +8 DEFERRED/20 91:1 1 1 +8 DISTINCT/14 92:1 1 1 +2 IS/18 93:1 1 1 +6 UPDATE/18 94:1 1 1 +6 VALUES/19 95:1 1 1 +7 VIRTUAL/18 96:1 1 1 +6 ALWAYS/16 97:1 1 1 +4 WHEN/17 98:1 1 1 +5 WHERE/21 99:1 1 1 +9 RECURSIVE/18 100:1 1 1 +5 ABORT/18 101:1 1 1 +5 AFTER/19 102:1 1 1 +6 RENAME/16 103:1 1 1 +3 AND/17 104:1 1 1 +4 DROP/22 105:1 1 1 +9 PARTITION/27 106:1 1 1 +13 AUTOINCREMENT/15 107:1 1 1 +2 TO/15 108:1 1 1 +2 IN/17 109:1 1 1 +4 CAST/19 110:1 1 1 +6 COLUMN/19 111:1 1 1 +6 COMMIT/21 112:1 1 1 +8 CONFLICT/18 113:1 1 1 +5 CROSS/31 114:1 1 1 +17 CURRENT_TIMESTAMP/26 115:1 1 1 +12 CURRENT_TIME/20 116:1 1 1 +7 CURRENT/22 117:1 1 1 +9 PRECEDING/17 118:1 1 1 +4 FAIL/17 119:1 1 1 +4 LAST/19 120:1 1 1 +6 FILTER/20 121:1 1 1 +7 REPLACE/18 122:1 1 1 +5 FIRST/22 123:1 1 1 +9 FOLLOWING/17 124:1 1 1 +4 FROM/17 125:1 1 1 +4 FULL/18 126:1 1 1 +5 LIMIT/15 127:1 1 1 +2 IF/18 128:1 1 1 +5 ORDER/21 129:1 1 1 +8 RESTRICT/19 130:1 1 1 +6 OTHERS/17 131:1 1 1 +4 OVER/22 132:1 1 1 +9 RETURNING/18 133:1 1 1 +5 RIGHT/21 134:1 1 1 +8 ROLLBACK/17 135:1 1 1 +4 ROWS/16 136:1 1 1 +3 ROW/22 137:1 1 1 +9 UNBOUNDED/18 138:1 1 1 +5 UNION/18 139:1 1 1 +5 USING/19 140:1 1 1 +6 VACUUM/17 141:1 1 1 +4 VIEW/19 142:1 1 1 +6 WINDOW/15 143:1 1 1 +2 DO/15 144:1 1 1 +2 BY/22 145:1 1 1 +9 INITIALLY/16 146:1 1 1 +3 ALL/20 147:1 1 1 +7 PRIMARY/6 0 0 0 '
+)
+
+describe('parseRowsetChunks', () => {
+  it('should extract rowset from single buffer', () => {
+    const rowset = parseRowsetChunks([CHUNKED_RESPONSE])
+    expect(rowset.length).toBe(147)
+    expect(rowset[0]['key']).toBe('REINDEX')
+    expect(rowset[146]['key']).toBe('PRIMARY')
+  })
+
+  it('should extract rowset from segmented buffers', () => {
+    // split CHUNKED_RESPONSE into 3 random sized buffers
+    const buffer1 = CHUNKED_RESPONSE.slice(0, 100)
+    const buffer2 = CHUNKED_RESPONSE.slice(100, 200)
+    const buffer3 = CHUNKED_RESPONSE.slice(200)
+
+    const rowset = parseRowsetChunks([buffer1, buffer2, buffer3])
+    expect(rowset.length).toBe(147)
+    expect(rowset[0]['key']).toBe('REINDEX')
+    expect(rowset[146]['key']).toBe('PRIMARY')
+  })
+})