Local builds: fix logging to the module build log file

Ever since local builds were changed to call handlers directly
instead of going through the scheduler, the current module ID wasn't
set, causing no logs to be written to the module build log file.

module_build_service/scheduler/consumer.py

@@ -6,6 +6,7 @@ to use.
 """
 
 from __future__ import absolute_import
+from contextlib import contextmanager
 import itertools
 
 try:
@@ -254,39 +255,49 @@ class MBSConsumer(fedmsg.consumers.FedmsgConsumer):
             log.debug("No module associated with msg %s", event_info["msg_id"])
             return
 
-        MBSConsumer.current_module_build_id = build.id
+        with set_current_module_build_id(build.id):
+            log.info("Calling %s", idx)
 
-        log.info("Calling %s", idx)
+            kwargs = event_info.copy()
+            kwargs.pop("event")
 
-        kwargs = event_info.copy()
-        kwargs.pop("event")
+            try:
+                if conf.celery_broker_url:
+                    # handlers are also Celery tasks, when celery_broker_url is configured,
+                    # call "delay" method to run the handlers as Celery async tasks
+                    func = getattr(handler, "delay")
+                    func(**kwargs)
+                else:
+                    handler(**kwargs)
+            except Exception as e:
+                log.exception("Could not process message handler.")
+                db_session.rollback()
+                db_session.refresh(build)
+                build.transition(
+                    db_session,
+                    conf,
+                    state=models.BUILD_STATES["failed"],
+                    state_reason=str(e),
+                    failure_type="infra",
+                )
+                db_session.commit()
 
-        try:
-            if conf.celery_broker_url:
-                # handlers are also Celery tasks, when celery_broker_url is configured,
-                # call "delay" method to run the handlers as Celery async tasks
-                func = getattr(handler, "delay")
-                func(**kwargs)
-            else:
-                handler(**kwargs)
-        except Exception as e:
-            log.exception("Could not process message handler.")
-            db_session.rollback()
-            db_session.refresh(build)
-            build.transition(
-                db_session,
-                conf,
-                state=models.BUILD_STATES["failed"],
-                state_reason=str(e),
-                failure_type="infra",
-            )
-            db_session.commit()
+                # Allow caller to do something when error is occurred.
+                raise
+            finally:
+                log.debug("Done with %s", idx)
 
-            # Allow caller to do something when error is occurred.
-            raise
-        finally:
-            MBSConsumer.current_module_build_id = None
-            log.debug("Done with %s", idx)
+
+@contextmanager
+def set_current_module_build_id(build_id):
+    """Context manager to temporarily set the module build ID - this is
+    needed to get logging to go to the right build log.
+    """
+    MBSConsumer.current_module_build_id = build_id
+    try:
+        yield
+    finally:
+        MBSConsumer.current_module_build_id = None
 
 
 def get_global_consumer():

module_build_service/scheduler/local.py

@@ -4,6 +4,7 @@ from __future__ import absolute_import
 import logging
 
 from module_build_service.common import models
+from module_build_service.scheduler.consumer import set_current_module_build_id
 from module_build_service.scheduler.db_session import db_session
 from module_build_service.scheduler.handlers.modules import init as modules_init_handler
 from module_build_service.scheduler.handlers.modules import wait as modules_wait_handler
@@ -33,7 +34,8 @@ def raise_for_failed_build(module_build_ids):
     has_failed_build = False
     for build in builds:
         if build.state == models.BUILD_STATES["failed"]:
-            modules_failed_handler("fake_msg_id", build.id, "failed")
+            with set_current_module_build_id(build.id):
+                modules_failed_handler("fake_msg_id", build.id, "failed")
             has_failed_build = True
     if has_failed_build:
         raise ValueError("Local module build failed.")
@@ -52,14 +54,17 @@ def main(module_build_ids):
     # the right order manually.
     # We only need to ensure that we won't call futher handlers in case the module build failed.
     for module_build_id in module_build_ids:
-        modules_init_handler("fake_msg_id", module_build_id, "init")
+        with set_current_module_build_id(module_build_id):
+            modules_init_handler("fake_msg_id", module_build_id, models.BUILD_STATES["init"])
 
     raise_for_failed_build(module_build_ids)
     for module_build_id in module_build_ids:
-        modules_wait_handler("fake_msg_id", module_build_id, "wait")
+        with set_current_module_build_id(module_build_id):
+            modules_wait_handler("fake_msg_id", module_build_id, models.BUILD_STATES["wait"])
 
     raise_for_failed_build(module_build_ids)
     for module_build_id in module_build_ids:
-        modules_done_handler("fake_msg_id", module_build_id, "done")
+        with set_current_module_build_id(module_build_id):
+            modules_done_handler("fake_msg_id", module_build_id, models.BUILD_STATES["done"])
 
     raise_for_failed_build(module_build_ids)