From d5bbb5748176f54be40d960a8501dbd55bdbbe3e Mon Sep 17 00:00:00 2001
From: Joseph Schorr <joseph.schorr@coreos.com>
Date: Thu, 6 Nov 2014 18:00:52 -0500
Subject: [PATCH] Change registry code to disconnect from the DB before long
 I/O operations

---
 data/database.py      | 17 +++++++++++++++++
 endpoints/registry.py | 13 ++++++-------
 endpoints/verbs.py    |  3 +++
 3 files changed, 26 insertions(+), 7 deletions(-)

diff --git a/data/database.py b/data/database.py
index 3966c30cc..c466a6d3f 100644
--- a/data/database.py
+++ b/data/database.py
@@ -35,6 +35,22 @@ class CallableProxy(Proxy):
       raise AttributeError('Cannot use uninitialized Proxy.')
     return self.obj(*args, **kwargs)
 
+
+class CloseForLongOperation(object):
+  """ Helper object which disconnects the database then reconnects after the nested operation
+      completes.
+  """
+
+  def __init__(self, config_object):
+    self.config_object = config_object
+
+  def __enter__(self):
+    close_db_filter(None)
+
+  def __exit__(self, type, value, traceback):
+    configure(self.config_object)
+
+
 class UseThenDisconnect(object):
   """ Helper object for conducting work with a database and then tearing it down. """
 
@@ -69,6 +85,7 @@ def _db_from_url(url, db_kwargs):
 
 
 def configure(config_object):
+  logger.debug('Configuring database')
   db_kwargs = dict(config_object['DB_CONNECTION_ARGS'])
   write_db_uri = config_object['DB_URI']
   db.initialize(_db_from_url(write_db_uri, db_kwargs))
diff --git a/endpoints/registry.py b/endpoints/registry.py
index 85b4b7963..6e6821ed0 100644
--- a/endpoints/registry.py
+++ b/endpoints/registry.py
@@ -13,7 +13,7 @@ from util import checksums, changes
 from util.http import abort, exact_abort
 from auth.permissions import (ReadRepositoryPermission,
                               ModifyRepositoryPermission)
-from data import model
+from data import model, database
 from util import gzipstream
 
 
@@ -152,8 +152,9 @@ def get_image_layer(namespace, repository, image_id, headers):
 
       profile.debug('Streaming layer data')
 
-      # TODO: DATABASE: We should disconnect from the database here, so that
-      # we're not holding the DB handle during this long download.
+      # Close the database handle here for this process before we send the long download.
+      database.close_db_filter(None)
+
       return Response(store.stream_read(repo_image.storage.locations, path), headers=headers)
     except (IOError, AttributeError):
       profile.debug('Image not found')
@@ -215,11 +216,9 @@ def put_image_layer(namespace, repository, image_id):
   h, sum_hndlr = checksums.simple_checksum_handler(json_data)
   sr.add_handler(sum_hndlr)
 
-  # TODO: DATABASE: We should disconnect from the database here and reconnect AFTER, so that
-  # we're not holding the DB handle during this long upload.
-
   # Stream write the data to storage.
-  store.stream_write(repo_image.storage.locations, layer_path, sr)
+  with database.CloseForLongOperation(app.config):
+    store.stream_write(repo_image.storage.locations, layer_path, sr)
 
   # Append the computed checksum.
   csums = []
diff --git a/endpoints/verbs.py b/endpoints/verbs.py
index b38a90d6f..92f3fa24d 100644
--- a/endpoints/verbs.py
+++ b/endpoints/verbs.py
@@ -101,6 +101,9 @@ def get_squashed_tag(namespace, repository, tag):
         logger.debug('Redirecting to download URL for derived image %s', derived.uuid)
         return redirect(download_url)
 
+      # Close the database handle here for this process before we send the long download.
+      database.close_db_filter(None)
+
       logger.debug('Sending cached derived image %s', derived.uuid)
       return send_file(store.stream_read_file(derived.locations, derived_layer_path))