REFACTOR: Use connection pooling

Signed-off-by: baalajimaestro <me@baalajimaestro.me>

src/nim_censor_bot.nim

@@ -12,14 +12,13 @@ import std/[asyncdispatch,
             options,
             strutils,
             random,
-            with,
             os,
             tables,
             times,
             sequtils,
             json]
 
-import norm/[model, sqlite]
+import norm/[model, sqlite, pool]
 
 # Logging Level
 var L = newConsoleLogger(levelThreshold=lvlError, fmtStr="$levelname, [$time] ")
@@ -49,6 +48,7 @@ var GroupMedia = initTable[int, string]()
 let AdminID = getEnv("ADMIN_ID").split(",")
 
 let dbConn* = getDb()
+var connPool = newPool[DbConn](15)
 
 # Functions to assist adding/deleting entries from tables
 func NewCensoredData*(ftype = ""; fhash = ""; fileid = ""; time = 0.0; caption = ""):
@@ -58,8 +58,9 @@ func NewBannedUsers*(userid = int64(0), bantime = 0.0, bantype = ""):
   BannedUsers = BannedUsers(userid: userid, bantime: bantime, bantype: bantype)
 
 # Create tables if they dont exist
-dbConn.createTables(NewCensoredData())
-dbConn.createTables(NewBannedUsers())
+withDb(connPool):
+  db.createTables(NewCensoredData())
+  db.createTables(NewBannedUsers())
 
 # Ratelimits a user if there is more than 30 timestamps within 60 secs. 
 # Resets their ratelimit counter if 60s has passed since first timestamp
@@ -70,24 +71,26 @@ proc ManageRateLimit(): void=
     elif len(RateLimiter[i]) >= 30:
       var BannedUser = NewBannedUsers(i, epochTime(), "auto")
       RateLimiter.del(i)
-      with dbConn:
-        insert BannedUser
-  if dbConn.exists(BannedUsers, "bantype = ? and ? - bantime >= 1800", "auto", epochTime()):
-    var TempData = @[NewBannedUsers()]
-    dbConn.select(TempData,  "bantype = ? and ? - bantime >= 1800", "auto", epochTime())
-    for i in TempData:
-      var e = i
-      dbConn.delete(e)
+      withDb(connPool):
+        db.insert(BannedUser)
+  withDb(connPool):
+    if db.exists(BannedUsers, "bantype = ? and ? - bantime >= 1800", "auto", epochTime()):
+      var TempData = @[NewBannedUsers()]
+      db.select(TempData,  "bantype = ? and ? - bantime >= 1800", "auto", epochTime())
+      for i in TempData:
+        var e = i
+        db.delete(e)
 
 # Cleanup old data to save space
 # Deletes old data after 6 months of entry
 proc OldDataCleanup(): void=
-  if dbConn.exists(CensoredData, "? - time >= 15780000", epochTime()):
-    var TempData = @[NewCensoredData()]
-    dbConn.select(TempData,  "? - time >= 15780000", epochTime())
-    for i in TempData:
-      var e = i
-      dbConn.delete(e)
+  withDb(connPool):
+    if db.exists(CensoredData, "? - time >= 15780000", epochTime()):
+      var TempData = @[NewCensoredData()]
+      db.select(TempData,  "? - time >= 15780000", epochTime())
+      for i in TempData:
+        var e = i
+        db.delete(e)
 
 # Generates a 6-character magic that will be used to identify the file
 proc generate_hash(): string=
@@ -101,86 +104,89 @@ proc generate_hash(): string=
 proc startHandler(b: Telebot, c: Command): Future[bool] {.gcsafe, async.} =
   let param = c.params
   ManageRateLimit()
