database: Add a custom SELECT method

This is necessary because SQLModel only spits out the first column of
any custom queries.

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

diff --git a/src/dnsbl/database.py b/src/dnsbl/database.py
index 0eaff3befc3ef3873054db57af68493aee747d4d..2b165d5a2d10f8f8c0ba2c7d463bf5360672cee2 100644 (file)
--- a/src/dnsbl/database.py
+++ b/src/dnsbl/database.py
@@ -18,6 +18,7 @@
 #                                                                             #
 ###############################################################################
 
+import collections
 import functools
 import logging
 import sqlalchemy.orm
@@ -120,6 +121,22 @@ class Database(object):
                # Return the object
                return object
 
+       def select(self, stmt, batch_size=128):
+               """
+                       Selects custom queries
+               """
+               result = self.execute(stmt)
+
+               # Process the result in batches
+               if batch_size:
+                       result = result.yield_per(batch_size)
+
+               # Return all objects
+               for object in result.mappings():
+                       row = collections.namedtuple("Row", object.keys())
+
+                       yield row(**object)
+
        def fetch(self, stmt, batch_size=128):
                """
                        Fetches objects of the given type