Merge pull request #2 from djuretic/feature/main_function

If worker has a main function, call it when executing runner
dperezrada · Dec 16, 2016 · bfbc707 · bfbc707
2 parents 0159ce9 + 080b14b
commit bfbc707
Showing 1 changed file with 13 additions and 8 deletions.
diff --git a/poliglo/runner.py b/poliglo/runner.py
@@ -125,11 +125,16 @@ def default_main_inside(
     info = imp.find_module(module_name, [basepath])
     module = imp.load_module(module_name, *info)
 
-    config = get_config(environ.get('POLIGLO_SERVER_URL'), 'all')
-    connection = get_connection(config)
-    default_main(
-        environ.get('POLIGLO_SERVER_URL'),
-        module_name,
-        module.process,
-        {'connection': connection}
-    )
+    # wait_jobs.py has a main function
+    main_function = getattr(module, 'main', None)
+    if main_function:
+        main_function()
+    else:
+        config = get_config(environ.get('POLIGLO_SERVER_URL'), 'all')
+        connection = get_connection(config)
+        default_main(
+            environ.get('POLIGLO_SERVER_URL'),
+            module_name,
+            module.process,
+            {'connection': connection}
+        )