Add device_class and state_class to sql (#85418)

homeassistant/components/sql/__init__.py

@@ -4,8 +4,14 @@ from __future__ import annotations
 import voluptuous as vol
 
 from homeassistant.components.recorder import CONF_DB_URL
+from homeassistant.components.sensor import (
+    CONF_STATE_CLASS,
+    DEVICE_CLASSES_SCHEMA,
+    STATE_CLASSES_SCHEMA,
+)
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import (
+    CONF_DEVICE_CLASS,
     CONF_NAME,
     CONF_UNIQUE_ID,
     CONF_UNIT_OF_MEASUREMENT,
@@ -36,6 +42,8 @@ QUERY_SCHEMA = vol.Schema(
         vol.Optional(CONF_VALUE_TEMPLATE): cv.template,
         vol.Optional(CONF_UNIQUE_ID): cv.string,
         vol.Optional(CONF_DB_URL): cv.string,
+        vol.Optional(CONF_DEVICE_CLASS): DEVICE_CLASSES_SCHEMA,
+        vol.Optional(CONF_STATE_CLASS): STATE_CLASSES_SCHEMA,
     }
 )
 

homeassistant/components/sql/sensor.py

@@ -11,9 +11,15 @@ from sqlalchemy.exc import SQLAlchemyError
 from sqlalchemy.orm import Session, scoped_session, sessionmaker
 
 from homeassistant.components.recorder import CONF_DB_URL, DEFAULT_DB_FILE, DEFAULT_URL
-from homeassistant.components.sensor import SensorEntity
+from homeassistant.components.sensor import (
+    CONF_STATE_CLASS,
+    SensorDeviceClass,
+    SensorEntity,
+    SensorStateClass,
+)
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import (
+    CONF_DEVICE_CLASS,
     CONF_NAME,
     CONF_UNIQUE_ID,
     CONF_UNIT_OF_MEASUREMENT,
@@ -54,6 +60,8 @@ async def async_setup_platform(
     column_name: str = conf[CONF_COLUMN_NAME]
     unique_id: str | None = conf.get(CONF_UNIQUE_ID)
     db_url: str | None = conf.get(CONF_DB_URL)
+    device_class: SensorDeviceClass | None = conf.get(CONF_DEVICE_CLASS)
+    state_class: SensorStateClass | None = conf.get(CONF_STATE_CLASS)
 
     if value_template is not None:
         value_template.hass = hass
@@ -68,6 +76,8 @@ async def async_setup_platform(
         unique_id,
         db_url,
         True,
+        device_class,
+        state_class,
         async_add_entities,
     )
 
@@ -104,6 +114,8 @@ async def async_setup_entry(
         entry.entry_id,
         db_url,
         False,
+        None,
+        None,
         async_add_entities,
     )
 
@@ -118,6 +130,8 @@ async def async_setup_sensor(
     unique_id: str | None,
     db_url: str | None,
     yaml: bool,
+    device_class: SensorDeviceClass | None,
+    state_class: SensorStateClass | None,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up the SQL sensor."""
@@ -163,6 +177,8 @@ async def async_setup_sensor(
                 value_template,
                 unique_id,
                 yaml,
+                device_class,
+                state_class,
             )
         ],
         True,
@@ -185,11 +201,15 @@ class SQLSensor(SensorEntity):
         value_template: Template | None,
         unique_id: str | None,
         yaml: bool,
+        device_class: SensorDeviceClass | None,
+        state_class: SensorStateClass | None,
     ) -> None:
         """Initialize the SQL sensor."""
         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
         self.sessionmaker = sessmaker

tests/components/sql/__init__.py

@@ -4,9 +4,15 @@ from __future__ import annotations
 from typing import Any
 
 from homeassistant.components.recorder import CONF_DB_URL
+from homeassistant.components.sensor import (
+    CONF_STATE_CLASS,
+    SensorDeviceClass,
+    SensorStateClass,
+)
 from homeassistant.components.sql.const import CONF_COLUMN_NAME, CONF_QUERY, DOMAIN
 from homeassistant.config_entries import SOURCE_USER
 from homeassistant.const import (
+    CONF_DEVICE_CLASS,
     CONF_NAME,
     CONF_UNIQUE_ID,
     CONF_UNIT_OF_MEASUREMENT,
@@ -56,6 +62,8 @@ YAML_CONFIG = {
         CONF_UNIT_OF_MEASUREMENT: "MiB",
         CONF_UNIQUE_ID: "unique_id_12345",
         CONF_VALUE_TEMPLATE: "{{ value }}",
+        CONF_DEVICE_CLASS: SensorDeviceClass.DATA_RATE,
+        CONF_STATE_CLASS: SensorStateClass.MEASUREMENT,
     }
 }
 

tests/components/sql/test_sensor.py

@@ -9,6 +9,7 @@ from sqlalchemy import text as sql_text
 from sqlalchemy.exc import SQLAlchemyError
 
 from homeassistant.components.recorder import Recorder
+from homeassistant.components.sensor import SensorDeviceClass, SensorStateClass
 from homeassistant.components.sql.const import DOMAIN
 from homeassistant.config_entries import SOURCE_USER
 from homeassistant.const import STATE_UNKNOWN
@@ -300,3 +301,19 @@ async def test_invalid_url_setup_from_yaml(
         assert pattern not in caplog.text
     for pattern in expected_patterns:
         assert pattern in caplog.text
+
+
+async def test_attributes_from_yaml_setup(
+    recorder_mock: Recorder, hass: HomeAssistant
+) -> None:
+    """Test attributes from yaml config."""
+
+    assert await async_setup_component(hass, DOMAIN, YAML_CONFIG)
+    await hass.async_block_till_done()
+
+    state = hass.states.get("sensor.get_value")
+
+    assert state.state == "5"
+    assert state.attributes["device_class"] == SensorDeviceClass.DATA_RATE
+    assert state.attributes["state_class"] == SensorStateClass.MEASUREMENT
+    assert state.attributes["unit_of_measurement"] == "MiB"