X-Git-Url: https://git.decadent.org.uk/gitweb/?a=blobdiff_plain;f=daklib%2Fdbconn.py;h=224935dd5883504a59c9d1b27f44175ce227e287;hb=3a55501ad9f5eaed322f5ea8365bb799e468b2e2;hp=2f5fef30f9866d2ee158b010fe89ed9cae5d9bb2;hpb=28f3e6f0659504214d482d682f4b2b4791872cc6;p=dak.git

diff --git a/daklib/dbconn.py b/daklib/dbconn.py
index 2f5fef30..224935dd 100755
--- a/daklib/dbconn.py
+++ b/daklib/dbconn.py
@@ -39,7 +39,7 @@ import traceback
 
 from inspect import getargspec
 
-from sqlalchemy import create_engine, Table, MetaData, select
+from sqlalchemy import create_engine, Table, MetaData
 from sqlalchemy.orm import sessionmaker, mapper, relation
 
 # Don't remove this, we re-export the exceptions to scripts which import us
@@ -73,9 +73,9 @@ def session_wrapper(fn):
         private_transaction = False
 
         # Find the session object
-        try:
-            session = kwargs['session']
-        except KeyError:
+        session = kwargs.get('session')
+
+        if session is None:
             if len(args) <= len(getargspec(fn)[0]) - 1:
                 # No session specified as last argument or in kwargs
                 private_transaction = True
@@ -83,6 +83,9 @@ def session_wrapper(fn):
             else:
                 # Session is last argument in args
                 session = args[-1]
+                if session is None:
+                    session = args[-1] = DBConn().session()
+                    private_transaction = True
 
         if private_transaction:
             session.commit_or_flush = session.commit