create pool methods to reserve and release connections for transaction use

Author: leafo

README.md

@@ -470,11 +470,47 @@ Additional pool-specific methods:
 
 * `pool:pool_size()` — Returns the current number of connections in the pool
 * `pool:active_connections()` — Returns the number of connections currently executing queries
+* `pool:reserve()` — Returns a raw `Postgres` connection for exclusive use (see below)
+* `pool:release(pg)` — Returns a reserved connection back to the pool (see below)
 
 > **Note:** `wait_for_notification` is not supported with `PostgresPool` because
 > notifications are tied to the specific socket connection that issued the
 > `LISTEN` command.
 
+### Transaction Support with reserve/release
+
+When you need to run a transaction across multiple queries, use `reserve()` to
+get exclusive access to a connection. This prevents other coroutines from using
+the same connection between your transaction queries.
+
+```lua
+local pg = pool:reserve()  -- Returns raw Postgres connection, marked as reserved
+pg:query("BEGIN")
+pg:query("INSERT INTO users (name) VALUES ('test')")
+pg:query("COMMIT")
+pool:release(pg)  -- Returns connection to pool
+```
+
+#### `pool:reserve()`
+
+Returns a raw `Postgres` connection object from the pool, marked as reserved.
+While reserved, the connection will not be used by `pool:query()` or other pool
+methods.
+
+* Returns: `Postgres` connection object on success
+* Returns: `nil, "not connected"` if pool has no connections
+* Returns: `nil, "pool exhausted, max_pool_size reached"` if at capacity
+
+#### `pool:release(pg)`
+
+Returns a reserved connection back to the pool for reuse.
+
+* Validates the connection belongs to this pool
+* If the connection is still in a transaction (`transaction_status` is `"T"` or `"E"`), automatically issues `ROLLBACK` to reset it to idle state
+* Returns: `true` on success
+* Returns: `nil, "connection not from this pool"` if connection belongs to a different pool
+* Returns: `nil, "connection not reserved"` if the connection was not reserved
+
 ## Extended and simple query protocols
 
 pgmoon will issue your query to the database server using either the simple or

pgmoon/pool.lua

@@ -14,7 +14,7 @@ do
       local _list_0 = self.pool
       for _index_0 = 1, #_list_0 do
         local pg = _list_0[_index_0]
-        if not (pg.busy) then
+        if not (pg.busy or pg.reserved) then
           return pg
         end
       end
@@ -125,6 +125,55 @@ do
     encode_bytea = Postgres.encode_bytea,
     decode_bytea = Postgres.decode_bytea,
     setup_hstore = Postgres.setup_hstore,
+    reserve = function(self)
+      if #self.pool == 0 then
+        return nil, "not connected"
+      end
+      local _list_0 = self.pool
+      for _index_0 = 1, #_list_0 do
+        local pg = _list_0[_index_0]
+        if not (pg.busy or pg.reserved) then
+          pg.reserved = true
+          return pg
+        end
+      end
+      if self.config.max_pool_size and #self.pool >= self.config.max_pool_size then
+        return nil, "pool exhausted, max_pool_size reached"
+      end
+      local pg = self:_create_instance()
+      local ok, err = pg:connect()
+      if not (ok) then
+        return nil, err
+      end
+      table.insert(self.pool, pg)
+      pg.reserved = true
+      return pg
+    end,
+    release = function(self, pg)
+      if not (pg.parent_pool == self) then
+        return nil, "connection not from this pool"
+      end
+      if not (pg.reserved) then
+        return nil, "connection not reserved"
+      end
+      if pg.transaction_status == "T" or pg.transaction_status == "E" then
+        local ok, err = pcall(function()
+          return pg:query("ROLLBACK")
+        end)
+        if not (ok and pg.transaction_status == "I") then
+          pg:disconnect()
+          for i, conn in ipairs(self.pool) do
+            if conn == pg then
+              table.remove(self.pool, i)
+              break
+            end
+          end
+          return nil, "rollback failed, connection removed from pool"
+        end
+      end
+      pg.reserved = false
+      return true
+    end,
     pool_size = function(self)
       return #self.pool
     end,

pgmoon/pool.moon

@@ -14,9 +14,9 @@ class PostgresPool
     -- Must call connect() first
     return nil, "not connected" if #@pool == 0
 
-    -- Find first non-busy instance
+    -- Find first non-busy and non-reserved instance
     for pg in *@pool
-      unless pg.busy
+      unless pg.busy or pg.reserved
         return pg
 
     -- All busy, check max_pool_size
@@ -108,6 +108,48 @@ class PostgresPool
   decode_bytea: Postgres.decode_bytea
   setup_hstore: Postgres.setup_hstore
 
