Refactor SQL with ManualTriggerEntity (#95116)

* First go

* Finalize sensor

* Add tests

* Remove not need _attr_name

* device_class

* _process_manual_data allow Any as value

homeassistant/components/sql/__init__.py

@@ -14,6 +14,7 @@ from homeassistant.components.sensor import (
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import (
     CONF_DEVICE_CLASS,
+    CONF_ICON,
     CONF_NAME,
     CONF_UNIQUE_ID,
     CONF_UNIT_OF_MEASUREMENT,
@@ -23,6 +24,7 @@ from homeassistant.const import (
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers import discovery
 import homeassistant.helpers.config_validation as cv
+from homeassistant.helpers.template_entity import CONF_AVAILABILITY, CONF_PICTURE
 from homeassistant.helpers.typing import ConfigType
 
 from .const import CONF_COLUMN_NAME, CONF_QUERY, DOMAIN, PLATFORMS
@@ -41,7 +43,7 @@ def validate_sql_select(value: str) -> str:
 QUERY_SCHEMA = vol.Schema(
     {
         vol.Required(CONF_COLUMN_NAME): cv.string,
-        vol.Required(CONF_NAME): cv.string,
+        vol.Required(CONF_NAME): cv.template,
         vol.Required(CONF_QUERY): vol.All(cv.string, validate_sql_select),
         vol.Optional(CONF_UNIT_OF_MEASUREMENT): cv.string,
         vol.Optional(CONF_VALUE_TEMPLATE): cv.template,
@@ -49,6 +51,9 @@ QUERY_SCHEMA = vol.Schema(
         vol.Optional(CONF_DB_URL): cv.string,
         vol.Optional(CONF_DEVICE_CLASS): DEVICE_CLASSES_SCHEMA,
         vol.Optional(CONF_STATE_CLASS): STATE_CLASSES_SCHEMA,
+        vol.Optional(CONF_AVAILABILITY): cv.template,
+        vol.Optional(CONF_ICON): cv.template,
+        vol.Optional(CONF_PICTURE): cv.template,
     }
 )
 

homeassistant/components/sql/sensor.py

@@ -4,6 +4,7 @@ from __future__ import annotations
 from datetime import date
 import decimal
 import logging
+from typing import Any
 
 import sqlalchemy
 from sqlalchemy import lambda_stmt
@@ -27,6 +28,7 @@ from homeassistant.components.sensor import (
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import (
     CONF_DEVICE_CLASS,
+    CONF_ICON,
     CONF_NAME,
     CONF_UNIQUE_ID,
     CONF_UNIT_OF_MEASUREMENT,
@@ -40,6 +42,11 @@ 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.template_entity import (
+    CONF_AVAILABILITY,
+    CONF_PICTURE,
+    ManualTriggerEntity,
+)
 from homeassistant.helpers.typing import ConfigType, DiscoveryInfoType
 
 from .const import CONF_COLUMN_NAME, CONF_QUERY, DOMAIN
@@ -61,7 +68,7 @@ async def async_setup_platform(
     if (conf := discovery_info) is None:
         return
 
-    name: str = conf[CONF_NAME]
+    name: Template = conf[CONF_NAME]
     query_str: str = conf[CONF_QUERY]
     unit: str | None = conf.get(CONF_UNIT_OF_MEASUREMENT)
     value_template: Template | None = conf.get(CONF_VALUE_TEMPLATE)
@@ -70,13 +77,24 @@ async def async_setup_platform(
     db_url: str = resolve_db_url(hass, conf.get(CONF_DB_URL))
     device_class: SensorDeviceClass | None = conf.get(CONF_DEVICE_CLASS)
     state_class: SensorStateClass | None = conf.get(CONF_STATE_CLASS)
+    availability: Template | None = conf.get(CONF_AVAILABILITY)
+    icon: Template | None = conf.get(CONF_ICON)
+    picture: Template | None = conf.get(CONF_PICTURE)
 
     if value_template is not None:
         value_template.hass = hass
 
+    trigger_entity_config = {CONF_NAME: name, CONF_DEVICE_CLASS: device_class}
+    if availability:
+        trigger_entity_config[CONF_AVAILABILITY] = availability
+    if icon:
+        trigger_entity_config[CONF_ICON] = icon
+    if picture:
+        trigger_entity_config[CONF_PICTURE] = picture
+
     await async_setup_sensor(
         hass,
-        name,
+        trigger_entity_config,
         query_str,
         column_name,
         unit,
@@ -84,7 +102,6 @@ async def async_setup_platform(
         unique_id,
         db_url,
         True,
-        device_class,
         state_class,
         async_add_entities,
     )
@@ -114,9 +131,12 @@ async def async_setup_entry(
         if value_template is not None:
             value_template.hass = hass
 
+    name_template = Template(name, hass)
+    trigger_entity_config = {CONF_NAME: name_template, CONF_DEVICE_CLASS: device_class}
+
     await async_setup_sensor(
         hass,
-        name,
+        trigger_entity_config,
         query_str,
         column_name,
         unit,
@@ -124,7 +144,6 @@ async def async_setup_entry(
         entry.entry_id,
         db_url,
         False,
-        device_class,
         state_class,
         async_add_entities,
     )
@@ -162,7 +181,7 @@ def _async_get_or_init_domain_data(hass: HomeAssistant) -> SQLData:
 
 async def async_setup_sensor(
     hass: HomeAssistant,
-    name: str,
+    trigger_entity_config: ConfigType,
     query_str: str,
     column_name: str,
     unit: str | None,
@@ -170,7 +189,6 @@ async def async_setup_sensor(
     unique_id: str | None,
     db_url: str,
     yaml: bool,
-    device_class: SensorDeviceClass | None,
     state_class: SensorStateClass | None,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
@@ -245,7 +263,7 @@ async def async_setup_sensor(
     async_add_entities(
         [
             SQLSensor(
-                name,
+                trigger_entity_config,
                 sessmaker,
                 query_str,
                 column_name,
@@ -253,12 +271,10 @@ async def async_setup_sensor(
                 value_template,
                 unique_id,
                 yaml,
-                device_class,
                 state_class,
                 use_database_executor,
             )
         ],
-        True,
     )
 
 
@@ -295,15 +311,12 @@ def _generate_lambda_stmt(query: str) -> StatementLambdaElement:
     return lambda_stmt(lambda: text, lambda_cache=_SQL_LAMBDA_CACHE)
 
 
-class SQLSensor(SensorEntity):
+class SQLSensor(ManualTriggerEntity, SensorEntity):
     """Representation of an SQL sensor."""
 
-    _attr_icon = "mdi:database-search"
-    _attr_has_entity_name = True
-
     def __init__(
         self,
-        name: str,
+        trigger_entity_config: ConfigType,
         sessmaker: scoped_session,
         query: str,
         column: str,
@@ -311,15 +324,13 @@ class SQLSensor(SensorEntity):
         value_template: Template | None,
         unique_id: str | None,
         yaml: bool,
-        device_class: SensorDeviceClass | None,
         state_class: SensorStateClass | None,
         use_database_executor: bool,
     ) -> None:
         """Initialize the SQL sensor."""
+        super().__init__(self.hass, trigger_entity_config)
         self._query = query
-        self._attr_name = name if yaml else None
         self._attr_native_unit_of_measurement = unit
-        self._attr_device_class = device_class
         self._attr_state_class = state_class
         self._template = value_template
         self._column_name = column
@@ -328,22 +339,34 @@ class SQLSensor(SensorEntity):
         self._attr_unique_id = unique_id
         self._use_database_executor = use_database_executor
         self._lambda_stmt = _generate_lambda_stmt(query)
+        self._attr_has_entity_name = not yaml
         if not yaml and unique_id:
             self._attr_device_info = DeviceInfo(
                 entry_type=DeviceEntryType.SERVICE,
                 identifiers={(DOMAIN, unique_id)},
                 manufacturer="SQL",
-                name=name,
+                name=trigger_entity_config[CONF_NAME].template,
             )
 
+    async def async_added_to_hass(self) -> None:
+        """Call when entity about to be added to hass."""
+        await super().async_added_to_hass()
+        await self.async_update()
+
+    @property
+    def extra_state_attributes(self) -> dict[str, Any] | None:
+        """Return extra attributes."""
+        return dict(self._attr_extra_state_attributes)
+
     async def async_update(self) -> None:
         """Retrieve sensor data from the query using the right executor."""
         if self._use_database_executor:
-            await get_instance(self.hass).async_add_executor_job(self._update)
+            data = await get_instance(self.hass).async_add_executor_job(self._update)
         else:
-            await self.hass.async_add_executor_job(self._update)
+            data = await self.hass.async_add_executor_job(self._update)
+        self._process_manual_data(data)
 
-    def _update(self) -> None:
+    def _update(self) -> Any:
         """Retrieve sensor data from the query."""
         data = None
         self._attr_extra_state_attributes = {}
@@ -384,3 +407,4 @@ class SQLSensor(SensorEntity):
             _LOGGER.warning("%s returned no results", self._query)
 
         sess.close()
+        return data

homeassistant/helpers/template_entity.py

@@ -624,7 +624,7 @@ class ManualTriggerEntity(TriggerBaseEntity):
         TriggerBaseEntity.__init__(self, hass, config)
 
     @callback
-    def _process_manual_data(self, value: str | None = None) -> None:
+    def _process_manual_data(self, value: Any | None = None) -> None:
         """Process new data manually.
 
         Implementing class should call this last in update method to render templates.

tests/components/sql/__init__.py

@@ -13,12 +13,14 @@ from homeassistant.components.sql.const import CONF_COLUMN_NAME, CONF_QUERY, DOM
 from homeassistant.config_entries import SOURCE_USER
 from homeassistant.const import (
     CONF_DEVICE_CLASS,
+    CONF_ICON,
     CONF_NAME,
     CONF_UNIQUE_ID,
     CONF_UNIT_OF_MEASUREMENT,
     CONF_VALUE_TEMPLATE,
 )
 from homeassistant.core import HomeAssistant
+from homeassistant.helpers.template_entity import CONF_AVAILABILITY, CONF_PICTURE
 
 from tests.common import MockConfigEntry
 
@@ -148,6 +150,23 @@ YAML_CONFIG_NO_DB = {
     }
 }
 
+YAML_CONFIG_ALL_TEMPLATES = {
+    "sql": {
+        CONF_DB_URL: "sqlite://",
+        CONF_NAME: "Get values with template",
+        CONF_QUERY: "SELECT 5 as output",
+        CONF_COLUMN_NAME: "output",
+        CONF_UNIT_OF_MEASUREMENT: "MiB/s",
+        CONF_UNIQUE_ID: "unique_id_123456",
+        CONF_VALUE_TEMPLATE: "{{ value }}",
+        CONF_ICON: '{% if states("sensor.input1")=="on" %} mdi:on {% else %} mdi:off {% endif %}',
+        CONF_PICTURE: '{% if states("sensor.input1")=="on" %} /local/picture1.jpg {% else %} /local/picture2.jpg {% endif %}',
+        CONF_AVAILABILITY: '{{ states("sensor.input2")=="on" }}',
+        CONF_DEVICE_CLASS: SensorDeviceClass.DATA_RATE,
+        CONF_STATE_CLASS: SensorStateClass.MEASUREMENT,
+    }
+}
+
 
 async def init_integration(
     hass: HomeAssistant,

tests/components/sql/test_sensor.py

@@ -13,7 +13,12 @@ from homeassistant.components.recorder import Recorder
 from homeassistant.components.sensor import SensorDeviceClass, SensorStateClass
 from homeassistant.components.sql.const import CONF_QUERY, DOMAIN
 from homeassistant.config_entries import SOURCE_USER
-from homeassistant.const import CONF_UNIQUE_ID, STATE_UNKNOWN
+from homeassistant.const import (
+    CONF_ICON,
+    CONF_UNIQUE_ID,
+    STATE_UNAVAILABLE,
+    STATE_UNKNOWN,
+)
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers import issue_registry as ir
 from homeassistant.setup import async_setup_component
@@ -21,6 +26,7 @@ from homeassistant.util import dt as dt_util
 
 from . import (
     YAML_CONFIG,
+    YAML_CONFIG_ALL_TEMPLATES,
     YAML_CONFIG_BINARY,
     YAML_CONFIG_FULL_TABLE_SCAN,
     YAML_CONFIG_FULL_TABLE_SCAN_NO_UNIQUE_ID,
@@ -32,13 +38,14 @@ from . import (
 from tests.common import MockConfigEntry, async_fire_time_changed
 
 
-async def test_query(recorder_mock: Recorder, hass: HomeAssistant) -> None:
+async def test_query_basic(recorder_mock: Recorder, hass: HomeAssistant) -> None:
     """Test the SQL sensor."""
     config = {
         "db_url": "sqlite://",
         "query": "SELECT 5 as value",
         "column": "value",
         "name": "Select value SQL query",
+        "unique_id": "very_unique_id",
     }
     await init_integration(hass, config)
 
@@ -235,6 +242,65 @@ async def test_query_from_yaml(recorder_mock: Recorder, hass: HomeAssistant) ->
     assert state.state == "5"
 
 
+async def test_templates_with_yaml(
+    recorder_mock: Recorder, hass: HomeAssistant
+) -> None:
+    """Test the SQL sensor from yaml config with templates."""
+
+    hass.states.async_set("sensor.input1", "on")
+    hass.states.async_set("sensor.input2", "on")
+    await hass.async_block_till_done()
+
+    assert await async_setup_component(hass, DOMAIN, YAML_CONFIG_ALL_TEMPLATES)
+    await hass.async_block_till_done()
+
+    state = hass.states.get("sensor.get_values_with_template")
+    assert state.state == "5"
+    assert state.attributes[CONF_ICON] == "mdi:on"
+    assert state.attributes["entity_picture"] == "/local/picture1.jpg"
+
+    hass.states.async_set("sensor.input1", "off")
+    await hass.async_block_till_done()
+
+    async_fire_time_changed(
+        hass,
+        dt_util.utcnow() + timedelta(minutes=1),
+    )
+    await hass.async_block_till_done()
+
+    state = hass.states.get("sensor.get_values_with_template")
+    assert state.state == "5"
+    assert state.attributes[CONF_ICON] == "mdi:off"
+    assert state.attributes["entity_picture"] == "/local/picture2.jpg"
+
+    hass.states.async_set("sensor.input2", "off")
+    await hass.async_block_till_done()
+
+    async_fire_time_changed(
+        hass,
+        dt_util.utcnow() + timedelta(minutes=2),
+    )
+    await hass.async_block_till_done()
+
+    state = hass.states.get("sensor.get_values_with_template")
+    assert state.state == STATE_UNAVAILABLE
+
+    hass.states.async_set("sensor.input1", "on")
+    hass.states.async_set("sensor.input2", "on")
+    await hass.async_block_till_done()
+
+    async_fire_time_changed(
+        hass,
+        dt_util.utcnow() + timedelta(minutes=3),
+    )
+    await hass.async_block_till_done()
+
+    state = hass.states.get("sensor.get_values_with_template")
+    assert state.state == "5"
+    assert state.attributes[CONF_ICON] == "mdi:on"
+    assert state.attributes["entity_picture"] == "/local/picture1.jpg"
+
+
 async def test_config_from_old_yaml(
     recorder_mock: Recorder, hass: HomeAssistant
 ) -> None: