From 700e7b74e4379f376beeee25f18932f6d4fa22bb Mon Sep 17 00:00:00 2001 From: Joseph Schorr Date: Fri, 22 Jul 2016 14:04:13 -0400 Subject: [PATCH] Enable automatic retry for the database --- data/database.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/data/database.py b/data/database.py index a00b8a139..fc3f9d7bd 100644 --- a/data/database.py +++ b/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):