- # No session specified as last argument or in kwargs, create one.
- if session is None or len(args) == len(getargspec(fn)[0]) - 1:
- private_transaction = True
- kwargs['session'] = DBConn().session()
+ # Find the session object
+ try:
+ session = kwargs['session']
+ except KeyError:
+ if len(args) <= len(getargspec(fn)[0]) - 1:
+ # No session specified as last argument or in kwargs
+ private_transaction = True
+ session = kwargs['session'] = DBConn().session()
+ else:
+ # Session is last argument in args
+ session = args[-1]
+
+ if private_transaction:
+ session.commit_or_flush = session.commit
+ else:
+ session.commit_or_flush = session.flush