From 7ab23820b65c587994e01952a7a1bfa0629c0375 Mon Sep 17 00:00:00 2001
From: Sam Willis <sam.willis@gmail.com>
Date: Tue, 8 Oct 2024 18:56:34 +0100
Subject: [PATCH] Add tests for large result sets

---
 packages/pglite/tests/large-query.test.ts | 80 +++++++++++++++++++++++
 1 file changed, 80 insertions(+)

diff --git a/packages/pglite/tests/large-query.test.ts b/packages/pglite/tests/large-query.test.ts
index f91bf9a96..64fa0a771 100644
--- a/packages/pglite/tests/large-query.test.ts
+++ b/packages/pglite/tests/large-query.test.ts
@@ -97,4 +97,84 @@ describe('large query', () => {
     // sleep for GC to collect
     await new Promise((resolve) => setTimeout(resolve, 100))
   })
+
+  it('select 10k rows ~ 10mb', async () => {
+    const pg = await PGlite.create()
+
+    await pg.exec(`
+      CREATE TABLE test (id SERIAL PRIMARY KEY, data TEXT);
+    `)
+
+    // 1kb
+    const value = 'a'.repeat(1000)
+
+    await pg.exec(`
+      INSERT INTO test (data)
+      SELECT 'Row ' || generate_series || '${value}'
+      FROM generate_series(1, 10000);
+    `)
+
+    const res = await pg.query<{ data: string }>(`
+      SELECT * FROM test;
+    `)
+
+    expect(res.rows.length).toBe(10000)
+
+    pg.close()
+
+    // sleep for GC to collect
+    await new Promise((resolve) => setTimeout(resolve, 100))
+  })
+
+  it('select 100k rows ~ 100mb', async () => {
+    const pg = await PGlite.create()
+
+    await pg.exec(`
+      CREATE TABLE test (id SERIAL PRIMARY KEY, data TEXT);
+    `)
+
+    // 1kb
+    const value = 'a'.repeat(1000)
+
+    await pg.exec(`
+      INSERT INTO test (data)
+      SELECT 'Row ' || generate_series || '${value}'
+      FROM generate_series(1, 100000);
+    `)
+
+    const res = await pg.query<{ data: string }>(`
+      SELECT * FROM test;
+    `)
+
+    expect(res.rows.length).toBe(100000)
+
+    pg.close()
+
+    // sleep for GC to collect
+    await new Promise((resolve) => setTimeout(resolve, 100))
+  })
+
+  it('select 1m rows ~ 1gb', async () => {
+    const pg = await PGlite.create()
+
+    await pg.exec(`
+      CREATE TABLE test (id SERIAL PRIMARY KEY, data TEXT);
+    `)
+
+    // 1kb
+    const value = 'a'.repeat(1000)
+
+    // This time only select from the series, not insert into the table
+    const res = await pg.query(`
+      SELECT 'Row ' || generate_series || '${value}'
+      FROM generate_series(1, 1000000);
+    `)
+
+    expect(res.rows.length).toBe(1000000)
+
+    pg.close()
+
+    // sleep for GC to collect
+    await new Promise((resolve) => setTimeout(resolve, 100))
+  })
 })