implement database connection pools

Signed-off-by: George Lemon <georgelemon@protonmail.com>

Author: georgelemon

src/ozark/database.nim

@@ -4,7 +4,7 @@
 #          Made by Humans from OpenPeeps
 #          https://github.com/openpeeps/ozark
 
-import std/[macros, net, strutils, tables]
+import std/[macros, net, strutils, tables, locks, os]
 
 import pkg/threading/once
 import pkg/db_connector/db_postgres
@@ -14,28 +14,27 @@ export Port, strVal, `%`
 
 type
   DBConnectionPool* = ref object
-    connections*: seq[DBConn]
-    busyConnections*: seq[DBConn]
+    connections*: seq[DBConn]      # available
+    busyConnections*: seq[DBConn]  # checked-out
+    maxSize*: int
+    lock: Lock
 
   DBDriver* = enum
     PostgreSQLDriver
     MYSQLDriver
     SQLiteDriver
 
-  DBConnection = ref object
-    driver: DBDriver
+  DBConnection* = ref object
+    driver*: DBDriver
     address*, name*, user*, password*: string
-    # dbConn: DBConn
-    # pool: DBConnectionPool
-    port: Port
+    port*: Port
 
-  DBConnections = OrderedTableRef[string, DBConnection]
+  DBConnections* = OrderedTableRef[string, DBConnection]
 
   Ozark = object
     dbs: DBConnections
-      # holds credentials for multiple Database Connections
     maindb: DBConnection
-      # credentials for the main database connection
+    mainPool: DBConnectionPool
 
 var
   DB: ptr Ozark
@@ -45,13 +44,12 @@ proc getInstance*(): ptr Ozark =
   ## Get the singleton instance of the database manager
   once(o):
     DB = createShared(Ozark)
+    DB[].dbs = newOrderedTable[string, DBConnection]() # init map
   result = DB
 
 proc initOzarkDatabase*(address, name, user, password: string,
                 port: Port = Port(5432),
                 driver: DBDriver = DBDriver.PostgreSQLDriver) =
-  ## Initializes the singleton instance of the database manager
-  ## using provided credentials as main database
   let db = getInstance()
   db[].maindb = DBConnection(
     address: address,
@@ -102,5 +100,93 @@ macro withDatabase*(id: static string, body: untyped) =
                 db[id].password, db[$id].name)
     defer:
       dbcon.close()
+    block:
+      `body`
+
+
+proc openConn(cfg: DBConnection): DBConn =
+  case cfg.driver
+  of PostgreSQLDriver:
+    open(cfg.address, cfg.user, cfg.password, cfg.name)
+  else:
+    raise newException(ValueError, "Only PostgreSQL driver pool is currently implemented.")
+
+proc initOzarkPool*(size: Positive = 10) =
+  ## Initialize main DB connection pool.
+  let db = getInstance()
+  assert db[].maindb != nil, "Main DB credentials not initialized. Call initOzarkDatabase first."
+
+  var pool = DBConnectionPool(
+    maxSize: size.int,
+    connections: @[],
+    busyConnections: @[]
+  )
+  initLock(pool.lock)
+
+  for _ in 0..<size.int:
+    pool.connections.add(openConn(db[].maindb))
+
+  db[].mainPool = pool
+
+proc closeOzarkPool*() =
+  ## Close all pooled connections.
+  let db = getInstance()
+  if db[].mainPool.isNil: return
+
+  acquire(db[].mainPool.lock)
+  defer: release(db[].mainPool.lock)
+
+  for c in db[].mainPool.connections:
+    c.close()
+  for c in db[].mainPool.busyConnections:
+    c.close()
+
+  db[].mainPool.connections.setLen(0)
+  db[].mainPool.busyConnections.setLen(0)
+
+proc acquireConn*(pool: DBConnectionPool, timeoutMs: int = 5000): DBConn =
+  ## Borrow one connection from pool, waiting up to timeoutMs.
+  let stepMs = 10
+  var waited = 0
+  while waited <= timeoutMs:
+    acquire(pool.lock)
+    if pool.connections.len > 0:
+      result = pool.connections.pop()
+      pool.busyConnections.add(result)
+      release(pool.lock)
+      return
+    release(pool.lock)
+    sleep(stepMs)
+    inc(waited, stepMs)
+
+  raise newException(ValueError, "Timed out waiting for a DB connection from pool.")
+
+proc releaseConn*(pool: DBConnectionPool, conn: DBConn) =
+  ## Return a connection to pool.
+  acquire(pool.lock)
+  defer: release(pool.lock)
+
+  var idx = -1
+  for i, c in pool.busyConnections:
+    if c == conn:
+      idx = i
+      break
+
+  if idx >= 0:
+    pool.busyConnections.del(idx)
+    pool.connections.add(conn)
+
+macro withDBPool*(body: untyped) =
+  ## Run queries using a pooled connection.
+  result = newStmtList()
+  add result, quote do:
+    let db = getInstance()
+    assert db != nil, "Database manager not initialized. Call initOzarkDatabase first."
+    assert db[].mainPool != nil, "DB pool not initialized. Call initOzarkPool first."
+
+    let dbcon {.inject.} = acquireConn(db[].mainPool)
+    defer:
+      releaseConn(db[].mainPool, dbcon)
+
     block:
       `body`

