Merge pull request #2824 from coreos-inc/close-build-connection

Close the database connection after operations in buildman

buildman/component/buildcomponent.py

@@ -15,8 +15,9 @@ from buildman.jobutil.buildjob import BuildJobLoadException
 from buildman.jobutil.buildstatus import StatusHandler
 from buildman.jobutil.workererror import WorkerError
 
+from app import app
 from data import model
-from data.database import BUILD_PHASE
+from data.database import BUILD_PHASE, UseThenDisconnect
 from data.model import InvalidRepositoryBuildException
 from util import slash_join
 
@@ -358,6 +359,7 @@ class BuildComponent(BaseComponent):
       # Label the pushed manifests with the build metadata.
       manifest_digests = kwargs.get('digests') or []
       for digest in manifest_digests:
+        with UseThenDisconnect(app.config):
           try:
             manifest = model.tag.load_manifest_by_digest(self._current_job.namespace,
                                                          self._current_job.repo_name, digest)

buildman/jobutil/buildjob.py

@@ -1,9 +1,11 @@
 import json
 import logging
 
+from app import app
 from cachetools import lru_cache
 from notifications import spawn_notification
 from data import model
+from data.database import UseThenDisconnect
 from util.imagetree import ImageTree
 from util.morecollections import AttrDict
 
@@ -40,6 +42,7 @@ class BuildJob(object):
 
   @lru_cache(maxsize=1)
   def _load_repo_build(self):
+    with UseThenDisconnect(app.config):
       try:
         return model.build.get_repository_build(self.build_uuid)
       except model.InvalidRepositoryBuildException:
@@ -108,6 +111,7 @@ class BuildJob(object):
         starting at the given base_image_id. This mimics the Docker cache checking, so it should,
         in theory, provide "perfect" caching.
     """
+    with UseThenDisconnect(app.config):
       # Lookup the base image in the repository. If it doesn't exist, nothing more to do.
       repo_build = self.repo_build
       repo_namespace = repo_build.repository.namespace_user.username
@@ -147,6 +151,7 @@ class BuildJob(object):
     """ Determines the cached tag by looking for one of the tags being built, and seeing if it
         exists in the repository. This is a fallback for when no comment information is available.
     """
+    with UseThenDisconnect(app.config):
       tags = self.build_config.get('docker_tags', ['latest'])
       repository = self.repo_build.repository
       existing_tags = model.tag.list_repository_tags(repository.namespace_user.username,
@@ -186,6 +191,7 @@ class BuildJobNotifier(object):
       )
 
   def send_notification(self, kind, error_message=None, image_id=None, manifest_digests=None):
+    with UseThenDisconnect(app.config):
       tags = self.build_config.get('docker_tags', ['latest'])
       event_data = {
         'build_id': self.repo_build.uuid,

data/database.py

@@ -168,6 +168,7 @@ class CloseForLongOperation(object):
     self.config_object = config_object
 
   def __enter__(self):
+    # TODO(jschorr): Remove this stupid hack.
     if self.config_object.get('TESTING') is True:
       return
 
@@ -185,9 +186,17 @@ class UseThenDisconnect(object):
     self.config_object = config_object
 
   def __enter__(self):
+    # TODO(jschorr): Remove this stupid hack.
+    if self.config_object.get('TESTING') is True:
+      return
+
     configure(self.config_object)
 
   def __exit__(self, typ, value, traceback):
+    # TODO(jschorr): Remove this stupid hack.
+    if self.config_object.get('TESTING') is True:
+      return
+
     close_db_filter(None)