Add support for aiomysql, aiosqlite and asyncpg

Signed-off-by: David Rapan <[email protected]>

Author: davidrapan

homeassistant/components/sql/manifest.json

@@ -6,5 +6,12 @@
   "config_flow": true,
   "documentation": "https://www.home-assistant.io/integrations/sql",
   "iot_class": "local_polling",
-  "requirements": ["SQLAlchemy==2.0.41", "sqlparse==0.5.0"]
+  "requirements": [
+    "aiomysql==0.3.2",
+    "aiosqlite==0.21.0",
+    "asyncpg==0.30.0",
+    "greenlet==3.2.4",
+    "SQLAlchemy==2.0.41",
+    "sqlparse==0.5.0"
+  ]
 }

homeassistant/components/sql/models.py

@@ -4,7 +4,8 @@
 
 from dataclasses import dataclass
 
-from sqlalchemy.orm import scoped_session
+from sqlalchemy.ext.asyncio import AsyncSession, async_scoped_session
+from sqlalchemy.orm import Session, scoped_session
 
 from homeassistant.core import CALLBACK_TYPE
 
@@ -14,4 +15,6 @@ class SQLData:
     """Data for the sql integration."""
 
     shutdown_event_cancel: CALLBACK_TYPE
-    session_makers_by_db_url: dict[str, scoped_session]
+    session_makers_by_db_url: dict[
+        str, async_scoped_session[AsyncSession] | scoped_session[Session]
+    ]

homeassistant/components/sql/sensor.py

@@ -5,11 +5,12 @@
 from datetime import date
 import decimal
 import logging
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 from sqlalchemy.engine import Result
 from sqlalchemy.exc import SQLAlchemyError
-from sqlalchemy.orm import scoped_session
+from sqlalchemy.ext.asyncio import AsyncSession, async_scoped_session
+from sqlalchemy.orm import Session, scoped_session
 
 from homeassistant.components.recorder import CONF_DB_URL, get_instance
 from homeassistant.components.sensor import CONF_STATE_CLASS
@@ -200,7 +201,7 @@ class SQLSensor(ManualTriggerSensorEntity):
     def __init__(
         self,
         trigger_entity_config: ConfigType,
-        sessmaker: scoped_session,
+        sessmaker: async_scoped_session[AsyncSession] | scoped_session[Session],
         query: str,
         column: str,
         value_template: ValueTemplate | None,
@@ -243,31 +244,10 @@ 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)
-        else:
-            await self.hass.async_add_executor_job(self._update)
-
-    def _update(self) -> None:
-        """Retrieve sensor data from the query."""
+    def _process(self, result: Result) -> None:
+        """Process the SQL result."""
         data = None
         extra_state_attributes = {}
-        self._attr_extra_state_attributes = {}
-        sess: scoped_session = self.sessionmaker()
-        try:
-            result: Result = sess.execute(self._lambda_stmt)
-        except SQLAlchemyError as err:
-            _LOGGER.error(
-                "Error executing query %s: %s",
-                self._query,
-                redact_credentials(str(err)),
-            )
-            sess.rollback()
-            sess.close()
-            return
-
         for res in result.mappings():
             _LOGGER.debug("Query %s result in %s", self._query, res.items())
             data = res[self._column_name]
@@ -298,4 +278,35 @@ def _update(self) -> None:
         if data is None:
             _LOGGER.warning("%s returned no results", self._query)
 
-        sess.close()
+    def _update(self) -> None:
+        """Retrieve sensor data from the query."""
+        self._attr_extra_state_attributes = {}
+        try:
+            if TYPE_CHECKING:
+                assert isinstance(self.sessionmaker, scoped_session)
+            with self.sessionmaker() as session:
+                self._process(session.execute(self._lambda_stmt))
+        except SQLAlchemyError as err:
+            _LOGGER.error(
+                "Error executing query %s: %s",
+                self._query,
+                redact_credentials(str(err)),
+            )
+
+    async def async_update(self) -> None:
+        """Retrieve sensor data from the query using the right executor."""
+        if isinstance(self.sessionmaker, async_scoped_session):
+            self._attr_extra_state_attributes = {}
+            try:
+                async with self.sessionmaker() as session:
+                    self._process(await session.execute(self._lambda_stmt))
+            except SQLAlchemyError as err:
+                _LOGGER.error(
+                    "Error executing query %s: %s",
+                    self._query,
+                    redact_credentials(str(err)),
+                )
+        elif self._use_database_executor:
+            await get_instance(self.hass).async_add_executor_job(self._update)
+        else:
+            await self.hass.async_add_executor_job(self._update)

