diff --git a/rhodecode/lib/celerylib/__init__.py b/rhodecode/lib/celerylib/__init__.py --- a/rhodecode/lib/celerylib/__init__.py +++ b/rhodecode/lib/celerylib/__init__.py @@ -7,7 +7,7 @@ :created_on: Nov 27, 2010 :author: marcink - :copyright: (C) 2009-2011 Marcin Kuzminski + :copyright: (C) 2010-2012 Marcin Kuzminski :license: GPLv3, see COPYING for more details. """ # This program is free software: you can redistribute it and/or modify @@ -29,19 +29,23 @@ import socket import traceback import logging from os.path import dirname as dn, join as jn +from pylons import config from hashlib import md5 from decorator import decorator -from pylons import config from vcs.utils.lazy import LazyProperty from rhodecode.lib import str2bool, safe_str from rhodecode.lib.pidlock import DaemonLock, LockHeld +from rhodecode.model import init_model +from rhodecode.model import meta +from rhodecode.model.db import Statistics, Repository, User + +from sqlalchemy import engine_from_config from celery.messaging import establish_connection - log = logging.getLogger(__name__) try: @@ -107,3 +111,23 @@ def locked_task(func): return 'Task with key %s already running' % lockkey return decorator(__wrapper, func) + + +def get_session(): + if CELERY_ON: + engine = engine_from_config(config, 'sqlalchemy.db1.') + init_model(engine) + sa = meta.Session + return sa + + +def dbsession(func): + def __wrapper(func, *fargs, **fkwargs): + try: + ret = func(*fargs, **fkwargs) + return ret + finally: + if CELERY_ON: + meta.Session.remove() + + return decorator(__wrapper, func)