-  if not dbConn.exists(BannedUsers, "userid = ?", c.message.chat.id):
-    # Update rate-limit counter
-    if not AdminID.contains($c.message.chat.id):
-      if RateLimiter.contains(c.message.chat.id):
-        RateLimiter[c.message.chat.id].insert(epochTime())
-      else:
-        RateLimiter[c.message.chat.id] = @[epochTime()]
-    if param == "":
-      discard await b.sendMessage(c.message.chat.id,
-              "Hey, To create a censored post, you can share any album, video, photo, gif, sticker, etc. " &
-              "The response of the bot can be used to share the media with other chats.")
-    else:
-      if not dbConn.exists(CensoredData, "fhash = ?", param):
-        discard await b.sendMessage(c.message.chat.id,
-                "Media does not exist on database, ask the sender to censor this again!")
-      else:
-        var TempData = @[NewCensoredData()]
-        dbConn.select(TempData, "fhash = ?", param)
-        if len(TempData) > 1:
-          var inputmedia = newSeq[InputMediaPhoto]()
-          for i in TempData:
-            if i.caption != "":
-              inputmedia.insert(InputMediaPhoto(kind: i.ftype, media: i.fileid, caption: some(i.caption)))
-            else:
-              inputmedia.insert(InputMediaPhoto(kind: i.ftype, media: i.fileid))
-          discard await b.sendMediaGroup($c.message.chat.id, media=inputmedia)
+  withDb(connPool):
+    if not db.exists(BannedUsers, "userid = ?", c.message.chat.id):
+      # Update rate-limit counter
+      if not AdminID.contains($c.message.chat.id):
+        if RateLimiter.contains(c.message.chat.id):
+          RateLimiter[c.message.chat.id].insert(epochTime())
         else:
-          if TempData[0].ftype == "photo":
-            discard await b.sendPhoto(chatId=c.message.chat.id, photo=TempData[0].fileid, caption=TempData[0].caption)
-          elif TempData[0].ftype == "document":
-            discard await b.sendDocument(c.message.chat.id, TempData[0].fileid, caption=TempData[0].caption)
-          elif TempData[0].ftype == "video":
-            discard await b.sendVideo(c.message.chat.id, TempData[0].fileid, caption=TempData[0].caption)
-          elif TempData[0].ftype == "videoNote":
-            discard await b.sendVideoNote(c.message.chat.id, TempData[0].fileid)
-          elif TempData[0].ftype == "animation":
-            discard await b.sendAnimation(c.message.chat.id, TempData[0].fileid, caption=TempData[0].caption)
-          elif TempData[0].ftype == "sticker":
-            discard await b.sendSticker(c.message.chat.id, TempData[0].fileid)
+          RateLimiter[c.message.chat.id] = @[epochTime()]
+      if param == "":
+        discard await b.sendMessage(c.message.chat.id,
+                "Hey, To create a censored post, you can share any album, video, photo, gif, sticker, etc. " &
+                "The response of the bot can be used to share the media with other chats.")
+      else:
+        if not db.exists(CensoredData, "fhash = ?", param):
+          discard await b.sendMessage(c.message.chat.id,
+                  "Media does not exist on database, ask the sender to censor this again!")
+        else:
+          var TempData = @[NewCensoredData()]
+          db.select(TempData, "fhash = ?", param)
+          if len(TempData) > 1:
+            var inputmedia = newSeq[InputMediaPhoto]()
+            for i in TempData:
+              if i.caption != "":
+                inputmedia.insert(InputMediaPhoto(kind: i.ftype, media: i.fileid, caption: some(i.caption)))
+              else:
+                inputmedia.insert(InputMediaPhoto(kind: i.ftype, media: i.fileid))
+            discard await b.sendMediaGroup($c.message.chat.id, media=inputmedia)
+          else:
+            if TempData[0].ftype == "photo":
+              discard await b.sendPhoto(chatId=c.message.chat.id, photo=TempData[0].fileid, caption=TempData[0].caption)
+            elif TempData[0].ftype == "document":
+              discard await b.sendDocument(c.message.chat.id, TempData[0].fileid, caption=TempData[0].caption)
+            elif TempData[0].ftype == "video":
+              discard await b.sendVideo(c.message.chat.id, TempData[0].fileid, caption=TempData[0].caption)
+            elif TempData[0].ftype == "videoNote":
+              discard await b.sendVideoNote(c.message.chat.id, TempData[0].fileid)
+            elif TempData[0].ftype == "animation":
+              discard await b.sendAnimation(c.message.chat.id, TempData[0].fileid, caption=TempData[0].caption)
+            elif TempData[0].ftype == "sticker":
+              discard await b.sendSticker(c.message.chat.id, TempData[0].fileid)
 
 # Give them source url
 proc sourceHandler(b: Telebot, c: Command): Future[bool] {.gcsafe, async.} =
   ManageRateLimit()
