From 0091d374690157a41c72d44aa09bbda9e7a3866b Mon Sep 17 00:00:00 2001
From: mvdbeek <m.vandenbeek@gmail.com>
Date: Mon, 29 Nov 2021 10:23:47 +0100
Subject: [PATCH] Add interactive shell for celery tasks

That's pretty useful for developing new tasks or just running one-off
tasks.
---
 scripts/celery_shell.py | 30 +++++++++++++++++++++++++++---
 scripts/db_shell.py     |  3 ++-
 2 files changed, 29 insertions(+), 4 deletions(-)

diff --git a/scripts/celery_shell.py b/scripts/celery_shell.py
index 34c98e8ff12..9df3cb8b1f9 100644
--- a/scripts/celery_shell.py
+++ b/scripts/celery_shell.py
@@ -1,8 +1,32 @@
-from scripts.db_shell import *
+"""
+Run this script from galaxy's root with
+```
+ipython -i scripts/celery_shell.py -- -c config/galaxy.yml
+```
+"""
+import logging
+import os
+
+WARNING_MODULES = ['parso', 'asyncio', 'galaxy.datatypes']
+for mod in WARNING_MODULES:
+    logger = logging.getLogger(mod)
+    logger.setLevel('WARNING')
+
+from scripts.db_shell import config
+os.environ['GALAXY_CONFIG_FILE'] = os.environ.get('GALAXY_CONFIG_FILE', config['config_file'])
 
 from galaxy.celery import get_galaxy_app
+from galaxy.celery import tasks  # noqa: F401
 
-from galaxy.celery.tasks import *
+HELP = """
+============
+Run celery tasks interactively.
+tasks are collected in task module.
+
+To run recalculate_user_disk_usage for user 1 in a celery worker
+type
+>>> tasks.recalculate_user_disk_usage.deley(user_id=1)
+"""
 
 app = get_galaxy_app()
-assert app, 'Set GALAXY_CONFIG_FILE to the path of your galaxy.yml file'
+print(HELP)
diff --git a/scripts/db_shell.py b/scripts/db_shell.py
index b5754ee00c0..07ac89956bb 100644
--- a/scripts/db_shell.py
+++ b/scripts/db_shell.py
@@ -36,7 +36,8 @@ from galaxy.model.orm.scripts import get_config
 registry = Registry()
 registry.load_datatypes()
 set_datatypes_registry(registry)
-db_url = get_config(sys.argv)['db_url']
+config = get_config(sys.argv)
+db_url = config['db_url']
 sa_session = init('/tmp/', db_url).context