Improve sql config flow (home-assistant#150757)

Co-authored-by: Joost Lekkerkerker <[email protected]>

Author: gjohansson-ST

homeassistant/components/sql/__init__.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import logging
+from typing import Any
 
 import sqlparse
 import voluptuous as vol
@@ -32,7 +33,13 @@
 )
 from homeassistant.helpers.typing import ConfigType
 
-from .const import CONF_COLUMN_NAME, CONF_QUERY, DOMAIN, PLATFORMS
+from .const import (
+    CONF_ADVANCED_OPTIONS,
+    CONF_COLUMN_NAME,
+    CONF_QUERY,
+    DOMAIN,
+    PLATFORMS,
+)
 from .util import redact_credentials
 
 _LOGGER = logging.getLogger(__name__)
@@ -75,18 +82,6 @@ def validate_sql_select(value: str) -> str:
 )
 
 
-def remove_configured_db_url_if_not_needed(
-    hass: HomeAssistant, entry: ConfigEntry
-) -> None:
-    """Remove db url from config if it matches recorder database."""
-    hass.config_entries.async_update_entry(
-        entry,
-        options={
-            key: value for key, value in entry.options.items() if key != CONF_DB_URL
-        },
-    )
-
-
 async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
     """Set up SQL from yaml config."""
     if (conf := config.get(DOMAIN)) is None:
@@ -107,8 +102,6 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
         redact_credentials(entry.options.get(CONF_DB_URL)),
         redact_credentials(get_instance(hass).db_url),
     )
-    if entry.options.get(CONF_DB_URL) == get_instance(hass).db_url:
-        remove_configured_db_url_if_not_needed(hass, entry)
 
     await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
 