homeassistant/components/sql/services.py

@@ -5,10 +5,12 @@
 import datetime
 import decimal
 import logging
+from typing import TYPE_CHECKING
 
 from sqlalchemy.engine import Result
 from sqlalchemy.exc import SQLAlchemyError
-from sqlalchemy.orm import Session
+from sqlalchemy.ext.asyncio import async_scoped_session
+from sqlalchemy.orm import scoped_session
 import voluptuous as vol
 
 from homeassistant.components.recorder import CONF_DB_URL, get_instance
@@ -70,39 +72,52 @@ async def _async_query_service(
             translation_placeholders={"db_url": redact_credentials(db_url)},
         )
 
+    def _process(result: Result) -> list[JsonValueType]:
+        rows: list[JsonValueType] = []
+        for row in result.mappings():
+            processed_row: dict[str, JsonValueType] = {}
+            for key, value in row.items():
+                if isinstance(value, decimal.Decimal):
+                    processed_row[key] = float(value)
+                elif isinstance(value, datetime.date):
+                    processed_row[key] = value.isoformat()
+                elif isinstance(value, (bytes, bytearray)):
+                    processed_row[key] = f"0x{value.hex()}"
+                else:
+                    processed_row[key] = value
+            rows.append(processed_row)
+        return rows
+
     def _execute_and_convert_query() -> list[JsonValueType]:
         """Execute the query and return the results with converted types."""
-        sess: Session = sessmaker()
         try:
-            result: Result = sess.execute(generate_lambda_stmt(query_str))
+            if TYPE_CHECKING:
+                assert isinstance(sessmaker, scoped_session)
+            with sessmaker() as session:
+                return _process(session.execute(generate_lambda_stmt(query_str)))
         except SQLAlchemyError as err:
             _LOGGER.debug(
                 "Error executing query %s: %s",
                 query_str,
                 redact_credentials(str(err)),
             )
-            sess.rollback()
             raise
-        else:
-            rows: list[JsonValueType] = []
-            for row in result.mappings():
-                processed_row: dict[str, JsonValueType] = {}
-                for key, value in row.items():
-                    if isinstance(value, decimal.Decimal):
-                        processed_row[key] = float(value)
-                    elif isinstance(value, datetime.date):
-                        processed_row[key] = value.isoformat()
-                    elif isinstance(value, (bytes, bytearray)):
-                        processed_row[key] = f"0x{value.hex()}"
-                    else:
-                        processed_row[key] = value
-                rows.append(processed_row)
-            return rows
-        finally:
-            sess.close()
 
     try:
-        if use_database_executor:
+        if isinstance(sessmaker, async_scoped_session):
+            try:
+                async with sessmaker() as session:
+                    result = _process(
+                        await session.execute(generate_lambda_stmt(query_str))
+                    )
+            except SQLAlchemyError as err:
+                _LOGGER.debug(
+                    "Error executing query %s: %s",
+                    query_str,
+                    redact_credentials(str(err)),
+                )
+                raise
+        elif use_database_executor:
             result = await get_instance(call.hass).async_add_executor_job(
                 _execute_and_convert_query
             )

homeassistant/components/sql/util.py

@@ -2,11 +2,18 @@
 
 from __future__ import annotations
 
+import asyncio
 import logging
 
 import sqlalchemy
 from sqlalchemy import lambda_stmt
 from sqlalchemy.exc import SQLAlchemyError
+from sqlalchemy.ext.asyncio import (
+    AsyncSession,
+    async_scoped_session,
+    async_sessionmaker,
+    create_async_engine,
+)
 from sqlalchemy.orm import Session, scoped_session, sessionmaker
 from sqlalchemy.sql.lambdas import StatementLambdaElement
 from sqlalchemy.util import LRUCache
