Add simple & supported way to setup MRQ's context. Fixes #51

Author: sylvinus

docs/memory-leaks.md

@@ -20,9 +20,8 @@ First, initialize a REPL with MRQ configured and guppy loaded:
 ```
 $ pip install guppy
 $ python
->>> from mrq import config
->>> from mrq.context import set_current_config, run_task
->>> set_current_config(config.get_config(sources=("file", "env"), config_type="run"))
+>>> from mrq.context import setup_context, run_task
+>>> setup_context()
 >>> from guppy import hpy
 >>> hp = hpy()
 ```

mrq/bin/mrq_run.py

@@ -27,7 +27,7 @@ def main():
 
     parser = argparse.ArgumentParser(description='Runs a task')
 
-    cfg = config.get_config(parser=parser, config_type="run")
+    cfg = config.get_config(parser=parser, config_type="run", sources=("file", "env", "args"))
     cfg["is_cli"] = True
     set_current_config(cfg)
 

mrq/bin/mrq_worker.py

@@ -21,13 +21,14 @@
 
 from mrq import config
 from mrq.utils import load_class_by_path
+from mrq.context import set_current_config
 
 
 def main():
 
     parser = argparse.ArgumentParser(description='Start a RQ worker')
 
-    cfg = config.get_config(parser=parser, config_type="worker")
+    cfg = config.get_config(parser=parser, config_type="worker", sources=("file", "env", "args"))
 
     # If we are launching with a --processes option and without the SUPERVISOR_ENABLED env
     # then we should just call supervisord.
@@ -100,7 +101,9 @@ def sigterm_handler(signum, frame):  # pylint: disable=unused-argument
 
         worker_class = load_class_by_path(cfg["worker_class"])
 
-        w = worker_class(cfg)
+        set_current_config(cfg)
+
+        w = worker_class()
 
         exitcode = w.work_loop()
 

mrq/config.py

@@ -300,12 +300,11 @@ def add_parser_args(parser, config_type):
 def get_config(
         sources=(
             "file",
-            "env",
-            "args"),
+            "env"),
         env_prefix="MRQ_",
         file_path=None,
         parser=None,
-        config_type="worker"):
+        config_type=None):
     """ Returns a config dict merged from several possible sources """
 
     if not parser:
@@ -314,7 +313,7 @@ def get_config(
     add_parser_args(parser, config_type)
 
     if config_type in ["run"]:
-        default_config = parser.parse_args(["x"]).__dict__
+        default_config = parser.parse_args(["notask"]).__dict__
     else:
         default_config = parser.parse_args([]).__dict__
 

mrq/context.py

@@ -6,7 +6,9 @@
 import time
 from .utils import LazyObject, load_class_by_path, group_iter
 from itertools import count as itertools_count
+from .config import get_config
 
+# This should be MRQ's only Python object shared by all the jobs in the same process
 _GLOBAL_CONTEXT = {
 
     # Contains all the running greenlets for this worker. greenletid => Job object
@@ -19,10 +21,18 @@
     "config": None
 }
 
-# Global log object, usable from all tasks
+# Global log object, usable from all jobs
 log = Logger(None, job="current")
 
 
+def setup_context(**kwargs):
+    """ Setup MRQ's environment.
+
+        Note: gevent should probably be initialized too if you want to use concurrency.
+    """
+    set_current_config(get_config(**kwargs))
+
+
 def set_current_job(job):
     current = gevent.getcurrent()
 
@@ -44,7 +54,6 @@ def get_current_job(greenlet_id=None):
 
 def set_current_worker(worker):
     _GLOBAL_CONTEXT["worker"] = worker
-    set_current_config(worker.config)
 
 
 def get_current_worker():

mrq/dashboard/app.py

@@ -24,7 +24,7 @@
 
 parser = argparse.ArgumentParser(description='Start the MRQ dashboard')
 
-cfg = get_config(parser=parser, config_type="dashboard")
+cfg = get_config(parser=parser, config_type="dashboard", sources=("file", "env", "args"))
 set_current_config(cfg)
 
 app = Flask(

mrq/worker.py

@@ -16,7 +16,7 @@
 from .job import Job
 from .exceptions import (TimeoutInterrupt, StopRequested, JobInterrupt, AbortInterrupt,
                          RetryInterrupt, MaxRetriesInterrupt)
-from .context import (set_current_worker, set_current_job, get_current_job,
+from .context import (set_current_worker, set_current_job, get_current_job, get_current_config,
                       connections, enable_greenlet_tracing)
 from .queue import Queue
 
@@ -35,9 +35,7 @@ class Worker(object):
     mongodb_logs = None
     redis = None
 
-    def __init__(self, config):
-
-        self.config = config
+    def __init__(self):
 
         set_current_worker(self)
 
@@ -57,7 +55,7 @@ def __init__(self, config):
         self.graceful_stop = None
 
         self.id = ObjectId()
-        if config["name"]:
+        if self.config.get("name"):
             self.name = self.config["name"]
         else:
             # Generate a somewhat human-readable name for this worker
@@ -85,6 +83,10 @@ def __init__(self, config):
             "total": 0
         }
 
+    @property
+    def config(self):
+        return get_current_config()
+
     def connect(self, force=False):
 
         if self.connected and not force:

tests/fixtures/standalone_script1.py

@@ -0,0 +1,8 @@
+from mrq.context import setup_context, run_task, get_current_config
+
+# Autoconfigure MRQ's environment
+setup_context()
+
+print run_task("tests.tasks.general.Add", {"a": 41, "b": 1})
+
+print get_current_config()["name"]

tests/test_context.py

@@ -1,5 +1,6 @@
 import json
 import time
+import os
 
 
 def test_context_get(worker):
@@ -80,3 +81,18 @@ def test_context_metric_failed(worker):
     assert metrics.get("jobs.status.started") == 2
     assert metrics.get("jobs.status.failed") == 1
     assert metrics.get("jobs.status.success") is None
+
+
+def test_context_setup():
+
+    try:
+        import subprocess32 as subprocess
+    except:
+        import subprocess
+
+    process = subprocess.Popen("python tests/fixtures/standalone_script1.py",
+                               shell=True, close_fds=True, env={"MRQ_NAME": "testname1", "PYTHONPATH": os.getcwd()}, cwd=os.getcwd(), stdout=subprocess.PIPE)
+
+    out, err = process.communicate()
+
+    assert out.endswith("42\ntestname1\n")