Implement config flow for SQL integration (#68700)

homeassistant/components/sql/sensor.py

@@ -4,9 +4,9 @@ from __future__ import annotations
 from datetime import date
 import decimal
 import logging
-import re
 
 import sqlalchemy
+from sqlalchemy.exc import SQLAlchemyError
 from sqlalchemy.orm import scoped_session, sessionmaker
 import voluptuous as vol
 
@@ -15,39 +15,32 @@ from homeassistant.components.sensor import (
     PLATFORM_SCHEMA as PARENT_PLATFORM_SCHEMA,
     SensorEntity,
 )
+from homeassistant.config_entries import SOURCE_IMPORT, ConfigEntry
 from homeassistant.const import CONF_NAME, CONF_UNIT_OF_MEASUREMENT, CONF_VALUE_TEMPLATE
 from homeassistant.core import HomeAssistant
+from homeassistant.exceptions import TemplateError
 import homeassistant.helpers.config_validation as cv
+from homeassistant.helpers.device_registry import DeviceEntryType
+from homeassistant.helpers.entity import DeviceInfo
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.template import Template
 from homeassistant.helpers.typing import ConfigType, DiscoveryInfoType
 
+from .const import CONF_COLUMN_NAME, CONF_QUERIES, CONF_QUERY, DB_URL_RE, DOMAIN
+
 _LOGGER = logging.getLogger(__name__)
 
-CONF_COLUMN_NAME = "column"
-CONF_QUERIES = "queries"
-CONF_QUERY = "query"
-
-DB_URL_RE = re.compile("//.*:.*@")
-
 
 def redact_credentials(data: str) -> str:
     """Redact credentials from string data."""
     return DB_URL_RE.sub("//****:****@", data)
 
 
-def validate_sql_select(value: str) -> str:
-    """Validate that value is a SQL SELECT query."""
-    if not value.lstrip().lower().startswith("select"):
-        raise vol.Invalid("Only SELECT queries allowed")
-    return value
-
-
 _QUERY_SCHEME = vol.Schema(
     {
         vol.Required(CONF_COLUMN_NAME): cv.string,
         vol.Required(CONF_NAME): cv.string,
-        vol.Required(CONF_QUERY): vol.All(cv.string, validate_sql_select),
+        vol.Required(CONF_QUERY): cv.string,
         vol.Optional(CONF_UNIT_OF_MEASUREMENT): cv.string,
         vol.Optional(CONF_VALUE_TEMPLATE): cv.template,
     }
@@ -58,67 +51,101 @@ PLATFORM_SCHEMA = PARENT_PLATFORM_SCHEMA.extend(
 )
 
 
-def setup_platform(
+async def async_setup_platform(
     hass: HomeAssistant,
     config: ConfigType,
-    add_entities: AddEntitiesCallback,
+    async_add_entities: AddEntitiesCallback,
     discovery_info: DiscoveryInfoType | None = None,
 ) -> None:
     """Set up the SQL sensor platform."""
-    if not (db_url := config.get(CONF_DB_URL)):
-        db_url = DEFAULT_URL.format(hass_config_path=hass.config.path(DEFAULT_DB_FILE))
+    _LOGGER.warning(
+        # SQL config flow added in 2022.4 and should be removed in 2022.6
+        "Configuration of the SQL sensor platform in YAML is deprecated and "
+        "will be removed in Home Assistant 2022.6; Your existing configuration "
+        "has been imported into the UI automatically and can be safely removed "
+        "from your configuration.yaml file"
+    )
 
-    sess: scoped_session | None = None
-    try:
-        engine = sqlalchemy.create_engine(db_url)
-        sessmaker = scoped_session(sessionmaker(bind=engine))
-
-        # Run a dummy query just to test the db_url
-        sess = sessmaker()
-        sess.execute("SELECT 1;")
-
-    except sqlalchemy.exc.SQLAlchemyError as err:
-        _LOGGER.error(
-            "Couldn't connect using %s DB_URL: %s",
-            redact_credentials(db_url),
-            redact_credentials(str(err)),
-        )
-        return
-    finally:
-        if sess:
-            sess.close()
-
-    queries = []
+    default_db_url = DEFAULT_URL.format(
+        hass_config_path=hass.config.path(DEFAULT_DB_FILE)
+    )
 
     for query in config[CONF_QUERIES]:
-        name: str = query[CONF_NAME]
-        query_str: str = query[CONF_QUERY]
-        unit: str | None = query.get(CONF_UNIT_OF_MEASUREMENT)
-        value_template: Template | None = query.get(CONF_VALUE_TEMPLATE)
-        column_name: str = query[CONF_COLUMN_NAME]
+        new_config = {
+            CONF_DB_URL: config.get(CONF_DB_URL, default_db_url),
+            CONF_NAME: query.get(CONF_NAME),
+            CONF_QUERY: query.get(CONF_QUERY),
+            CONF_UNIT_OF_MEASUREMENT: query.get(CONF_UNIT_OF_MEASUREMENT),
+            CONF_VALUE_TEMPLATE: query.get(CONF_VALUE_TEMPLATE),
+            CONF_COLUMN_NAME: query.get(CONF_COLUMN_NAME),
+        }
+        hass.async_create_task(
+            hass.config_entries.flow.async_init(
+                DOMAIN,
+                context={"source": SOURCE_IMPORT},
+                data=new_config,
+            )
+        )
 
+
+async def async_setup_entry(
+    hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
+) -> None:
+    """Set up the SQL sensor entry."""
+
+    db_url: str = entry.options[CONF_DB_URL]
+    name: str = entry.options[CONF_NAME]
+    query_str: str = entry.options[CONF_QUERY]
+    unit: str | None = entry.options.get(CONF_UNIT_OF_MEASUREMENT)
+    template: str | None = entry.options.get(CONF_VALUE_TEMPLATE)
+    column_name: str = entry.options[CONF_COLUMN_NAME]
+
+    value_template: Template | None = None
+    if template is not None:
+        try:
+            value_template = Template(template)
+            value_template.ensure_valid()
+        except TemplateError:
+            value_template = None
         if value_template is not None:
             value_template.hass = hass
 
-        # MSSQL uses TOP and not LIMIT
-        if not ("LIMIT" in query_str.upper() or "SELECT TOP" in query_str.upper()):
-            query_str = (
-                query_str.replace("SELECT", "SELECT TOP 1")
-                if "mssql" in db_url
-                else query_str.replace(";", " LIMIT 1;")
-            )
+    try:
+        engine = sqlalchemy.create_engine(db_url)
+        sessmaker = scoped_session(sessionmaker(bind=engine))
+    except SQLAlchemyError as err:
+        _LOGGER.error("Can not open database %s", {redact_credentials(str(err))})
+        return
 
-        sensor = SQLSensor(
-            name, sessmaker, query_str, column_name, unit, value_template
+    # MSSQL uses TOP and not LIMIT
+    if not ("LIMIT" in query_str.upper() or "SELECT TOP" in query_str.upper()):
+        query_str = (
+            query_str.replace("SELECT", "SELECT TOP 1")
+            if "mssql" in db_url
+            else query_str.replace(";", " LIMIT 1;")
         )
-        queries.append(sensor)
 
-    add_entities(queries, True)
+    async_add_entities(
+        [
+            SQLSensor(
+                name,
+                sessmaker,
+                query_str,
+                column_name,
+                unit,
+                value_template,
+                entry.entry_id,
+            )
+        ],
+        True,
+    )
 
 
 class SQLSensor(SensorEntity):
     """Representation of an SQL sensor."""
 
+    _attr_icon = "mdi:database-search"
+
     def __init__(
         self,
         name: str,
@@ -127,6 +154,7 @@ class SQLSensor(SensorEntity):
         column: str,
         unit: str | None,
         value_template: Template | None,
+        entry_id: str,
     ) -> None:
         """Initialize the SQL sensor."""
         self._attr_name = name
@@ -136,6 +164,13 @@ class SQLSensor(SensorEntity):
         self._column_name = column
         self.sessionmaker = sessmaker
         self._attr_extra_state_attributes = {}
+        self._attr_unique_id = entry_id
+        self._attr_device_info = DeviceInfo(
+            entry_type=DeviceEntryType.SERVICE,
+            identifiers={(DOMAIN, entry_id)},
+            manufacturer="SQL",
+            name=name,
+        )
 
     def update(self) -> None:
         """Retrieve sensor data from the query."""
@@ -145,7 +180,7 @@ class SQLSensor(SensorEntity):
         sess: scoped_session = self.sessionmaker()
         try:
             result = sess.execute(self._query)
-        except sqlalchemy.exc.SQLAlchemyError as err:
+        except SQLAlchemyError as err:
             _LOGGER.error(
                 "Error executing query %s: %s",
                 self._query,