tests/test2.nim

@@ -29,17 +29,15 @@ test "init embedded postgres and create tables":
 
   initOzarkDatabase("localhost", "postgres", "postgres", "postgres", Port(5432))
   withDB do:
-    Models.rawSQL("""
-CREATE TABLE IF NOT EXISTS users (
-  id SERIAL PRIMARY KEY,
-  username VARCHAR(50),
-  name VARCHAR(100),
-  email VARCHAR(100)
-)""").exec()
+    Models.table(Users).prepareTable().exec()
+    Models.table(Users).dropTable(cascade = true).exec()
+    Models.table(Users).prepareTable().exec()
+
+  initOzarkPool(15)
 
 suite "INSERT and SELECT queries":
   test "insert and select data":
-    withDB do:
+    withDBPool do:
       var name = "John"
       let id = Models.table(Users).insert({
         name: name,
@@ -58,7 +56,7 @@ suite "INSERT and SELECT queries":
       check res.entries[0].email == "test@example.com"
 
   test "select specific columns":
-    withDB do:
+    withDBPool do:
       let res = Models.table(Users)
                       .select(["name", "email"])
                       .where("name", "John").get()
@@ -68,22 +66,22 @@ suite "INSERT and SELECT queries":
 
 suite "WHERE queries":
   test "where query":
-    withDB do:
+    withDBPool do:
       let res = Models.table(Users)
                       .select("name")
                       .where("name", "John").get()
       check res.isEmpty == false
       # check res.get(0).name == "John"
 
   test "whereNot query":
-    withDB do:
+    withDBPool do:
       let res = Models.table(Users)
                       .select("name")
                       .whereNot("name", "John").get()
       check res.isEmpty
 
   test "orWhere query":
-    withDB do:
+    withDBPool do:
       let res = Models.table(Users)
                       .select("name")
                       .where("name", "Ghost")
@@ -92,7 +90,7 @@ suite "WHERE queries":
       check res.get(0).name == "John"
 
   test "orWhereNot query":
-    withDB do:
+    withDBPool do:
       let res = Models.table(Users)
                       .select("name")
                       .whereNot("name", "John")
@@ -103,51 +101,51 @@ suite "WHERE queries":
 
 suite "LIKE queries":
   test "like query":
-    withDB do:
+    withDBPool do:
       let res = Models.table(Users).select("name")
                       .whereLike("name", "Jo").get()
       check res.isEmpty == false
       check res.get(0).name == "John"
 
   test "whereStartsLike query":
-    withDB do:
+    withDBPool do:
       let res = Models.table(Users).select("name")
                       .whereStartsLike("name", "Jo").get()
       check res.isEmpty == false
       check res.get(0).name == "John"
 
 
   test "whereEndsLike query":
-    withDB do:
+    withDBPool do:
       let res = Models.table(Users).select("name")
                       .whereEndsLike("name", "hn").get()
       check res.isEmpty == false
       check res.get(0).name == "John"
 
   test "whereNotLike query":
-    withDB do:
+    withDBPool do:
       let res = Models.table(Users).select("name")
                       .whereNot("name", "Ghost").get()
       check res.isEmpty == false
       check res.get(0).name == "John"
 
   test "wereNotStartsLike query":
-    withDB do:
+    withDBPool do:
       let res = Models.table(Users).select("name")
                       .whereNotStartsLike("name", "Gh").get()
       check res.isEmpty == false
       check res.get(0).name == "John"
 
   test "whereNotEndsLike query":
-    withDB do:
+    withDBPool do:
       let res = Models.table(Users).select("name")
                       .whereNotEndsLike("name", "st").get()
       check res.isEmpty == false
       check res.get(0).name == "John"
 
 suite "IN queries":
   test "whereNotIn query":
-    withDB do:
+    withDBPool do:
       let res = Models.table(Users).select("name")
                       .whereNotIn("name", "John").get()
       check res.isEmpty == true