From 150b376cf90f60c08ace286f5f08549933e03ce6 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jonathan=20=C3=98strup?=
 <7877957+TechnicallyJoe@users.noreply.github.com>
Date: Sun, 9 Feb 2020 17:43:09 +0100
Subject: [PATCH] Add recorder vars db_max_retries and db_retry_wait (#31561)

* added recorder vars db_max_retries and db_retry_wait

* fixed test_recorder_setup_failure

I failed because it was missing the two new variables. I simply added these with default values.

* fixed syntax error in test_recorder_setup_failure

* fixed formatting error in test_init_py for recorder component

* fixed typo in test case

* Updated the way the default keys for db_,max_wait and db_retry_wait is set

Implemented based on suggestions from @springstan

* Updated config_schema call to adhere to Black

* changed conf.get to conf[dict] for var retrieval

* removed 2 blank lines
---
 homeassistant/components/recorder/__init__.py | 32 ++++++++++++++-----
 tests/components/recorder/test_init.py        |  9 +++++-
 2 files changed, 32 insertions(+), 9 deletions(-)

diff --git a/homeassistant/components/recorder/__init__.py b/homeassistant/components/recorder/__init__.py
index ab56a5fc33b..af34d4dd9f6 100644
--- a/homeassistant/components/recorder/__init__.py
+++ b/homeassistant/components/recorder/__init__.py
@@ -59,14 +59,16 @@ SERVICE_PURGE_SCHEMA = vol.Schema(
 
 DEFAULT_URL = "sqlite:///{hass_config_path}"
 DEFAULT_DB_FILE = "home-assistant_v2.db"
+DEFAULT_DB_MAX_RETRIES = 10
+DEFAULT_DB_RETRY_WAIT = 3
 
 CONF_DB_URL = "db_url"
+CONF_DB_MAX_RETRIES = "db_max_retries"
+CONF_DB_RETRY_WAIT = "db_retry_wait"
 CONF_PURGE_KEEP_DAYS = "purge_keep_days"
 CONF_PURGE_INTERVAL = "purge_interval"
 CONF_EVENT_TYPES = "event_types"
 
-CONNECT_RETRY_WAIT = 3
-
 FILTER_SCHEMA = vol.Schema(
     {
         vol.Optional(CONF_EXCLUDE, default={}): vol.Schema(
@@ -96,6 +98,12 @@ CONFIG_SCHEMA = vol.Schema(
                     vol.Coerce(int), vol.Range(min=0)
                 ),
                 vol.Optional(CONF_DB_URL): cv.string,
+                vol.Optional(
+                    CONF_DB_MAX_RETRIES, default=DEFAULT_DB_MAX_RETRIES
+                ): cv.positive_int,
+                vol.Optional(
+                    CONF_DB_RETRY_WAIT, default=DEFAULT_DB_RETRY_WAIT
+                ): cv.positive_int,
             }
         )
     },
@@ -133,6 +141,8 @@ async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
     conf = config[DOMAIN]
     keep_days = conf.get(CONF_PURGE_KEEP_DAYS)
     purge_interval = conf.get(CONF_PURGE_INTERVAL)
+    db_max_retries = conf[CONF_DB_MAX_RETRIES]
+    db_retry_wait = conf[CONF_DB_RETRY_WAIT]
 
     db_url = conf.get(CONF_DB_URL, None)
     if not db_url:
@@ -145,6 +155,8 @@ async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
         keep_days=keep_days,
         purge_interval=purge_interval,
         uri=db_url,
+        db_max_retries=db_max_retries,
+        db_retry_wait=db_retry_wait,
         include=include,
         exclude=exclude,
     )
@@ -174,6 +186,8 @@ class Recorder(threading.Thread):
         keep_days: int,
         purge_interval: int,
         uri: str,
+        db_max_retries: int,
+        db_retry_wait: int,
         include: Dict,
         exclude: Dict,
     ) -> None:
@@ -186,6 +200,8 @@ class Recorder(threading.Thread):
         self.queue: Any = queue.Queue()
         self.recording_start = dt_util.utcnow()
         self.db_url = uri
+        self.db_max_retries = db_max_retries
+        self.db_retry_wait = db_retry_wait
         self.async_db_ready = asyncio.Future()
         self.engine: Any = None
         self.run_info: Any = None
@@ -217,9 +233,9 @@ class Recorder(threading.Thread):
         tries = 1
         connected = False
 
-        while not connected and tries <= 10:
+        while not connected and tries <= self.db_max_retries:
             if tries != 1:
-                time.sleep(CONNECT_RETRY_WAIT)
+                time.sleep(self.db_retry_wait)
             try:
                 self._setup_connection()
                 migration.migrate_schema(self)
@@ -230,7 +246,7 @@ class Recorder(threading.Thread):
                 _LOGGER.error(
                     "Error during connection setup: %s (retrying in %s seconds)",
                     err,
-                    CONNECT_RETRY_WAIT,
+                    self.db_retry_wait,
                 )
                 tries += 1
 
@@ -337,9 +353,9 @@ class Recorder(threading.Thread):
 
             tries = 1
             updated = False
-            while not updated and tries <= 10:
+            while not updated and tries <= self.db_max_retries:
                 if tries != 1:
-                    time.sleep(CONNECT_RETRY_WAIT)
+                    time.sleep(self.db_retry_wait)
                 try:
                     with session_scope(session=self.get_session()) as session:
                         try:
@@ -367,7 +383,7 @@ class Recorder(threading.Thread):
                         "Error in database connectivity: %s. "
                         "(retrying in %s seconds)",
                         err,
-                        CONNECT_RETRY_WAIT,
+                        self.db_retry_wait,
                     )
                     tries += 1
 
diff --git a/tests/components/recorder/test_init.py b/tests/components/recorder/test_init.py
index ae04066651f..a21ef578ca9 100644
--- a/tests/components/recorder/test_init.py
+++ b/tests/components/recorder/test_init.py
@@ -198,7 +198,14 @@ def test_recorder_setup_failure():
     ):
         setup.side_effect = ImportError("driver not found")
         rec = Recorder(
-            hass, keep_days=7, purge_interval=2, uri="sqlite://", include={}, exclude={}
+            hass,
+            keep_days=7,
+            purge_interval=2,
+            uri="sqlite://",
+            db_max_retries=10,
+            db_retry_wait=3,
+            include={},
+            exclude={},
         )
         rec.start()
         rec.join()