Add device_class and state_class in config flow for SQL (#95020)

* Add device_class and state_class in config flow for SQL

* Update when selected NONE_SENTINEL

* Add tests

* Use SensorDeviceClass and SensorStateClass in tests

* Add volatile_organic_compounds_parts in strings selector

* Add test_attributes_from_entry_config

* Remove test_attributes_from_entry_config and complement test_device_state_class

* Add test_attributes_from_entry_config in test_sensor.py

homeassistant/components/sql/config_flow.py

@@ -12,7 +12,17 @@ import voluptuous as vol
 
 from homeassistant import config_entries
 from homeassistant.components.recorder import CONF_DB_URL, get_instance
-from homeassistant.const import CONF_NAME, CONF_UNIT_OF_MEASUREMENT, CONF_VALUE_TEMPLATE
+from homeassistant.components.sensor import (
+    CONF_STATE_CLASS,
+    SensorDeviceClass,
+    SensorStateClass,
+)
+from homeassistant.const import (
+    CONF_DEVICE_CLASS,
+    CONF_NAME,
+    CONF_UNIT_OF_MEASUREMENT,
+    CONF_VALUE_TEMPLATE,
+)
 from homeassistant.core import callback
 from homeassistant.data_entry_flow import FlowResult
 from homeassistant.helpers import selector
@@ -22,6 +32,8 @@ from .util import resolve_db_url
 
 _LOGGER = logging.getLogger(__name__)
 
+NONE_SENTINEL = "none"
+
 OPTIONS_SCHEMA: vol.Schema = vol.Schema(
     {
         vol.Optional(
@@ -39,6 +51,34 @@ OPTIONS_SCHEMA: vol.Schema = vol.Schema(
         vol.Optional(
             CONF_VALUE_TEMPLATE,
         ): selector.TemplateSelector(),
+        vol.Optional(
+            CONF_DEVICE_CLASS,
+            default=NONE_SENTINEL,
+        ): selector.SelectSelector(
+            selector.SelectSelectorConfig(
+                options=[NONE_SENTINEL]
+                + sorted(
+                    [
+                        cls.value
+                        for cls in SensorDeviceClass
+                        if cls != SensorDeviceClass.ENUM
+                    ]
+                ),
+                mode=selector.SelectSelectorMode.DROPDOWN,
+                translation_key="device_class",
+            )
+        ),
+        vol.Optional(
+            CONF_STATE_CLASS,
+            default=NONE_SENTINEL,
+        ): selector.SelectSelector(
+            selector.SelectSelectorConfig(
+                options=[NONE_SENTINEL]
+                + sorted([cls.value for cls in SensorStateClass]),
+                mode=selector.SelectSelectorMode.DROPDOWN,
+                translation_key="state_class",
+            )
+        ),
     }
 )
 
@@ -139,6 +179,10 @@ class SQLConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
                 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[CONF_DEVICE_CLASS]) != NONE_SENTINEL:
+                options[CONF_DEVICE_CLASS] = device_class
+            if (state_class := user_input[CONF_STATE_CLASS]) != NONE_SENTINEL:
+                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
 
@@ -204,6 +248,10 @@ class SQLOptionsFlowHandler(config_entries.OptionsFlowWithConfigEntry):
                     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[CONF_DEVICE_CLASS]) != NONE_SENTINEL:
+                    options[CONF_DEVICE_CLASS] = device_class
+                if (state_class := user_input[CONF_STATE_CLASS]) != NONE_SENTINEL:
+                    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
 

homeassistant/components/sql/sensor.py

@@ -101,6 +101,8 @@ async def async_setup_entry(
     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]
+    device_class: SensorDeviceClass | None = entry.options.get(CONF_DEVICE_CLASS, None)
+    state_class: SensorStateClass | None = entry.options.get(CONF_STATE_CLASS, None)
 
     value_template: Template | None = None
     if template is not None:
@@ -122,8 +124,8 @@ async def async_setup_entry(
         entry.entry_id,
         db_url,
         False,
-        None,
+        device_class,
-        None,
+        state_class,
         async_add_entities,
     )
 

homeassistant/components/sql/strings.json

@@ -16,7 +16,9 @@
           "query": "Select Query",
           "column": "Column",
           "unit_of_measurement": "Unit of Measure",
-          "value_template": "Value Template"
+          "value_template": "Value Template",
+          "device_class": "Device Class",
+          "state_class": "State Class"
         },
         "data_description": {
           "db_url": "Database URL, leave empty to use HA recorder database",
@@ -24,7 +26,9 @@
           "query": "Query to run, needs to start with 'SELECT'",
           "column": "Column for returned query to present as state",
           "unit_of_measurement": "Unit of Measure (optional)",
-          "value_template": "Value Template (optional)"
+          "value_template": "Value Template (optional)",
+          "device_class": "The type/class of the sensor to set the icon in the frontend",
+          "state_class": "The state_class of the sensor"
         }
       }
     }
