celery: create custom celery task class which loads context/threadlocals

of the original request required for tasks like create_repo which needs
to generate a url after the repo is created. fixes  #4139

rhodecode/lib/celerylib/__init__.py

@@ -1,147 +1,228 b''
 # -*- coding: utf-8 -*-
 
 # Copyright (C) 2010-2016  RhodeCode GmbH
 #
 # This program is free software: you can redistribute it and/or modify
 # it under the terms of the GNU Affero General Public License, version 3
 # (only), as published by the Free Software Foundation.
 #
 # This program is distributed in the hope that it will be useful,
 # but WITHOUT ANY WARRANTY; without even the implied warranty of
 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 # GNU General Public License for more details.
 #
 # You should have received a copy of the GNU Affero General Public License
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
 #
 # This program is dual-licensed. If you wish to learn more about the
 # RhodeCode Enterprise Edition, including its added features, Support services,
 # and proprietary license terms, please see https://rhodecode.com/licenses/
 """
 celery libs for RhodeCode
 """
 
 
+import pylons
 import socket
 import logging
 
 import rhodecode
 
 from os.path import join as jn
 from pylons import config
+from celery.task import Task
+from pyramid.request import Request
+from pyramid.scripting import prepare
+from pyramid.threadlocal import get_current_request
 
 from decorator import decorator
 
 from zope.cachedescriptors.property import Lazy as LazyProperty
 
 from rhodecode.config import utils
-from rhodecode.lib.utils2 import safe_str, md5_safe, aslist
+from rhodecode.lib.utils2 import (
+    safe_str, md5_safe, aslist, get_routes_generator_for_server_url,
+    get_server_url)
 from rhodecode.lib.pidlock import DaemonLock, LockHeld
 from rhodecode.lib.vcs import connect_vcs
 from rhodecode.model import meta
+from rhodecode.lib.auth import AuthUser
 
 log = logging.getLogger(__name__)
 
 
 class ResultWrapper(object):
     def __init__(self, task):
         self.task = task
 
     @LazyProperty
     def result(self):
         return self.task
 
 