-  if not dbConn.exists(BannedUsers, "userid = ?", c.message.chat.id):
-    # Update rate-limit counter
-    if not AdminID.contains($c.message.chat.id):
-      if RateLimiter.contains(c.message.chat.id):
-        RateLimiter[c.message.chat.id].insert(epochTime())
-      else:
-        RateLimiter[c.message.chat.id] = @[epochTime()]
-    discard await b.sendMessage(c.message.chat.id,
-            "Hey, this bot is open-source and licensed under AGPL-v3! " &
-            "\n\nYou are welcome to selfhost your own instance of this bot" &
-            "\n\nThe source code is [here](https://git.baalajimaestro.me/baalajimaestro/nim-censor-bot)",
-            parseMode="Markdown",
-            disableWebPagePreview = true)
+  withDb(connPool):
+    if not db.exists(BannedUsers, "userid = ?", c.message.chat.id):
+      # Update rate-limit counter
+      if not AdminID.contains($c.message.chat.id):
+        if RateLimiter.contains(c.message.chat.id):
+          RateLimiter[c.message.chat.id].insert(epochTime())
+        else:
+          RateLimiter[c.message.chat.id] = @[epochTime()]
+      discard await b.sendMessage(c.message.chat.id,
+              "Hey, this bot is open-source and licensed under AGPL-v3! " &
+              "\n\nYou are welcome to selfhost your own instance of this bot" &
+              "\n\nThe source code is [here](https://git.baalajimaestro.me/baalajimaestro/nim-censor-bot)",
+              parseMode="Markdown",
+              disableWebPagePreview = true)
 
 # UnBan Handler
 proc unbanHandler(b: Telebot, c: Command): Future[bool] {.gcsafe, async.} =
   if $c.message.chat.id in AdminID:
     let user = c.params
-    if dbConn.exists(BannedUsers, "userid = ?", int64(parseInt(user))):
-      var TempData = @[NewBannedUsers()]
-      dbConn.select(TempData,  "userid = ?", int64(parseInt(user)))
-      for i in TempData:
-        var e = i
-        dbConn.delete(e)
-      discard await b.sendMessage(c.message.chat.id, "Unbanned!")
+    withDb(connPool):
+      if db.exists(BannedUsers, "userid = ?", int64(parseInt(user))):
+        var TempData = @[NewBannedUsers()]
+        db.select(TempData,  "userid = ?", int64(parseInt(user)))
+        for i in TempData:
+          var e = i
+          db.delete(e)
+        discard await b.sendMessage(c.message.chat.id, "Unbanned!")
 
 # ban Handler
 proc banHandler(b: Telebot, c: Command): Future[bool] {.gcsafe, async.} =
   if $c.message.chat.id in AdminID:
     let user = c.params
     var BannedUser = NewBannedUsers(int64(parseInt(user)), epochTime(), "permanent")
-    with dbConn:
-      insert BannedUser
+    withDb(connPool):
+        db.insert(BannedUser)
     discard await b.sendMessage(c.message.chat.id, "Banned!")
 
 # Inline share handler
-proc inlineHandler(b: Telebot, u: InlineQuery): Future[bool]{.async.} =
+proc inlineHandler(b: Telebot, u: InlineQuery): Future[bool]{.gcsafe, async.} =
     var TempData = @[NewCensoredData()]
     var ftype = ""
     var res: InlineQueryResultArticle
@@ -188,25 +194,26 @@ proc inlineHandler(b: Telebot, u: InlineQuery): Future[bool]{.async.} =
     res.kind = "article"
     res.id = "1"
     if u.query != "":
