From 81d3161a5e2ced733e4b3febc6baf58f9936f41c Mon Sep 17 00:00:00 2001
From: Paulus Schoutsen <paulus@paulusschoutsen.nl>
Date: Sun, 19 Aug 2018 17:22:09 +0200
Subject: [PATCH] Add forgiving add column (#16057)

* Add forgiving add column

* Lint
---
 .../components/recorder/migration.py          | 27 ++++++++++++++-----
 tests/components/recorder/test_migrate.py     | 25 +++++++++++++----
 2 files changed, 41 insertions(+), 11 deletions(-)

diff --git a/homeassistant/components/recorder/migration.py b/homeassistant/components/recorder/migration.py
index 939985ebfb1..5633830dc51 100644
--- a/homeassistant/components/recorder/migration.py
+++ b/homeassistant/components/recorder/migration.py
@@ -117,7 +117,13 @@ def _drop_index(engine, table_name, index_name):
 def _add_columns(engine, table_name, columns_def):
     """Add columns to a table."""
     from sqlalchemy import text
-    from sqlalchemy.exc import SQLAlchemyError
+    from sqlalchemy.exc import OperationalError
+
+    _LOGGER.info("Adding columns %s to table %s. Note: this can take several "
+                 "minutes on large databases and slow computers. Please "
+                 "be patient!",
+                 ', '.join(column.split(' ')[0] for column in columns_def),
+                 table_name)
 
     columns_def = ['ADD COLUMN {}'.format(col_def) for col_def in columns_def]
 
@@ -126,13 +132,22 @@ def _add_columns(engine, table_name, columns_def):
             table=table_name,
             columns_def=', '.join(columns_def))))
         return
-    except SQLAlchemyError:
-        pass
+    except OperationalError:
+        # Some engines support adding all columns at once,
+        # this error is when they dont'
+        _LOGGER.info('Unable to use quick column add. Adding 1 by 1.')
 
     for column_def in columns_def:
-        engine.execute(text("ALTER TABLE {table} {column_def}".format(
-            table=table_name,
-            column_def=column_def)))
+        try:
+            engine.execute(text("ALTER TABLE {table} {column_def}".format(
+                table=table_name,
+                column_def=column_def)))
+        except OperationalError as err:
+            if 'duplicate' not in str(err).lower():
+                raise
+
+            _LOGGER.warning('Column %s already exists on %s, continueing',
+                            column_def.split(' ')[0], table_name)
 
 
 def _apply_update(engine, new_version, old_version):
diff --git a/tests/components/recorder/test_migrate.py b/tests/components/recorder/test_migrate.py
index 5ac9b3adb81..1c48c261372 100644
--- a/tests/components/recorder/test_migrate.py
+++ b/tests/components/recorder/test_migrate.py
@@ -5,11 +5,11 @@ from unittest.mock import patch, call
 
 import pytest
 from sqlalchemy import create_engine
+from sqlalchemy.pool import StaticPool
 
 from homeassistant.bootstrap import async_setup_component
-from homeassistant.components.recorder import wait_connection_ready, migration
-from homeassistant.components.recorder.models import SCHEMA_VERSION
-from homeassistant.components.recorder.const import DATA_INSTANCE
+from homeassistant.components.recorder import (
+    wait_connection_ready, migration, const, models)
 from tests.components.recorder import models_original
 
 
@@ -37,8 +37,8 @@ def test_schema_update_calls(hass):
         yield from wait_connection_ready(hass)
 
     update.assert_has_calls([
-        call(hass.data[DATA_INSTANCE].engine, version+1, 0) for version
-        in range(0, SCHEMA_VERSION)])
+        call(hass.data[const.DATA_INSTANCE].engine, version+1, 0) for version
+        in range(0, models.SCHEMA_VERSION)])
 
 
 @asyncio.coroutine
@@ -65,3 +65,18 @@ def test_invalid_update():
     """Test that an invalid new version raises an exception."""
     with pytest.raises(ValueError):
         migration._apply_update(None, -1, 0)
+
+
+def test_forgiving_add_column():
+    """Test that add column will continue if column exists."""
+    engine = create_engine(
+        'sqlite://',
+        poolclass=StaticPool
+    )
+    engine.execute('CREATE TABLE hello (id int)')
+    migration._add_columns(engine, 'hello', [
+        'context_id CHARACTER(36)',
+    ])
+    migration._add_columns(engine, 'hello', [
+        'context_id CHARACTER(36)',
+    ])