+class RhodecodeCeleryTask(Task):
+    """
+    This is a celery task which will create a rhodecode app instance context
+    for the task, patch pyramid + pylons threadlocals with the original request
+    that created the task and also add the user to the context.
+
+    This class as a whole should be removed once the pylons port is complete
+    and a pyramid only solution for celery is implemented as per issue #4139
+    """
+
+    def apply_async(self, args=None, kwargs=None, task_id=None, producer=None,
+                    link=None, link_error=None, **options):
+        """ queue the job to run (we are in web request context here) """
+
+        request = get_current_request()
+
+        # we hook into kwargs since it is the only way to pass our data to the
+        # celery worker in celery 2.2
+        kwargs.update({
+            '_rhodecode_proxy_data': {
+                'environ': {
+                    'PATH_INFO': request.environ['PATH_INFO'],
+                    'SCRIPT_NAME': request.environ['SCRIPT_NAME'],
+                    'HTTP_HOST': request.environ.get('HTTP_HOST',
+                        request.environ['SERVER_NAME']),
+                    'SERVER_NAME': request.environ['SERVER_NAME'],
+                    'SERVER_PORT': request.environ['SERVER_PORT'],
+                    'wsgi.url_scheme': request.environ['wsgi.url_scheme'],
+                },
+                'auth_user': {
+                    'ip_addr': request.user.ip_addr,
+                    'user_id': request.user.user_id
+                },
+            }
+        })
+        return super(RhodecodeCeleryTask, self).apply_async(
+            args, kwargs, task_id, producer, link, link_error, **options)
+
+    def __call__(self, *args, **kwargs):
+        """ rebuild the context and then run task on celery worker """
+        proxy_data = kwargs.pop('_rhodecode_proxy_data', {})
+
+        if not proxy_data:
+            return super(RhodecodeCeleryTask, self).__call__(*args, **kwargs)
+
+        log.debug('using celery proxy data to run task: %r', proxy_data)
+
+        from rhodecode.config.routing import make_map
+        from rhodecode.config.middleware import make_pyramid_app
+
+        # TODO: this can be done once per worker versus per task
+        pyramid_app = make_pyramid_app(config, **config['app_conf'])
+
+        request = Request.blank('/', environ=proxy_data['environ'])
+        request.user = AuthUser(user_id=proxy_data['auth_user']['user_id'],
+                                ip_addr=proxy_data['auth_user']['ip_addr'])
+
+        pyramid_request = prepare(request) # set pyramid threadlocal request
+
+        # pylons routing
+        if not rhodecode.CONFIG.get('routes.map'):
+            rhodecode.CONFIG['routes.map'] = make_map(config)
+        pylons.url._push_object(get_routes_generator_for_server_url(
+            get_server_url(request.environ)
+        ))
+
+        try:
+            return super(RhodecodeCeleryTask, self).__call__(*args, **kwargs)
+        finally:
+            pyramid_request['closer']()
+            pylons.url._pop_object()
+
+
 def run_task(task, *args, **kwargs):
     if rhodecode.CELERY_ENABLED:
         celery_is_up = False
         try:
             t = task.apply_async(args=args, kwargs=kwargs)
             log.info('running task %s:%s', t.task_id, task)
             celery_is_up = True
             return t
 
         except socket.error as e:
             if isinstance(e, IOError) and e.errno == 111:
                 log.error('Unable to connect to celeryd. Sync execution')
             else:
                 log.exception("Exception while connecting to celeryd.")
         except KeyError as e:
             log.error('Unable to connect to celeryd. Sync execution')
         except Exception as e:
             log.exception(
                 "Exception while trying to run task asynchronous. "
                 "Fallback to sync execution.")
 
         # keep in mind there maybe a subtle race condition where something
         # depending on rhodecode.CELERY_ENABLED such as @dbsession decorator
         # will see CELERY_ENABLED as True before this has a chance to set False
         rhodecode.CELERY_ENABLED = celery_is_up
     else:
         log.debug('executing task %s in sync mode', task)
     return ResultWrapper(task(*args, **kwargs))
 
 
 def __get_lockkey(func, *fargs, **fkwargs):
     params = list(fargs)
     params.extend(['%s-%s' % ar for ar in fkwargs.items()])
 
     func_name = str(func.__name__) if hasattr(func, '__name__') else str(func)
     _lock_key = func_name + '-' + '-'.join(map(safe_str, params))
     return 'task_%s.lock' % (md5_safe(_lock_key),)
 
 
 def locked_task(func):
     def __wrapper(func, *fargs, **fkwargs):
         lockkey = __get_lockkey(func, *fargs, **fkwargs)
         lockkey_path = config['app_conf']['cache_dir']
 
         log.info('running task with lockkey %s' % lockkey)
         try:
             l = DaemonLock(file_=jn(lockkey_path, lockkey))
             ret = func(*fargs, **fkwargs)
             l.release()
             return ret
         except LockHeld:
             log.info('LockHeld')
             return 'Task with key %s already running' % lockkey
 
     return decorator(__wrapper, func)
 
 
 def get_session():
     if rhodecode.CELERY_ENABLED:
         utils.initialize_database(config)
     sa = meta.Session()
     return sa
 
 
 def dbsession(func):
     def __wrapper(func, *fargs, **fkwargs):
         try:
             ret = func(*fargs, **fkwargs)
             return ret
         finally:
             if rhodecode.CELERY_ENABLED and not rhodecode.CELERY_EAGER:
                 meta.Session.remove()
 
     return decorator(__wrapper, func)
 
 
 def vcsconnection(func):
     def __wrapper(func, *fargs, **fkwargs):
         if rhodecode.CELERY_ENABLED and not rhodecode.CELERY_EAGER:
             backends = config['vcs.backends'] = aslist(
                 config.get('vcs.backends', 'hg,git'), sep=',')
             for alias in rhodecode.BACKENDS.keys():
                 if alias not in backends:
                     del rhodecode.BACKENDS[alias]
             utils.configure_pyro4(config)
             utils.configure_vcs(config)
             connect_vcs(
                 config['vcs.server'],
                 utils.get_vcs_server_protocol(config))
         ret = func(*fargs, **fkwargs)
         return ret
 
     return decorator(__wrapper, func)

rhodecode/lib/celerylib/tasks.py

@@ -1,284 +1,284 b''
 # -*- coding: utf-8 -*-
 
 # Copyright (C) 2012-2016  RhodeCode GmbH
 #
 # This program is free software: you can redistribute it and/or modify
 # it under the terms of the GNU Affero General Public License, version 3
 # (only), as published by the Free Software Foundation.
 #
 # This program is distributed in the hope that it will be useful,
 # but WITHOUT ANY WARRANTY; without even the implied warranty of
 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 # GNU General Public License for more details.
 #
 # You should have received a copy of the GNU Affero General Public License
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
 #
 # This program is dual-licensed. If you wish to learn more about the
 # RhodeCode Enterprise Edition, including its added features, Support services,
 # and proprietary license terms, please see https://rhodecode.com/licenses/
 
 """
 RhodeCode task modules, containing all task that suppose to be run
 by celery daemon
 """
 
 
 import os
 import logging
 
 from celery.task import task
 from pylons import config
 
 import rhodecode
 from rhodecode.lib.celerylib import (
     run_task, dbsession, __get_lockkey, LockHeld, DaemonLock,
-    get_session, vcsconnection)
+    get_session, vcsconnection, RhodecodeCeleryTask)
 from rhodecode.lib.hooks_base import log_create_repository
 from rhodecode.lib.rcmail.smtp_mailer import SmtpMailer
 from rhodecode.lib.utils import add_cache, action_logger
 from rhodecode.lib.utils2 import safe_int, str2bool
 from rhodecode.model.db import Repository, User
 
 
 add_cache(config)  # pragma: no cover
 
 
 def get_logger(cls):
     if rhodecode.CELERY_ENABLED:
         try:
             log = cls.get_logger()
         except Exception:
             log = logging.getLogger(__name__)
     else:
         log = logging.getLogger(__name__)
 
     return log
 
 
-@task(ignore_result=True)
+@task(ignore_result=True, base=RhodecodeCeleryTask)
 @dbsession
 def send_email(recipients, subject, body='', html_body='', email_config=None):
     """
     Sends an email with defined parameters from the .ini files.
 
     :param recipients: list of recipients, it this is empty the defined email
         address from field 'email_to' is used instead
     :param subject: subject of the mail
     :param body: body of the mail
     :param html_body: html version of body
     """
     log = get_logger(send_email)
 
     email_config = email_config or config
     subject = "%s %s" % (email_config.get('email_prefix', ''), subject)
     if not recipients:
         # if recipients are not defined we send to email_config + all admins
         admins = [
             u.email for u in User.query().filter(User.admin == True).all()]
         recipients = [email_config.get('email_to')] + admins
 
     mail_server = email_config.get('smtp_server') or None
     if mail_server is None:
         log.error("SMTP server information missing. Sending email failed. "
                   "Make sure that `smtp_server` variable is configured "
                   "inside the .ini file")
         return False
 
     mail_from = email_config.get('app_email_from', 'RhodeCode')
     user = email_config.get('smtp_username')
     passwd = email_config.get('smtp_password')
     mail_port = email_config.get('smtp_port')
     tls = str2bool(email_config.get('smtp_use_tls'))
     ssl = str2bool(email_config.get('smtp_use_ssl'))
     debug = str2bool(email_config.get('debug'))
     smtp_auth = email_config.get('smtp_auth')
 
     try:
         m = SmtpMailer(mail_from, user, passwd, mail_server, smtp_auth,
                        mail_port, ssl, tls, debug=debug)
         m.send(recipients, subject, body, html_body)
     except Exception:
         log.exception('Mail sending failed')
         return False
     return True
 
 
-@task(ignore_result=False)
+@task(ignore_result=True, base=RhodecodeCeleryTask)
 @dbsession
 @vcsconnection
 def create_repo(form_data, cur_user):
     from rhodecode.model.repo import RepoModel
     from rhodecode.model.user import UserModel
     from rhodecode.model.settings import SettingsModel
 
     log = get_logger(create_repo)
     DBS = get_session()
 
     cur_user = UserModel(DBS)._get_user(cur_user)
     owner = cur_user
 
     repo_name = form_data['repo_name']
     repo_name_full = form_data['repo_name_full']
     repo_type = form_data['repo_type']
     description = form_data['repo_description']
     private = form_data['repo_private']
     clone_uri = form_data.get('clone_uri')
     repo_group = safe_int(form_data['repo_group'])
     landing_rev = form_data['repo_landing_rev']
     copy_fork_permissions = form_data.get('copy_permissions')
     copy_group_permissions = form_data.get('repo_copy_permissions')
     fork_of = form_data.get('fork_parent_id')
     state = form_data.get('repo_state', Repository.STATE_PENDING)
 
     # repo creation defaults, private and repo_type are filled in form
     defs = SettingsModel().get_default_repo_settings(strip_prefix=True)
     enable_statistics = form_data.get(
         'enable_statistics', defs.get('repo_enable_statistics'))
     enable_locking = form_data.get(
         'enable_locking', defs.get('repo_enable_locking'))
     enable_downloads = form_data.get(
         'enable_downloads', defs.get('repo_enable_downloads'))
 
     try:
         RepoModel(DBS)._create_repo(
             repo_name=repo_name_full,
             repo_type=repo_type,
             description=description,
             owner=owner,
             private=private,
             clone_uri=clone_uri,
             repo_group=repo_group,
             landing_rev=landing_rev,
             fork_of=fork_of,
             copy_fork_permissions=copy_fork_permissions,
             copy_group_permissions=copy_group_permissions,
             enable_statistics=enable_statistics,
             enable_locking=enable_locking,
             enable_downloads=enable_downloads,
             state=state
         )
 
         action_logger(cur_user, 'user_created_repo',
                       repo_name_full, '', DBS)
         DBS.commit()
 
         # now create this repo on Filesystem
         RepoModel(DBS)._create_filesystem_repo(
             repo_name=repo_name,
             repo_type=repo_type,
             repo_group=RepoModel(DBS)._get_repo_group(repo_group),
             clone_uri=clone_uri,
         )
         repo = Repository.get_by_repo_name(repo_name_full)
         log_create_repository(created_by=owner.username, **repo.get_dict())
 
         # update repo commit caches initially
         repo.update_commit_cache()
 
         # set new created state
         repo.set_state(Repository.STATE_CREATED)
         DBS.commit()
     except Exception as e:
         log.warning('Exception %s occurred when creating repository, '
                     'doing cleanup...', e)
         # rollback things manually !
         repo = Repository.get_by_repo_name(repo_name_full)
         if repo:
             Repository.delete(repo.repo_id)
             DBS.commit()
             RepoModel(DBS)._delete_filesystem_repo(repo)
         raise
 
     # it's an odd fix to make celery fail task when exception occurs
     def on_failure(self, *args, **kwargs):
         pass
 
     return True
 
 
-@task(ignore_result=False)
+@task(ignore_result=True, base=RhodecodeCeleryTask)
 @dbsession
 @vcsconnection
 def create_repo_fork(form_data, cur_user):
     """
     Creates a fork of repository using internal VCS methods
 
     :param form_data:
     :param cur_user:
     """
     from rhodecode.model.repo import RepoModel
     from rhodecode.model.user import UserModel
 
     log = get_logger(create_repo_fork)
     DBS = get_session()
 
     cur_user = UserModel(DBS)._get_user(cur_user)
     owner = cur_user
 
     repo_name = form_data['repo_name']  # fork in this case
     repo_name_full = form_data['repo_name_full']
     repo_type = form_data['repo_type']
     description = form_data['description']
     private = form_data['private']
     clone_uri = form_data.get('clone_uri')
     repo_group = safe_int(form_data['repo_group'])
     landing_rev = form_data['landing_rev']
     copy_fork_permissions = form_data.get('copy_permissions')
     fork_id = safe_int(form_data.get('fork_parent_id'))
 
     try:
         fork_of = RepoModel(DBS)._get_repo(fork_id)
         RepoModel(DBS)._create_repo(
             repo_name=repo_name_full,
             repo_type=repo_type,
             description=description,
             owner=owner,
             private=private,
             clone_uri=clone_uri,
             repo_group=repo_group,
             landing_rev=landing_rev,
             fork_of=fork_of,
             copy_fork_permissions=copy_fork_permissions
         )
         action_logger(cur_user, 'user_forked_repo:%s' % repo_name_full,
                       fork_of.repo_name, '', DBS)
         DBS.commit()
 
         base_path = Repository.base_path()
         source_repo_path = os.path.join(base_path, fork_of.repo_name)
 
         # now create this repo on Filesystem
         RepoModel(DBS)._create_filesystem_repo(
             repo_name=repo_name,
             repo_type=repo_type,
             repo_group=RepoModel(DBS)._get_repo_group(repo_group),
             clone_uri=source_repo_path,
         )
         repo = Repository.get_by_repo_name(repo_name_full)
         log_create_repository(created_by=owner.username, **repo.get_dict())
 
         # update repo commit caches initially
         config = repo._config
         config.set('extensions', 'largefiles', '')
         repo.update_commit_cache(config=config)
 
         # set new created state
         repo.set_state(Repository.STATE_CREATED)
         DBS.commit()
     except Exception as e:
         log.warning('Exception %s occurred when forking repository, '
                     'doing cleanup...', e)
         # rollback things manually !
         repo = Repository.get_by_repo_name(repo_name_full)
         if repo:
             Repository.delete(repo.repo_id)
             DBS.commit()
             RepoModel(DBS)._delete_filesystem_repo(repo)
         raise
 
     # it's an odd fix to make celery fail task when exception occurs
     def on_failure(self, *args, **kwargs):
         pass
 
     return True