-        if not dbConn.exists(CensoredData, "fhash = ?",u.query):
-          res.title = "Media Not Found"
-          res.inputMessageContent = InputTextMessageContent(
-                "Media does not exist on database, ask the sender to censor this again!").some
-        else:
-          dbConn.select(TempData, "fhash = ?", u.query)
-          if len(TempData) > 1:
-            ftype = "Album"
-          elif len(TempData) == 1:
-            ftype = TempData[0].ftype
-          res.title = "NSFW " & capitalizeAscii(ftype)
-          res.inputMessageContent = InputMessageContent(kind: TextMessage,
-                messageText:"*NSFW " &
-                capitalizeAscii(ftype) &
-                "*\n\n[Tap to View](tg://resolve?domain=" &
-                b.username & "&start=" &
-                u.query &
-                ")",
-                parseMode: some("Markdown")).some
+        withDb(connPool):
+          if not db.exists(CensoredData, "fhash = ?",u.query):
+            res.title = "Media Not Found"
+            res.inputMessageContent = InputTextMessageContent(
+                  "Media does not exist on database, ask the sender to censor this again!").some
+          else:
+            db.select(TempData, "fhash = ?", u.query)
+            if len(TempData) > 1:
+              ftype = "Album"
+            elif len(TempData) == 1:
+              ftype = TempData[0].ftype
+            res.title = "NSFW " & capitalizeAscii(ftype)
+            res.inputMessageContent = InputMessageContent(kind: TextMessage,
+                  messageText:"*NSFW " &
+                  capitalizeAscii(ftype) &
+                  "*\n\n[Tap to View](tg://resolve?domain=" &
+                  b.username & "&start=" &
+                  u.query &
+                  ")",
+                  parseMode: some("Markdown")).some
     else:
       res.title = "Waiting for File Hash"
       res.inputMessageContent = InputTextMessageContent(
@@ -220,82 +227,83 @@ proc updateHandler(b: Telebot, u: Update): Future[bool] {.async, gcsafe.} =
     let response = u.message.get
     # Refresh rate-limits
     ManageRateLimit()
-    # Dont bother about rate-limited/banned users
-    if not dbConn.exists(BannedUsers, "userid = ?", response.chat.id):
-      # Update rate-limit counter
-      if not AdminID.contains($response.chat.id):
-        if RateLimiter.contains(response.chat.id):
-          RateLimiter[response.chat.id].insert(epochTime())
-        else:
-          RateLimiter[response.chat.id] = @[epochTime()]
-      if not response.text.isSome:
-        var fileid = ""
-        var ftype = ""
-        var fcaption = ""
-        if response.caption.isSome:
-          fcaption = response.caption.get
-        if response.document.isSome:
-          fileid = response.document.get.fileId
-          ftype = "document"
-        elif response.video.isSome:
-          fileid = response.video.get.fileId
-          ftype = "video"
-        elif response.videoNote.isSome:
-          fileid = response.videoNote.get.fileId
-          ftype = "videoNote"
-        elif response.animation.isSome:
-          fileid = response.animation.get.fileId
-          ftype = "animation"
-        elif response.photo.isSome:
-          fileid = response.photo.get[0].fileId
-          ftype = "photo"
-        elif response.sticker.isSome:
-          fileid = response.sticker.get.fileId
-          ftype = "sticker"
-        # Workaround for groupmedia using hashtables
-        # Telegram is sending multiple updates, instead of 1, so just workaround it
-        if response.mediaGroupId.isSome:
-          if parseInt(response.mediaGroupId.get) notin GroupMedia.keys().toSeq():
+    withDb(connPool):
+      # Dont bother about rate-limited/banned users
+      if not db.exists(BannedUsers, "userid = ?", response.chat.id):
+        # Update rate-limit counter
+        if not AdminID.contains($response.chat.id):
+          if RateLimiter.contains(response.chat.id):
+            RateLimiter[response.chat.id].insert(epochTime())
+          else:
+            RateLimiter[response.chat.id] = @[epochTime()]
+        if not response.text.isSome:
+          var fileid = ""
+          var ftype = ""
+          var fcaption = ""
+          if response.caption.isSome:
+            fcaption = response.caption.get
+          if response.document.isSome:
+            fileid = response.document.get.fileId
+            ftype = "document"
+          elif response.video.isSome:
+            fileid = response.video.get.fileId
+            ftype = "video"
+          elif response.videoNote.isSome:
+            fileid = response.videoNote.get.fileId
+            ftype = "videoNote"
+          elif response.animation.isSome:
+            fileid = response.animation.get.fileId
+            ftype = "animation"
+          elif response.photo.isSome:
+            fileid = response.photo.get[0].fileId
+            ftype = "photo"
+          elif response.sticker.isSome:
+            fileid = response.sticker.get.fileId
+            ftype = "sticker"
+          # Workaround for groupmedia using hashtables
+          # Telegram is sending multiple updates, instead of 1, so just workaround it
+          if response.mediaGroupId.isSome:
+            if parseInt(response.mediaGroupId.get) notin GroupMedia.keys().toSeq():
+              let filehash = generate_hash()
+              GroupMedia[parseInt(response.mediaGroupId.get)] = filehash
+              var CensoredRow = NewCensoredData(ftype, filehash, fileid, epochTime(), fcaption)
+              withDb(connPool):
+                db.insert(CensoredRow)
+              var replybutton = InlineKeyboardButton(text: "Share", switchInlineQuery: some(filehash))
+              let replymark = newInlineKeyboardMarkup(@[replybutton])
+              discard await b.sendMessage(response.chat.id,
+                  "*NSFW " & 
+                  "Album" &
+                  "*\n\n[Tap to View](tg://resolve?domain=" &
+                  b.username &
+                  "&start=" &
+                  filehash &
+                  ")",
+                  replyMarkup = replymark,
+                  parseMode="Markdown")
+            else:
+              let filehash = GroupMedia[parseInt(response.mediaGroupId.get)]
+              var CensoredRow = NewCensoredData(ftype, filehash, fileid, epochTime(), fcaption)
+              withDb(connPool):
+                db.insert(CensoredRow)
+
+          else:
             let filehash = generate_hash()
-            GroupMedia[parseInt(response.mediaGroupId.get)] = filehash
             var CensoredRow = NewCensoredData(ftype, filehash, fileid, epochTime(), fcaption)
-            with dbConn:
-                  insert CensoredRow
+            withDb(connPool):
+                db.insert(CensoredRow)
             var replybutton = InlineKeyboardButton(text: "Share", switchInlineQuery: some(filehash))
             let replymark = newInlineKeyboardMarkup(@[replybutton])
             discard await b.sendMessage(response.chat.id,
                 "*NSFW " & 
-                "Album" &
+                capitalizeAscii(ftype) &
                 "*\n\n[Tap to View](tg://resolve?domain=" &
                 b.username &
                 "&start=" &
                 filehash &
                 ")",
                 replyMarkup = replymark,
-                 parseMode="Markdown")
-          else:
-            let filehash = GroupMedia[parseInt(response.mediaGroupId.get)]
-            var CensoredRow = NewCensoredData(ftype, filehash, fileid, epochTime(), fcaption)
-            with dbConn:
-                  insert CensoredRow
-
-        else:
-          let filehash = generate_hash()
-          var CensoredRow = NewCensoredData(ftype, filehash, fileid, epochTime(), fcaption)
-          with dbConn:
-                insert CensoredRow
-          var replybutton = InlineKeyboardButton(text: "Share", switchInlineQuery: some(filehash))
-          let replymark = newInlineKeyboardMarkup(@[replybutton])
-          discard await b.sendMessage(response.chat.id,
-              "*NSFW " & 
-              capitalizeAscii(ftype) &
-              "*\n\n[Tap to View](tg://resolve?domain=" &
-              b.username &
-              "&start=" &
-              filehash &
-              ")",
-              replyMarkup = replymark,
-              parseMode="Markdown")
+                parseMode="Markdown")
       OldDataCleanup()
 
 when isMainModule: