database: Make the backend available to all objects

Signed-off-by: Michael Tremer <michael.tremer@ipfire.org>

diff --git a/src/dnsbl/database.py b/src/dnsbl/database.py
index 6dc8b4468c4dcda8e1016f7cbaaff9c63b24e60d..b97eb89461f4f97d916d74dbe3355567587a9342 100644 (file)
--- a/src/dnsbl/database.py
+++ b/src/dnsbl/database.py
@@ -18,7 +18,9 @@
 #                                                                             #
 ###############################################################################
 
+import functools
 import logging
+import sqlalchemy.orm
 import sqlmodel
 
 # Setup logging
@@ -31,6 +33,15 @@ class Database(object):
                # Connect to the database
                self.engine = self.connect(uri)
 
+               # Create a session maker
+               self.sessionmaker = sqlalchemy.orm.sessionmaker(
+                       self.engine,
+                       expire_on_commit = False,
+                       info             = {
+                               "backend" : self.backend,
+                       },
+               )
+
                # Session
                self.__session = None
 
@@ -56,7 +67,7 @@ class Database(object):
                        Returns the current database session
                """
                if self.__session is None:
-                       self.__session = sqlmodel.Session(self.engine)
+                       self.__session = self.sessionmaker()
 
                return self.__session
 
@@ -138,3 +149,17 @@ class Database(object):
 
                # Return as set
                return set([o for o in objects])
+
+
+class BackendMixin:
+       @functools.cached_property
+       def backend(self):
+               # Fetch the session
+               session = sqlalchemy.orm.object_session(self)
+
+               # Return the backend
+               return session.info.get("backend")
+
+       @functools.cached_property
+       def db(self):
+               return self.backend.db

diff --git a/src/dnsbl/lists.py b/src/dnsbl/lists.py
index 13967858b49d5219d306e6bfcb2e3ca5c0dfa6c7..479ec87cb61217ede3be33371134e82ff761a193 100644 (file)
--- a/src/dnsbl/lists.py
+++ b/src/dnsbl/lists.py
@@ -22,6 +22,7 @@ import datetime
 import sqlmodel
 import typing
 
+from . import database
 from . import sources
 from . import util
 
@@ -92,7 +93,7 @@ class Lists(object):
                )
 
 
-class List(sqlmodel.SQLModel, table=True):
+class List(sqlmodel.SQLModel, database.BackendMixin, table=True):
        __tablename__ = "lists"
 
        # ID

diff --git a/src/dnsbl/sources.py b/src/dnsbl/sources.py
index ed30291cf9bdfa33b9bf0ab8dcfd58551742233b..8465b47d65e0b8910f4be49daa6ce40ab15a828d 100644 (file)
--- a/src/dnsbl/sources.py
+++ b/src/dnsbl/sources.py
@@ -21,6 +21,8 @@
 import datetime
 import sqlmodel
 
+from . import database
+
 class Sources(object):
        def __init__(self, backend):
                self.backend = backend
@@ -60,7 +62,7 @@ class Sources(object):
                return source
 
 
-class Source(sqlmodel.SQLModel, table=True):
+class Source(sqlmodel.SQLModel, database.BackendMixin, table=True):
        __tablename__ = "sources"
 
        # ID