From 409c7fa7b932c7caa49e241d23dc8dd33720fc4a Mon Sep 17 00:00:00 2001
From: Sam Willis <sam.willis@gmail.com>
Date: Tue, 8 Oct 2024 17:37:11 +0100
Subject: [PATCH] Tweeks and tests (failing)

---
 packages/pglite/src/pglite.ts             |  10 ++-
 packages/pglite/tests/large-query.test.ts | 100 ++++++++++++++++++++++
 2 files changed, 108 insertions(+), 2 deletions(-)
 create mode 100644 packages/pglite/tests/large-query.test.ts

diff --git a/packages/pglite/src/pglite.ts b/packages/pglite/src/pglite.ts
index 2e60ae0d8..7ed3a28d3 100644
--- a/packages/pglite/src/pglite.ts
+++ b/packages/pglite/src/pglite.ts
@@ -31,6 +31,7 @@ const SOCKET_FILE = {
   OLOCK: '/tmp/pglite/.s.PGSQL.5432.lock.out',
   OUT: '/tmp/pglite/.s.PGSQL.5432.out',
 }
+const PAGE_SIZE = 8192
 
 export class PGlite
   extends BasePGlite
@@ -461,7 +462,10 @@ export class PGlite
         }
         const contents = new Uint8Array(buf)
         if (position >= contents.length) return 0
-        const size = Math.min(contents.length - position, length)
+        const size = Math.min(
+          PAGE_SIZE,
+          Math.min(contents.length - position, length),
+        )
         for (let i = 0; i < size; i++) {
           buffer[offset + i] = contents[position + i]
         }
@@ -517,7 +521,9 @@ export class PGlite
         _position: number,
       ) => {
         this.#queryOutChunks ??= []
-        this.#queryOutChunks.push(buffer.slice(offset, offset + length))
+        if (length > 0) {
+          this.#queryOutChunks.push(buffer.slice(offset, offset + length))
+        }
         return length
       },
       llseek: (_stream: any, _offset: number, _whence: number) => {
diff --git a/packages/pglite/tests/large-query.test.ts b/packages/pglite/tests/large-query.test.ts
new file mode 100644
index 000000000..f91bf9a96
--- /dev/null
+++ b/packages/pglite/tests/large-query.test.ts
@@ -0,0 +1,100 @@
+import { describe, it, expect } from 'vitest'
+import { PGlite } from '../dist/index.js'
+
+describe('large query', () => {
+  it('8192 byte wire message', async () => {
+    const pg = await PGlite.create()
+
+    await pg.exec('CREATE TABLE test (id SERIAL PRIMARY KEY, data TEXT);')
+
+    const value = 'a'.repeat(8150)
+    const sql = `INSERT INTO test (data) VALUES ('${value}');`
+
+    // 8192 page size, 6 bytes for the wire protocol header
+    expect(sql.length).toBe(8192 - 6)
+
+    await pg.exec(sql)
+
+    const res = await pg.query<{ data: string }>(`
+      SELECT * FROM test;
+    `)
+
+    expect(res.rows.length).toBe(1)
+    expect(res.rows[0].data).toBe(value)
+
+    pg.close()
+  })
+
+  it('8193 byte wire message', async () => {
+    const pg = await PGlite.create()
+
+    await pg.exec('CREATE TABLE test (id SERIAL PRIMARY KEY, data TEXT);')
+
+    const value = 'a'.repeat(8151)
+    const sql = `INSERT INTO test (data) VALUES ('${value}');`
+
+    // 1 more than 8192 page size, 6 bytes for the wire protocol header
+    expect(sql.length).toBe(8193 - 6)
+
+    await pg.exec(sql)
+
+    const res = await pg.query<{ data: string }>(`
+      SELECT * FROM test;
+    `)
+
+    expect(res.rows.length).toBe(1)
+    expect(res.rows[0].data).toBe(value)
+
+    pg.close()
+  })
+
+  it('1mb value in insert and select', async () => {
+    const pg = await PGlite.create()
+
+    await pg.exec(`
+      CREATE TABLE test (id SERIAL PRIMARY KEY, data TEXT);
+    `)
+
+    // 1mb value
+    const value = 'a'.repeat(1_000_000)
+
+    await pg.query('INSERT INTO test (data) VALUES ($1);', [value])
+
+    await pg.exec(`
+      INSERT INTO test (data) VALUES (${value});
+    `)
+
+    const res = await pg.query<{ data: string }>(`
+      SELECT * FROM test;
+    `)
+
+    expect(res.rows.length).toBe(2)
+    expect(res.rows[0].data).toBe(value)
+
+    pg.close()
+
+    // sleep for GC to collect
+    await new Promise((resolve) => setTimeout(resolve, 100))
+  })
+
+  it('1mb of SQL in query', async () => {
+    const pg = await PGlite.create()
+
+    await pg.exec(`
+      CREATE TABLE test (id SERIAL PRIMARY KEY, data TEXT);
+    `)
+
+    let sql = ''
+    for (let i = 0; i < 26316; i++) {
+      // 26316 * 38 = 1,000,008 bytes
+      sql += `INSERT INTO test (data) VALUES ('a');\n` // 38b statement
+    }
+
+    await pg.exec(sql)
+
+    pg.close()
+
+    // sleep for GC to collect
+    await new Promise((resolve) => setTimeout(resolve, 100))
+  })
+})