+  -- Reserve/release for transaction support
+  reserve: =>
+    return nil, "not connected" if #@pool == 0
+
+    -- Find first available (non-busy and non-reserved) instance
+    for pg in *@pool
+      unless pg.busy or pg.reserved
+        pg.reserved = true
+        return pg
+
+    -- All busy/reserved, check max_pool_size
+    if @config.max_pool_size and #@pool >= @config.max_pool_size
+      return nil, "pool exhausted, max_pool_size reached"
+
+    -- Create and connect new instance
+    pg = @_create_instance!
+    ok, err = pg\connect!
+    return nil, err unless ok
+
+    table.insert @pool, pg
+    pg.reserved = true
+    pg
+
+  release: (pg) =>
+    return nil, "connection not from this pool" unless pg.parent_pool == @
+    return nil, "connection not reserved" unless pg.reserved
+
+    -- Ensure connection is idle before returning to pool
+    if pg.transaction_status == "T" or pg.transaction_status == "E"
+      ok, err = pcall -> pg\query "ROLLBACK"
+      -- If rollback failed, connection is likely broken - remove from pool
+      unless ok and pg.transaction_status == "I"
+        pg\disconnect!
+        for i, conn in ipairs @pool
+          if conn == pg
+            table.remove @pool, i
+            break
+        return nil, "rollback failed, connection removed from pool"
+
+    pg.reserved = false
+    true
+
   -- Pool info helpers
   pool_size: => #@pool
   active_connections: =>

spec/pgmoon_pool_spec.moon

@@ -194,3 +194,106 @@ describe "pgmoon.pool", ->
     it "has NULL constant", ->
       pool = PostgresPool {}
       assert.truthy pool.NULL
+
+  describe "reserve and release", ->
+    before_each ->
+      pool = PostgresPool { database: DB, user: USER, password: PASSWORD, host: HOST, port: PORT }
+      assert pool\connect!
+
+    after_each ->
+      pool\disconnect!
+
+    it "reserves and releases a connection for transaction", ->
+      pg = assert pool\reserve!
+      assert.truthy pg.reserved
+      assert.same 1, pool\pool_size!
+
+      assert pg\query "BEGIN"
+      assert.same "T", pg.transaction_status
+
+      assert pg\query "SELECT 1"
+      assert pg\query "COMMIT"
+      assert.same "I", pg.transaction_status
+
+      assert pool\release pg
+      assert.falsy pg.reserved
+
+    it "reserved connection is not reused by pool:query", ->
+      pg = assert pool\reserve!
+
+      -- pool:query should create a new connection since first is reserved
+      assert pool\query "SELECT 1"
+      assert.same 2, pool\pool_size!
+
+      assert pool\release pg
+
+    it "auto-rollback on release when in transaction (T state)", ->
+      pg = assert pool\reserve!
+
+      assert pg\query "BEGIN"
+      assert.same "T", pg.transaction_status
+
+      -- Release without committing
+      assert pool\release pg
+      assert.falsy pg.reserved
+      assert.same "I", pg.transaction_status
+
+    it "auto-rollback on release when in error state (E state)", ->
+      pg = assert pool\reserve!
+
+      assert pg\query "BEGIN"
+      -- Cause an error
+      res, err = pg\query "SELECT * FROM nonexistent_table_xyz"
+      assert.is_nil res
+      assert.same "E", pg.transaction_status
+
+      -- Release should auto-rollback and recover
+      assert pool\release pg
+      assert.falsy pg.reserved
+      assert.same "I", pg.transaction_status
+
+    it "respects max_pool_size when reserving", ->
+      pool\disconnect!
+      pool = PostgresPool {
+        database: DB, user: USER, password: PASSWORD, host: HOST, port: PORT
+        max_pool_size: 2
+      }
+      assert pool\connect!
+
+      pg1 = assert pool\reserve!
+      pg2 = assert pool\reserve!
+      assert.same 2, pool\pool_size!
+
+      res, err = pool\reserve!
+      assert.is_nil res
+      assert.same "pool exhausted, max_pool_size reached", err
+
+      pool\release pg1
+      pool\release pg2
+
+    it "errors when releasing connection not from this pool", ->
+      other_pool = PostgresPool { database: DB, user: USER, password: PASSWORD, host: HOST, port: PORT }
+      assert other_pool\connect!
+
+      pg = assert other_pool\reserve!
+
+      res, err = pool\release pg
+      assert.is_nil res
+      assert.same "connection not from this pool", err
+
+      other_pool\release pg
+      other_pool\disconnect!
+
+    it "errors when releasing connection that is not reserved", ->
+      pg = pool.pool[1]
+      assert.falsy pg.reserved
+
+      res, err = pool\release pg
+      assert.is_nil res
+      assert.same "connection not reserved", err
+
+    it "reserve returns error when not connected", ->
+      pool\disconnect!
+      res, err = pool\reserve!
+      assert.is_nil res
+      assert.same "not connected", err