@@ -119,3 +112,47 @@ async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Unload SQL config entry."""
 
     return await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
+
+
+async def async_migrate_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+    """Migrate old entry."""
+    _LOGGER.debug("Migrating from version %s.%s", entry.version, entry.minor_version)
+
+    if entry.version > 1:
+        # This means the user has downgraded from a future version
+        return False
+
+    if entry.version == 1:
+        old_options = {**entry.options}
+        new_data = {}
+        new_options: dict[str, Any] = {}
+
+        if (db_url := old_options.get(CONF_DB_URL)) and db_url != get_instance(
+            hass
+        ).db_url:
+            new_data[CONF_DB_URL] = db_url
+
+        new_options[CONF_COLUMN_NAME] = old_options.get(CONF_COLUMN_NAME)
+        new_options[CONF_QUERY] = old_options.get(CONF_QUERY)
+        new_options[CONF_ADVANCED_OPTIONS] = {}
+
+        for key in (
+            CONF_VALUE_TEMPLATE,
+            CONF_UNIT_OF_MEASUREMENT,
+            CONF_DEVICE_CLASS,
+            CONF_STATE_CLASS,
+        ):
+            if (value := old_options.get(key)) is not None:
+                new_options[CONF_ADVANCED_OPTIONS][key] = value
+
+        hass.config_entries.async_update_entry(
+            entry, data=new_data, options=new_options, version=2
+        )
+
+    _LOGGER.debug(
+        "Migration to version %s.%s successful",
+        entry.version,
+        entry.minor_version,
+    )
+
+    return True

homeassistant/components/sql/config_flow.py

@@ -6,7 +6,7 @@
 from typing import Any
 
 import sqlalchemy
-from sqlalchemy.engine import Result
+from sqlalchemy.engine import Engine, Result
 from sqlalchemy.exc import MultipleResultsFound, NoSuchColumnError, SQLAlchemyError
 from sqlalchemy.orm import Session, scoped_session, sessionmaker
 import sqlparse
@@ -32,59 +32,59 @@
     CONF_VALUE_TEMPLATE,
 )
 from homeassistant.core import callback
+from homeassistant.data_entry_flow import section
 from homeassistant.helpers import selector
 
-from .const import CONF_COLUMN_NAME, CONF_QUERY, DOMAIN
+from .const import CONF_ADVANCED_OPTIONS, CONF_COLUMN_NAME, CONF_QUERY, DOMAIN
 from .util import resolve_db_url
 
 _LOGGER = logging.getLogger(__name__)
 
 
 OPTIONS_SCHEMA: vol.Schema = vol.Schema(
     {
-        vol.Optional(
-            CONF_DB_URL,
-        ): selector.TextSelector(),
-        vol.Required(
-            CONF_COLUMN_NAME,
-        ): selector.TextSelector(),
-        vol.Required(
-            CONF_QUERY,
-        ): selector.TextSelector(selector.TextSelectorConfig(multiline=True)),
-        vol.Optional(
-            CONF_UNIT_OF_MEASUREMENT,
-        ): selector.TextSelector(),
-        vol.Optional(
-            CONF_VALUE_TEMPLATE,
-        ): selector.TemplateSelector(),
-        vol.Optional(CONF_DEVICE_CLASS): selector.SelectSelector(
-            selector.SelectSelectorConfig(
-                options=[
-                    cls.value
-                    for cls in SensorDeviceClass
-                    if cls != SensorDeviceClass.ENUM
-                ],
-                mode=selector.SelectSelectorMode.DROPDOWN,
-                translation_key="device_class",
-                sort=True,
-            )
+        vol.Required(CONF_QUERY): selector.TextSelector(
+            selector.TextSelectorConfig(multiline=True)
         ),
-        vol.Optional(CONF_STATE_CLASS): selector.SelectSelector(
-            selector.SelectSelectorConfig(
-                options=[cls.value for cls in SensorStateClass],
-                mode=selector.SelectSelectorMode.DROPDOWN,
-                translation_key="state_class",
-                sort=True,
-            )
+        vol.Required(CONF_COLUMN_NAME): selector.TextSelector(),
+        vol.Required(CONF_ADVANCED_OPTIONS): section(
+            vol.Schema(
+                {
+                    vol.Optional(CONF_VALUE_TEMPLATE): selector.TemplateSelector(),
+                    vol.Optional(CONF_UNIT_OF_MEASUREMENT): selector.TextSelector(),
+                    vol.Optional(CONF_DEVICE_CLASS): selector.SelectSelector(
+                        selector.SelectSelectorConfig(
+                            options=[
+                                cls.value
+                                for cls in SensorDeviceClass
+                                if cls != SensorDeviceClass.ENUM
+                            ],
+                            mode=selector.SelectSelectorMode.DROPDOWN,
+                            translation_key="device_class",
+                            sort=True,
+                        )
+                    ),
+                    vol.Optional(CONF_STATE_CLASS): selector.SelectSelector(
+                        selector.SelectSelectorConfig(
+                            options=[cls.value for cls in SensorStateClass],
+                            mode=selector.SelectSelectorMode.DROPDOWN,
+                            translation_key="state_class",
+                            sort=True,
+                        )
+                    ),
+                }
+            ),
+            {"collapsed": True},
         ),
     }
 )
 
 CONFIG_SCHEMA: vol.Schema = vol.Schema(
     {
         vol.Required(CONF_NAME, default="Select SQL Query"): selector.TextSelector(),
+        vol.Optional(CONF_DB_URL): selector.TextSelector(),
     }
-).extend(OPTIONS_SCHEMA.schema)
+)
 
 
 def validate_sql_select(value: str) -> str:
@@ -99,6 +99,31 @@ def validate_sql_select(value: str) -> str:
     return str(query[0])
 
 
+def validate_db_connection(db_url: str) -> bool:
+    """Validate db connection."""
+
+    engine: Engine | None = None
+    sess: Session | None = None
+    try:
+        engine = sqlalchemy.create_engine(db_url, future=True)
+        sessmaker = scoped_session(sessionmaker(bind=engine, future=True))
+        sess = sessmaker()
+        sess.execute(sqlalchemy.text("select 1 as value"))
+    except SQLAlchemyError as error:
+        _LOGGER.debug("Execution error %s", error)
+        if sess:
+            sess.close()
+        if engine:
+            engine.dispose()
+        raise
+
+    if sess:
+        sess.close()
+        engine.dispose()
+
+    return True
+
+
 def validate_query(db_url: str, query: str, column: str) -> bool:
     """Validate SQL query."""
 
@@ -136,7 +161,9 @@ def validate_query(db_url: str, query: str, column: str) -> bool:
 class SQLConfigFlow(ConfigFlow, domain=DOMAIN):
     """Handle a config flow for SQL integration."""
 
-    VERSION = 1
+    VERSION = 2
+
+    data: dict[str, Any]
 
     @staticmethod
     @callback
@@ -151,17 +178,46 @@ async def async_step_user(
     ) -> ConfigFlowResult:
         """Handle the user step."""
         errors = {}
-        description_placeholders = {}
 
         if user_input is not None:
             db_url = user_input.get(CONF_DB_URL)
+
+            try:
+                db_url_for_validation = resolve_db_url(self.hass, db_url)
+                await self.hass.async_add_executor_job(
+                    validate_db_connection, db_url_for_validation
+                )
+            except SQLAlchemyError:
+                errors["db_url"] = "db_url_invalid"
+
+            if not errors:
+                self.data = {CONF_NAME: user_input[CONF_NAME]}
+                if db_url and db_url_for_validation != get_instance(self.hass).db_url:
+                    self.data[CONF_DB_URL] = db_url
+                return await self.async_step_options()
+
+        return self.async_show_form(
+            step_id="user",
+            data_schema=self.add_suggested_values_to_schema(CONFIG_SCHEMA, user_input),
+            errors=errors,
+        )
+
+    async def async_step_options(
+        self, user_input: dict[str, Any] | None = None
+    ) -> ConfigFlowResult:
+        """Handle the user step."""
+        errors = {}
+        description_placeholders = {}
+
+        if user_input is not None:
             query = user_input[CONF_QUERY]
             column = user_input[CONF_COLUMN_NAME]
-            db_url_for_validation = None
 
             try:
                 query = validate_sql_select(query)
-                db_url_for_validation = resolve_db_url(self.hass, db_url)
+                db_url_for_validation = resolve_db_url(
+                    self.hass, self.data.get(CONF_DB_URL)
+                )
                 await self.hass.async_add_executor_job(
                     validate_query, db_url_for_validation, query, column
                 )
@@ -178,32 +234,25 @@ async def async_step_user(
                 _LOGGER.debug("Invalid query: %s", err)
                 errors["query"] = "query_invalid"
 
-            options = {
-                CONF_QUERY: query,
-                CONF_COLUMN_NAME: column,
-                CONF_NAME: user_input[CONF_NAME],
+            mod_advanced_options = {
+                k: v
+                for k, v in user_input[CONF_ADVANCED_OPTIONS].items()
+                if v is not None
             }
-            if uom := user_input.get(CONF_UNIT_OF_MEASUREMENT):
-                options[CONF_UNIT_OF_MEASUREMENT] = uom
-            if value_template := user_input.get(CONF_VALUE_TEMPLATE):
-                options[CONF_VALUE_TEMPLATE] = value_template
-            if device_class := user_input.get(CONF_DEVICE_CLASS):
-                options[CONF_DEVICE_CLASS] = device_class
-            if state_class := user_input.get(CONF_STATE_CLASS):
-                options[CONF_STATE_CLASS] = state_class
-            if db_url_for_validation != get_instance(self.hass).db_url:
-                options[CONF_DB_URL] = db_url_for_validation
+            user_input[CONF_ADVANCED_OPTIONS] = mod_advanced_options
 
             if not errors:
+                name = self.data[CONF_NAME]
+                self.data.pop(CONF_NAME)
                 return self.async_create_entry(
-                    title=user_input[CONF_NAME],
-                    data={},
-                    options=options,
+                    title=name,
+                    data=self.data,
+                    options=user_input,
                 )
 
         return self.async_show_form(
-            step_id="user",
-            data_schema=self.add_suggested_values_to_schema(CONFIG_SCHEMA, user_input),
+            step_id="options",
+            data_schema=self.add_suggested_values_to_schema(OPTIONS_SCHEMA, user_input),
             errors=errors,
             description_placeholders=description_placeholders,
         )
@@ -220,10 +269,9 @@ async def async_step_init(
         description_placeholders = {}
 
         if user_input is not None:
-            db_url = user_input.get(CONF_DB_URL)
+            db_url = self.config_entry.data.get(CONF_DB_URL)
             query = user_input[CONF_QUERY]
             column = user_input[CONF_COLUMN_NAME]
-            name = self.config_entry.options.get(CONF_NAME, self.config_entry.title)
 
             try:
                 query = validate_sql_select(query)
@@ -252,24 +300,15 @@ async def async_step_init(
                     recorder_db,
                 )
 
-                options = {
-                    CONF_QUERY: query,
-                    CONF_COLUMN_NAME: column,
-                    CONF_NAME: name,
+                mod_advanced_options = {
+                    k: v
+                    for k, v in user_input[CONF_ADVANCED_OPTIONS].items()
+                    if v is not None
                 }
-                if uom := user_input.get(CONF_UNIT_OF_MEASUREMENT):
-                    options[CONF_UNIT_OF_MEASUREMENT] = uom
-                if value_template := user_input.get(CONF_VALUE_TEMPLATE):
-                    options[CONF_VALUE_TEMPLATE] = value_template
-                if device_class := user_input.get(CONF_DEVICE_CLASS):
-                    options[CONF_DEVICE_CLASS] = device_class
-                if state_class := user_input.get(CONF_STATE_CLASS):
-                    options[CONF_STATE_CLASS] = state_class
-                if db_url_for_validation != get_instance(self.hass).db_url:
-                    options[CONF_DB_URL] = db_url_for_validation
+                user_input[CONF_ADVANCED_OPTIONS] = mod_advanced_options
 
                 return self.async_create_entry(
-                    data=options,
+                    data=user_input,
                 )
 
         return self.async_show_form(

homeassistant/components/sql/const.py

@@ -9,4 +9,5 @@
 
 CONF_COLUMN_NAME = "column"
 CONF_QUERY = "query"
+CONF_ADVANCED_OPTIONS = "advanced_options"
 DB_URL_RE = re.compile("//.*:.*@")