Merge pull request #1643 from coreos-inc/db-retry

Enable automatic retry for the database

data/database.py

@@ -14,6 +14,8 @@ import toposort
 
 from enum import Enum
 from peewee import *
+from playhouse.shortcuts import RetryOperationalError
+
 from sqlalchemy.engine.url import make_url
 
 from data.fields import ResumableSHA256Field, ResumableSHA1Field, JSONField, Base64BinaryField
@@ -24,7 +26,7 @@ from util.names import urn_generator
 logger = logging.getLogger(__name__)
 
 
-SCHEME_DRIVERS = {
+_SCHEME_DRIVERS = {
   'mysql': MySQLDatabase,
   'mysql+pymysql': MySQLDatabase,
   'sqlite': SqliteDatabase,
@@ -190,6 +192,10 @@ def validate_database_url(url, db_kwargs, connect_timeout=5):
   driver.close()
 
 
+def _wrap_for_retry(driver):
+  return type('Retrying' + driver.__class__.__name__, (RetryOperationalError, driver), {})
+
+
 def _db_from_url(url, db_kwargs):
   parsed_url = make_url(url)
 
@@ -206,7 +212,9 @@ def _db_from_url(url, db_kwargs):
   if parsed_url.drivername == 'sqlite' and 'connect_timeout' in db_kwargs:
     del db_kwargs['connect_timeout']
 
-  return SCHEME_DRIVERS[parsed_url.drivername](parsed_url.database, **db_kwargs)
+  driver = _SCHEME_DRIVERS[parsed_url.drivername]
+  wrapped_driver = _wrap_for_retry(driver)
+  return wrapped_driver(parsed_url.database, **db_kwargs)
 
 
 def configure(config_object):