Improve code quality sql (#65321)

homeassistant/components/sql/sensor.py

@@ -1,8 +1,7 @@
 """Sensor from an SQL Query."""
 from __future__ import annotations
 
-import datetime
-import decimal
+from datetime import date
 import logging
 import re
 
@@ -11,11 +10,15 @@ from sqlalchemy.orm import scoped_session, sessionmaker
 import voluptuous as vol
 
 from homeassistant.components.recorder import CONF_DB_URL, DEFAULT_DB_FILE, DEFAULT_URL
-from homeassistant.components.sensor import PLATFORM_SCHEMA, SensorEntity
+from homeassistant.components.sensor import (
+    PLATFORM_SCHEMA as PARENT_PLATFORM_SCHEMA,
+    SensorEntity,
+)
 from homeassistant.const import CONF_NAME, CONF_UNIT_OF_MEASUREMENT, CONF_VALUE_TEMPLATE
 from homeassistant.core import HomeAssistant
 import homeassistant.helpers.config_validation as cv
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
+from homeassistant.helpers.template import Template
 from homeassistant.helpers.typing import ConfigType, DiscoveryInfoType
 
 _LOGGER = logging.getLogger(__name__)
@@ -27,12 +30,12 @@ CONF_QUERY = "query"
 DB_URL_RE = re.compile("//.*:.*@")
 
 
-def redact_credentials(data):
+def redact_credentials(data: str) -> str:
     """Redact credentials from string data."""
     return DB_URL_RE.sub("//****:****@", data)
 
 
-def validate_sql_select(value):
+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")
@@ -49,7 +52,7 @@ _QUERY_SCHEME = vol.Schema(
     }
 )
 
-PLATFORM_SCHEMA = PLATFORM_SCHEMA.extend(
+PLATFORM_SCHEMA = PARENT_PLATFORM_SCHEMA.extend(
     {vol.Required(CONF_QUERIES): [_QUERY_SCHEME], vol.Optional(CONF_DB_URL): cv.string}
 )
 
@@ -64,7 +67,7 @@ def setup_platform(
     if not (db_url := config.get(CONF_DB_URL)):
         db_url = DEFAULT_URL.format(hass_config_path=hass.config.path(DEFAULT_DB_FILE))
 
-    sess = None
+    sess: scoped_session | None = None
     try:
         engine = sqlalchemy.create_engine(db_url)
         sessmaker = scoped_session(sessionmaker(bind=engine))
@@ -87,11 +90,11 @@ def setup_platform(
     queries = []
 
     for query in config[CONF_QUERIES]:
-        name = query.get(CONF_NAME)
-        query_str = query.get(CONF_QUERY)
-        unit = query.get(CONF_UNIT_OF_MEASUREMENT)
-        value_template = query.get(CONF_VALUE_TEMPLATE)
-        column_name = query.get(CONF_COLUMN_NAME)
+        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]
 
         if value_template is not None:
             value_template.hass = hass
@@ -115,60 +118,32 @@ def setup_platform(
 class SQLSensor(SensorEntity):
     """Representation of an SQL sensor."""
 
-    def __init__(self, name, sessmaker, query, column, unit, value_template):
+    def __init__(
+        self,
+        name: str,
+        sessmaker: scoped_session,
+        query: str,
+        column: str,
+        unit: str | None,
+        value_template: Template | None,
+    ) -> None:
         """Initialize the SQL sensor."""
-        self._name = name
+        self._attr_name = name
         self._query = query
-        self._unit_of_measurement = unit
+        self._attr_native_unit_of_measurement = unit
         self._template = value_template
         self._column_name = column
         self.sessionmaker = sessmaker
-        self._state = None
-        self._attributes = None
+        self._attr_extra_state_attributes = {}
 
-    @property
-    def name(self):
-        """Return the name of the query."""
-        return self._name
-
-    @property
-    def native_value(self):
-        """Return the query's current state."""
-        return self._state
-
-    @property
-    def native_unit_of_measurement(self):
-        """Return the unit of measurement."""
-        return self._unit_of_measurement
-
-    @property
-    def extra_state_attributes(self):
-        """Return the state attributes."""
-        return self._attributes
-
-    def update(self):
+    def update(self) -> None:
         """Retrieve sensor data from the query."""
 
         data = None
+        self._attr_extra_state_attributes = {}
+        sess: scoped_session = self.sessionmaker()
         try:
-            sess = self.sessionmaker()
             result = sess.execute(self._query)
-            self._attributes = {}
-
-            if not result.returns_rows or result.rowcount == 0:
-                _LOGGER.warning("%s returned no results", self._query)
-                self._state = None
-                return
-
-            for res in result.mappings():
-                _LOGGER.debug("result = %s", res.items())
-                data = res[self._column_name]
-                for key, value in res.items():
-                    if isinstance(value, decimal.Decimal):
-                        value = float(value)
-                    if isinstance(value, datetime.date):
-                        value = str(value)
-                    self._attributes[key] = value
         except sqlalchemy.exc.SQLAlchemyError as err:
             _LOGGER.error(
                 "Error executing query %s: %s",
@@ -176,12 +151,27 @@ class SQLSensor(SensorEntity):
                 redact_credentials(str(err)),
             )
             return
-        finally:
-            sess.close()
+
+        _LOGGER.debug("Result %s, ResultMapping %s", result, result.mappings())
+
+        for res in result.mappings():
+            _LOGGER.debug("result = %s", res.items())
+            data = res[self._column_name]
+            for key, value in res.items():
+                if isinstance(value, float):
+                    value = float(value)
+                if isinstance(value, date):
+                    value = value.isoformat()
+                self._attr_extra_state_attributes[key] = value
 
         if data is not None and self._template is not None:
-            self._state = self._template.async_render_with_possible_json_value(
-                data, None
+            self._attr_native_value = (
+                self._template.async_render_with_possible_json_value(data, None)
             )
         else:
-            self._state = data
+            self._attr_native_value = data
+
+        if not data:
+            _LOGGER.warning("%s returned no results", self._query)
+
+        sess.close()