From be5ad97f571b69c29e20b650eeeeb5c17c5724a7 Mon Sep 17 00:00:00 2001
From: Kamil Tomsik <info@tomsik.cz>
Date: Tue, 13 Aug 2024 23:46:08 +0200
Subject: [PATCH] pool: add basic test, fix restore session
 auto-closing/releasing, but only for owned/pooled connections

---
 src/pool.zig    | 29 +++++++++++++++++++++++++----
 src/session.zig | 17 +++++++++++++++++
 2 files changed, 42 insertions(+), 4 deletions(-)

diff --git a/src/pool.zig b/src/pool.zig
index a6f877d..383418d 100644
--- a/src/pool.zig
+++ b/src/pool.zig
@@ -35,7 +35,9 @@ pub const Pool = struct {
     }
 
     pub fn getSession(self: *Pool, allocator: std.mem.Allocator) Error!Session {
-        return Session.init(allocator, self.getConnection());
+        var sess = try Session.init(allocator, try self.getConnection());
+        sess.pool = self;
+        return sess;
     }
 
     /// Get a connection from the pool. If the pool is empty, this will block
@@ -94,15 +96,34 @@ pub const Pool = struct {
     }
 };
 
+const t = std.testing;
+
 test {
-    var pool = Pool.init(@import("sqlite.zig").SQLite3, std.testing.allocator, 3, &.{
+    var pool = Pool.init(@import("sqlite.zig").SQLite3, t.allocator, 3, &.{
         .filename = ":memory:",
     });
     defer pool.deinit();
 
     const c1 = try pool.getConnection();
-    defer pool.releaseConnection(c1);
+    try t.expectEqual(1, pool.count);
+    try t.expectEqual(0, pool.conns.items.len);
+
+    pool.releaseConnection(c1);
+    try t.expectEqual(1, pool.count);
+    try t.expectEqual(1, pool.conns.items.len);
 
     const c2 = try pool.getConnection();
-    defer pool.releaseConnection(c2);
+    try t.expectEqual(1, pool.count);
+    try t.expectEqual(0, pool.conns.items.len);
+    try t.expectEqual(c1.handle, c2.handle);
+
+    const c3 = try pool.getConnection();
+    try t.expectEqual(2, pool.count);
+    try t.expectEqual(0, pool.conns.items.len);
+    try t.expect(c1.handle != c3.handle);
+
+    pool.releaseConnection(c2);
+    pool.releaseConnection(c3);
+    try t.expectEqual(2, pool.count);
+    try t.expectEqual(2, pool.conns.items.len);
 }
diff --git a/src/session.zig b/src/session.zig
index a5bed5f..46adc48 100644
--- a/src/session.zig
+++ b/src/session.zig
@@ -2,12 +2,21 @@ const std = @import("std");
 const sqlite = @import("sqlite.zig");
 const util = @import("util.zig");
 const Connection = @import("connection.zig").Connection;
+const Pool = @import("pool.zig").Pool;
 const Statement = @import("statement.zig").Statement;
 const Query = @import("query.zig").Query;
 
 pub const Session = struct {
     arena: std.mem.Allocator,
     conn: Connection,
+    pool: ?*Pool = null,
+    close: bool = false,
+
+    pub fn open(comptime T: type, allocator: std.mem.Allocator, options: T.Options) !Session {
+        var sess = try Session.init(allocator, try Connection.open(T, options));
+        sess.close = true;
+        return sess;
+    }
 
     pub fn init(allocator: std.mem.Allocator, conn: Connection) !Session {
         const arena = try allocator.create(std.heap.ArenaAllocator);
@@ -23,6 +32,14 @@ pub const Session = struct {
         const arena: *std.heap.ArenaAllocator = @ptrCast(@alignCast(self.arena.ptr));
         arena.deinit();
         arena.child_allocator.destroy(arena);
+
+        if (self.pool) |pool| {
+            pool.releaseConnection(self.conn);
+        } else {
+            if (self.close) {
+                self.conn.close();
+            }
+        }
     }
 
     pub fn prepare(self: *Session, sql: []const u8, args: anytype) !Statement {