@@ -38,7 +42,9 @@
           "query": "[%key:component::sql::config::step::user::data::query%]",
           "column": "[%key:component::sql::config::step::user::data::column%]",
           "unit_of_measurement": "[%key:component::sql::config::step::user::data::unit_of_measurement%]",
-          "value_template": "[%key:component::sql::config::step::user::data::value_template%]"
+          "value_template": "[%key:component::sql::config::step::user::data::value_template%]",
+          "device_class": "[%key:component::sql::config::step::user::data::device_class%]",
+          "state_class": "[%key:component::sql::config::step::user::data::state_class%]"
         },
         "data_description": {
           "db_url": "[%key:component::sql::config::step::user::data_description::db_url%]",
@@ -46,7 +52,9 @@
           "query": "[%key:component::sql::config::step::user::data_description::query%]",
           "column": "[%key:component::sql::config::step::user::data_description::column%]",
           "unit_of_measurement": "[%key:component::sql::config::step::user::data_description::unit_of_measurement%]",
-          "value_template": "[%key:component::sql::config::step::user::data_description::value_template%]"
+          "value_template": "[%key:component::sql::config::step::user::data_description::value_template%]",
+          "device_class": "[%key:component::sql::config::step::user::data_description::device_class%]",
+          "state_class": "[%key:component::sql::config::step::user::data_description::state_class%]"
         }
       }
     },
@@ -56,6 +64,69 @@
       "column_invalid": "[%key:component::sql::config::error::column_invalid%]"
     }
   },
