Fix query result streaming in @effect/sql-pg (#5758)

Co-authored-by: Tim Smart <[email protected]>

Author: rocwang

.changeset/many-spies-raise.md

@@ -0,0 +1,5 @@
+---
+"@effect/sql-pg": patch
+---
+
+Fix query result streaming in @effect/sql-pg

packages/sql-pg/src/PgClient.ts

@@ -289,10 +289,10 @@ export const make = (
         // eslint-disable-next-line @typescript-eslint/no-this-alias
         const self = this
         return Effect.gen(function*() {
-          const cursor = yield* Effect.acquireRelease(
-            Effect.sync(() => (self.pg ?? pool).query(new Cursor(sql, params as any))),
-            (cursor) => Effect.sync(() => cursor.close())
-          )
+          const scope = yield* Effect.scope
+          const client = self.pg ?? (yield* reserveRaw)
+          yield* Scope.addFinalizer(scope, Effect.promise(() => cursor.close()))
+          const cursor = client.query(new Cursor(sql, params as any))
           const pull = Effect.async<Chunk.Chunk<any>, Option.Option<SqlError>>((resume) => {
             cursor.read(128, (err, rows) => {
               if (err) {

packages/sql-pg/test/Client.test.ts

@@ -3,6 +3,8 @@ import * as SqlClient from "@effect/sql/SqlClient"
 import * as Statement from "@effect/sql/Statement"
 import { assert, expect, it } from "@effect/vitest"
 import { Effect, Redacted, String } from "effect"
+import * as Chunk from "effect/Chunk"
+import * as Stream from "effect/Stream"
 import * as TestServices from "effect/TestServices"
 import { parse as parsePgConnectionString } from "pg-connection-string"
 import { PgContainer } from "./utils.js"
@@ -237,6 +239,20 @@ it.layer(PgContainer.ClientLive, { timeout: "30 seconds" })("PgClient", (it) =>
       const rows = yield* sql<{ json: unknown }>`select ${{ testValue: 123 }}::jsonb as json`
       expect(rows[0].json).toEqual({ testValue: 123 })
     }))
+
+  it.effect("stream", () =>
+    Effect.gen(function*() {
+      const sql = yield* SqlClient.SqlClient
+      const rows = yield* sql`SELECT generate_series(1, 3)`.stream.pipe(
+        Stream.runCollect,
+        Effect.map(Chunk.toReadonlyArray)
+      )
+      expect(rows).toEqual([
+        { "generate_series": 1 },
+        { "generate_series": 2 },
+        { "generate_series": 3 }
+      ])
+    }))
 })
 
 it.layer(PgContainer.ClientTransformLive, { timeout: "30 seconds" })("PgClient transforms", (it) => {