From 1e1ffabd72d627bab1c5b1e35733ee4cb198c8fc Mon Sep 17 00:00:00 2001
From: XANTRONIX Development <dev@xantronix.com>
Date: Tue, 26 Nov 2024 16:24:32 -0500
Subject: [PATCH] Use callbacks for database connections

---
 lib/nntp/tiny/server.py  | 15 +++++++++------
 lib/nntp/tiny/session.py |  2 +-
 2 files changed, 10 insertions(+), 7 deletions(-)

diff --git a/lib/nntp/tiny/server.py b/lib/nntp/tiny/server.py
index f96f5d8..0b6a525 100644
--- a/lib/nntp/tiny/server.py
+++ b/lib/nntp/tiny/server.py
@@ -1,6 +1,7 @@
 import enum
 
-from nntp.tiny.db        import Database
+from typing import Callable
+
 from nntp.tiny.newsgroup import Newsgroup
 
 class ServerCapability(enum.Flag):
@@ -9,13 +10,15 @@ class ServerCapability(enum.Flag):
     POST = enum.auto()
 
 class Server():
-    def __init__(self, db: Database):
-        self.db           = db
-        self.capabilities = ServerCapability.NONE
-        self.newsgroups   = dict()
+    def __init__(self, connect_to_db: Callable):
+        self.connect_to_db = connect_to_db
+        self.capabilities  = ServerCapability.NONE
+        self.newsgroups    = dict()
 
         self._init_newsgroups()
 
     def _init_newsgroups(self):
-        for newsgroup in self.db.query(Newsgroup).each():
+        db = self.connect_to_db()
+
+        for newsgroup in db.query(Newsgroup).each():
             self.newsgroups[newsgroup.name.casefold()] = newsgroup
diff --git a/lib/nntp/tiny/session.py b/lib/nntp/tiny/session.py
index bbf3180..08bff91 100644
--- a/lib/nntp/tiny/session.py
+++ b/lib/nntp/tiny/session.py
@@ -106,7 +106,7 @@ class Session():
 
     def __init__(self, server: Server, sock: socket.socket):
         self.server: Server      = server
-        self.db: Database        = server.db
+        self.db: Database        = server.connect_to_db()
         self.sock: socket.socket = sock
         self.buf: LineBuffer     = LineBuffer()
         self.state: SessionState = SessionState.ACTIVE