@@ -55,7 +62,9 @@ def validate_sql_select(value: str) -> str:
 
 async def async_create_sessionmaker(
     hass: HomeAssistant, db_url: str
-) -> tuple[scoped_session | None, bool, bool]:
+) -> tuple[
+    async_scoped_session[AsyncSession] | scoped_session[Session] | None, bool, bool
+]:
     """Create a session maker for the given db_url.
 
     This function gets or creates a SQLAlchemy `scoped_session` for the given
@@ -83,7 +92,7 @@ async def async_create_sessionmaker(
         uses_recorder_db = False
     else:
         uses_recorder_db = db_url == instance.db_url
-    sessmaker: scoped_session | None
+    sessmaker: async_scoped_session[AsyncSession] | scoped_session[Session] | None
     sql_data = _async_get_or_init_domain_data(hass)
     use_database_executor = False
     if uses_recorder_db and instance.dialect_name == SupportedDialect.SQLITE:
@@ -98,6 +107,9 @@ async def async_create_sessionmaker(
     # for every sensor.
     elif db_url in sql_data.session_makers_by_db_url:
         sessmaker = sql_data.session_makers_by_db_url[db_url]
+    elif "+aiomysql" in db_url or "+aiosqlite" in db_url or "+asyncpg" in db_url:
+        if sessmaker := await _async_validate_and_get_session_maker_for_db_url(db_url):
+            sql_data.session_makers_by_db_url[db_url] = sessmaker
     elif sessmaker := await hass.async_add_executor_job(
         _validate_and_get_session_maker_for_db_url, db_url
     ):
@@ -169,7 +181,9 @@ def _async_get_or_init_domain_data(hass: HomeAssistant) -> SQLData:
         sql_data: SQLData = hass.data[DOMAIN]
         return sql_data
 
-    session_makers_by_db_url: dict[str, scoped_session] = {}
+    session_makers_by_db_url: dict[
+        str, async_scoped_session[AsyncSession] | scoped_session[Session]
+    ] = {}
 
     #
     # Ensure we dispose of all engines at shutdown
@@ -178,10 +192,13 @@ def _async_get_or_init_domain_data(hass: HomeAssistant) -> SQLData:
     # Shutdown all sessions in the executor since they will
     # do blocking I/O
     #
-    def _shutdown_db_engines(event: Event) -> None:
+    async def _shutdown_db_engines(_: Event) -> None:
         """Shutdown all database engines."""
         for sessmaker in session_makers_by_db_url.values():
-            sessmaker.connection().engine.dispose()
+            if isinstance(sessmaker, async_scoped_session):
+                await (await sessmaker.connection()).engine.dispose()
+            else:
+                sessmaker.connection().engine.dispose()
 
     cancel_shutdown = hass.bus.async_listen_once(
         EVENT_HOMEASSISTANT_STOP, _shutdown_db_engines
@@ -192,18 +209,46 @@ def _shutdown_db_engines(event: Event) -> None:
     return sql_data
 
 
-def _validate_and_get_session_maker_for_db_url(db_url: str) -> scoped_session | None:
+async def _async_validate_and_get_session_maker_for_db_url(
+    db_url: str,
+) -> async_scoped_session[AsyncSession] | None:
+    """Validate the db_url and return a async session maker."""
+    try:
+        maker = async_scoped_session(
+            async_sessionmaker(bind=create_async_engine(db_url)),
+            scopefunc=asyncio.current_task,
+        )
+        # Run a dummy query just to test the db_url
+        async with maker() as session:
+            await session.execute(sqlalchemy.text("SELECT 1;"))
+
+    except SQLAlchemyError as err:
+        _LOGGER.error(
+            "Couldn't connect using %s DB_URL: %s",
+            redact_credentials(db_url),
+            redact_credentials(str(err)),
+        )
+        return None
+    else:
+        return maker
+
+
+def _validate_and_get_session_maker_for_db_url(
+    db_url: str,
+) -> scoped_session[Session] | None:
     """Validate the db_url and return a session maker.
 
     This does I/O and should be run in the executor.
     """
-    sess: Session | None = None
     try:
-        engine = sqlalchemy.create_engine(db_url, future=True)
-        sessmaker = scoped_session(sessionmaker(bind=engine, future=True))
+        maker = scoped_session(
+            sessionmaker(
+                bind=sqlalchemy.create_engine(db_url, future=True), future=True
+            )
+        )
         # Run a dummy query just to test the db_url
-        sess = sessmaker()
-        sess.execute(sqlalchemy.text("SELECT 1;"))
+        with maker() as session:
+            session.execute(sqlalchemy.text("SELECT 1;"))
 
     except SQLAlchemyError as err:
         _LOGGER.error(
@@ -213,10 +258,7 @@ def _validate_and_get_session_maker_for_db_url(db_url: str) -> scoped_session |
         )
         return None
     else:
-        return sessmaker
-    finally:
-        if sess:
-            sess.close()
+        return maker
 
 
 def generate_lambda_stmt(query: str) -> StatementLambdaElement: