celery: fixed bootstrap of beat workers

rhodecode/lib/celerylib/loader.py

@@ -124,7 +124,7 @@ def on_preload_parsed(options, **kwargs):
     ini_vars = options['ini_var']
 
     if ini_file is None:
-        print('You must provide the paste --ini argument')
+        print('You must provide the --ini argument to start celery')
         exit(-1)
 
     options = None
@@ -136,20 +136,25 @@ def on_preload_parsed(options, **kwargs):
     setup_logging(ini_file)
 
 
-@signals.celeryd_init.connect
-def on_celeryd_init(sender=None, conf=None, **kwargs):
+def _init_celery(app_type=''):
     from rhodecode.config.middleware import get_celery_config
 
-    log.debug('Bootstrapping RhodeCode application...')
+    log.debug('Bootstrapping RhodeCode application for %s...', app_type)
 
-    ini_file = conf['RC_INI_FILE']
-    options = conf['RC_INI_OPTIONS']
+    ini_file = celery_app.conf['RC_INI_FILE']
+    options = celery_app.conf['RC_INI_OPTIONS']
 
+    env = None
     try:
         env = bootstrap(ini_file, options=options)
     except Exception:
         log.exception('Failed to bootstrap RhodeCode APP')
 
+    if not env:
+        raise EnvironmentError(
+            'Failed to load pyramid ENV. '
+            'Probably there is another error present that prevents from running pyramid app')
+
     log.debug('Got Pyramid ENV: %s', env)
 
     celery_settings = get_celery_config(env['registry'].settings)
@@ -159,11 +164,21 @@ def on_celeryd_init(sender=None, conf=None, **kwargs):
         registry=env['registry'], closer=env['closer'],
         celery_settings=celery_settings)
 
+
+@signals.celeryd_init.connect
+def on_celeryd_init(sender=None, conf=None, **kwargs):
+    _init_celery('celery worker')
+
     # fix the global flag even if it's disabled via .ini file because this
     # is a worker code that doesn't need this to be disabled.
     rhodecode.CELERY_ENABLED = True
 
 
+@signals.beat_init.connect
+def on_beat_init(sender=None, conf=None, **kwargs):
+    _init_celery('celery beat')
+
+
 @signals.task_prerun.connect
 def task_prerun_signal(task_id, task, args, **kwargs):
     ping_db()