+  "selector": {
+    "device_class": {
+      "options": {
+        "none": "No device class",
+        "date": "[%key:component::sensor::entity_component::date::name%]",
+        "duration": "[%key:component::sensor::entity_component::duration::name%]",
+        "apparent_power": "[%key:component::sensor::entity_component::apparent_power::name%]",
+        "aqi": "[%key:component::sensor::entity_component::aqi::name%]",
+        "atmospheric_pressure": "[%key:component::sensor::entity_component::atmospheric_pressure::name%]",
+        "battery": "[%key:component::sensor::entity_component::battery::name%]",
+        "carbon_monoxide": "[%key:component::sensor::entity_component::carbon_monoxide::name%]",
+        "carbon_dioxide": "[%key:component::sensor::entity_component::carbon_dioxide::name%]",
+        "current": "[%key:component::sensor::entity_component::current::name%]",
+        "data_rate": "[%key:component::sensor::entity_component::data_rate::name%]",
+        "data_size": "[%key:component::sensor::entity_component::data_size::name%]",
+        "distance": "[%key:component::sensor::entity_component::distance::name%]",
+        "energy": "[%key:component::sensor::entity_component::energy::name%]",
+        "energy_storage": "[%key:component::sensor::entity_component::energy_storage::name%]",
+        "frequency": "[%key:component::sensor::entity_component::frequency::name%]",
+        "gas": "[%key:component::sensor::entity_component::gas::name%]",
+        "humidity": "[%key:component::sensor::entity_component::humidity::name%]",
+        "illuminance": "[%key:component::sensor::entity_component::illuminance::name%]",
+        "irradiance": "[%key:component::sensor::entity_component::irradiance::name%]",
+        "moisture": "[%key:component::sensor::entity_component::moisture::name%]",
+        "monetary": "[%key:component::sensor::entity_component::monetary::name%]",
+        "nitrogen_dioxide": "[%key:component::sensor::entity_component::nitrogen_dioxide::name%]",
+        "nitrogen_monoxide": "[%key:component::sensor::entity_component::nitrogen_monoxide::name%]",
+        "nitrous_oxide": "[%key:component::sensor::entity_component::nitrous_oxide::name%]",
+        "ozone": "[%key:component::sensor::entity_component::ozone::name%]",
+        "pm1": "[%key:component::sensor::entity_component::pm1::name%]",
+        "pm10": "[%key:component::sensor::entity_component::pm10::name%]",
+        "pm25": "[%key:component::sensor::entity_component::pm25::name%]",
+        "power_factor": "[%key:component::sensor::entity_component::power_factor::name%]",
+        "power": "[%key:component::sensor::entity_component::power::name%]",
+        "precipitation": "[%key:component::sensor::entity_component::precipitation::name%]",
+        "precipitation_intensity": "[%key:component::sensor::entity_component::precipitation_intensity::name%]",
+        "pressure": "[%key:component::sensor::entity_component::pressure::name%]",
+        "reactive_power": "[%key:component::sensor::entity_component::reactive_power::name%]",
+        "signal_strength": "[%key:component::sensor::entity_component::signal_strength::name%]",
+        "sound_pressure": "[%key:component::sensor::entity_component::sound_pressure::name%]",
+        "speed": "[%key:component::sensor::entity_component::speed::name%]",
+        "sulphur_dioxide": "[%key:component::sensor::entity_component::sulphur_dioxide::name%]",
+        "temperature": "[%key:component::sensor::entity_component::temperature::name%]",
+        "timestamp": "[%key:component::sensor::entity_component::timestamp::name%]",
+        "volatile_organic_compounds": "[%key:component::sensor::entity_component::volatile_organic_compounds::name%]",
+        "volatile_organic_compounds_parts": "[%key:component::sensor::entity_component::volatile_organic_compounds_parts::name%]",
+        "voltage": "[%key:component::sensor::entity_component::voltage::name%]",
+        "volume": "[%key:component::sensor::entity_component::volume::name%]",
+        "volume_storage": "[%key:component::sensor::entity_component::volume_storage::name%]",
+        "water": "[%key:component::sensor::entity_component::water::name%]",
+        "weight": "[%key:component::sensor::entity_component::weight::name%]",
+        "wind_speed": "[%key:component::sensor::entity_component::wind_speed::name%]"
+      }
+    },
+    "state_class": {
+      "options": {
+        "none": "No state class",
+        "measurement": "[%key:component::sensor::entity_component::_::state_attributes::state_class::state::measurement%]",
+        "total": "[%key:component::sensor::entity_component::_::state_attributes::state_class::state::total%]",
+        "total_increasing": "[%key:component::sensor::entity_component::_::state_attributes::state_class::state::total_increasing%]"
+      }
+    }
+  },
   "issues": {
     "entity_id_query_does_full_table_scan": {
       "title": "SQL query does full table scan",

tests/components/sql/__init__.py

@@ -27,6 +27,8 @@ ENTRY_CONFIG = {
     CONF_QUERY: "SELECT 5 as value",
     CONF_COLUMN_NAME: "value",
     CONF_UNIT_OF_MEASUREMENT: "MiB",
+    CONF_DEVICE_CLASS: SensorDeviceClass.DATA_SIZE,
+    CONF_STATE_CLASS: SensorStateClass.TOTAL,
 }
 
 ENTRY_CONFIG_WITH_VALUE_TEMPLATE = {

tests/components/sql/test_config_flow.py

@@ -7,6 +7,8 @@ from sqlalchemy.exc import SQLAlchemyError
 
 from homeassistant import config_entries
 from homeassistant.components.recorder import Recorder
+from homeassistant.components.sensor.const import SensorDeviceClass, SensorStateClass
+from homeassistant.components.sql.config_flow import NONE_SENTINEL
 from homeassistant.components.sql.const import DOMAIN
 from homeassistant.core import HomeAssistant
 from homeassistant.data_entry_flow import FlowResultType
@@ -50,6 +52,8 @@ async def test_form(recorder_mock: Recorder, hass: HomeAssistant) -> None:
         "query": "SELECT 5 as value",
         "column": "value",
         "unit_of_measurement": "MiB",
+        "device_class": SensorDeviceClass.DATA_SIZE,
+        "state_class": SensorStateClass.TOTAL,
     }
     assert len(mock_setup_entry.mock_calls) == 1
 
@@ -151,6 +155,8 @@ async def test_flow_fails_invalid_query(
         "query": "SELECT 5 as value",
         "column": "value",
         "unit_of_measurement": "MiB",
+        "device_class": SensorDeviceClass.DATA_SIZE,
+        "state_class": SensorStateClass.TOTAL,
     }
 
 
@@ -187,6 +193,8 @@ async def test_flow_fails_invalid_column_name(
         "query": "SELECT 5 as value",
         "column": "value",
         "unit_of_measurement": "MiB",
+        "device_class": SensorDeviceClass.DATA_SIZE,
+        "state_class": SensorStateClass.TOTAL,
     }
 
 
@@ -201,6 +209,8 @@ async def test_options_flow(recorder_mock: Recorder, hass: HomeAssistant) -> Non
             "query": "SELECT 5 as value",
             "column": "value",
             "unit_of_measurement": "MiB",
+            "device_class": SensorDeviceClass.DATA_SIZE,
+            "state_class": SensorStateClass.TOTAL,
         },
     )
     entry.add_to_hass(hass)
@@ -225,6 +235,8 @@ async def test_options_flow(recorder_mock: Recorder, hass: HomeAssistant) -> Non
             "column": "size",
             "unit_of_measurement": "MiB",
             "value_template": "{{ value }}",
+            "device_class": SensorDeviceClass.DATA_SIZE,
+            "state_class": SensorStateClass.TOTAL,
         },
     )
 
@@ -235,6 +247,8 @@ async def test_options_flow(recorder_mock: Recorder, hass: HomeAssistant) -> Non
         "column": "size",
         "unit_of_measurement": "MiB",
         "value_template": "{{ value }}",
+        "device_class": SensorDeviceClass.DATA_SIZE,
+        "state_class": SensorStateClass.TOTAL,
     }
 
 
@@ -594,3 +608,79 @@ async def test_full_flow_not_recorder_db(
         "column": "value",
         "unit_of_measurement": "MB",
     }
+
+
+async def test_device_state_class(recorder_mock: Recorder, hass: HomeAssistant) -> None:
+    """Test we get the form."""
+
+    entry = MockConfigEntry(
+        domain=DOMAIN,
+        data={},
+        options={
+            "name": "Get Value",
+            "query": "SELECT 5 as value",
+            "column": "value",
+            "unit_of_measurement": "MiB",
+        },
+    )
+    entry.add_to_hass(hass)
+
+    result = await hass.config_entries.options.async_init(entry.entry_id)
+    assert result["type"] == FlowResultType.FORM
+    assert result["step_id"] == "init"
+
+    with patch(
+        "homeassistant.components.sql.async_setup_entry",
+        return_value=True,
+    ):
+        result2 = await hass.config_entries.options.async_configure(
+            result["flow_id"],
+            user_input={
+                "query": "SELECT 5 as value",
+                "column": "value",
+                "unit_of_measurement": "MiB",
+                "device_class": SensorDeviceClass.DATA_SIZE,
+                "state_class": SensorStateClass.TOTAL,
+            },
+        )
+        await hass.async_block_till_done()
+
+    assert result2["type"] == FlowResultType.CREATE_ENTRY
+    assert result2["data"] == {
+        "name": "Get Value",
+        "query": "SELECT 5 as value",
+        "column": "value",
+        "unit_of_measurement": "MiB",
+        "device_class": SensorDeviceClass.DATA_SIZE,
+        "state_class": SensorStateClass.TOTAL,
+    }
+
+    result = await hass.config_entries.options.async_init(entry.entry_id)
+    assert result["type"] == FlowResultType.FORM
+    assert result["step_id"] == "init"
+
+    with patch(
+        "homeassistant.components.sql.async_setup_entry",
+        return_value=True,
+    ):
+        result3 = await hass.config_entries.options.async_configure(
+            result["flow_id"],
+            user_input={
+                "query": "SELECT 5 as value",
+                "column": "value",
+                "unit_of_measurement": "MiB",
+                "device_class": NONE_SENTINEL,
+                "state_class": NONE_SENTINEL,
+            },
+        )
+        await hass.async_block_till_done()
+
+    assert result3["type"] == FlowResultType.CREATE_ENTRY
+    assert "device_class" not in result3["data"]
+    assert "state_class" not in result3["data"]
+    assert result3["data"] == {
+        "name": "Get Value",
+        "query": "SELECT 5 as value",
+        "column": "value",
+        "unit_of_measurement": "MiB",
+    }

tests/components/sql/test_sensor.py

@@ -457,3 +457,47 @@ async def test_engine_is_disposed_at_stop(
         await hass.async_stop()
 
     assert mock_engine_dispose.call_count == 2
+
+
+async def test_attributes_from_entry_config(
+    recorder_mock: Recorder, hass: HomeAssistant
+) -> None:
+    """Test attributes from entry config."""
+
+    await init_integration(
+        hass,
+        config={
+            "name": "Get Value - With",
+            "query": "SELECT 5 as value",
+            "column": "value",
+            "unit_of_measurement": "MiB",
+            "device_class": SensorDeviceClass.DATA_SIZE,
+            "state_class": SensorStateClass.TOTAL,
+        },
+        entry_id="8693d4782ced4fb1ecca4743f29ab8f1",
+    )
+
+    state = hass.states.get("sensor.get_value_with")
+    assert state.state == "5"
+    assert state.attributes["value"] == 5
+    assert state.attributes["unit_of_measurement"] == "MiB"
+    assert state.attributes["device_class"] == SensorDeviceClass.DATA_SIZE
+    assert state.attributes["state_class"] == SensorStateClass.TOTAL
+
+    await init_integration(
+        hass,
+        config={
+            "name": "Get Value - Without",
+            "query": "SELECT 5 as value",
+            "column": "value",
+            "unit_of_measurement": "MiB",
+        },
+        entry_id="7aec7cd8045fba4778bb0621469e3cd9",
+    )
+
+    state = hass.states.get("sensor.get_value_without")
+    assert state.state == "5"
+    assert state.attributes["value"] == 5
+    assert state.attributes["unit_of_measurement"] == "MiB"
+    assert "device_class" not in state.attributes
+    assert "state_class" not in state.attributes