Merge util._async_validate_and_get_session_maker_for_db_url

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

Author: davidrapan

homeassistant/components/sql/sensor.py

@@ -289,12 +289,12 @@ def _update(self) -> None:
             try:
                 self._process(session.execute(self._lambda_stmt))
             except SQLAlchemyError as err:
+                session.rollback()
                 _LOGGER.error(
                     "Error executing query %s: %s",
                     self._query,
                     redact_credentials(str(err)),
                 )
-                session.rollback()
 
     async def async_update(self) -> None:
         """Retrieve sensor data from the query using the right executor."""
@@ -304,12 +304,12 @@ async def async_update(self) -> None:
                 try:
                     self._process(await session.execute(self._lambda_stmt))
                 except SQLAlchemyError as err:
+                    await session.rollback()
                     _LOGGER.error(
                         "Error executing query %s: %s",
                         self._query,
                         redact_credentials(str(err)),
                     )
-                    await session.rollback()
         elif self._use_database_executor:
             await get_instance(self.hass).async_add_executor_job(self._update)
         else:

homeassistant/components/sql/services.py

@@ -96,12 +96,12 @@ def _execute_and_convert_query() -> list[JsonValueType]:
             try:
                 return _process(session.execute(generate_lambda_stmt(query_str)))
             except SQLAlchemyError as err:
+                session.rollback()
                 _LOGGER.debug(
                     "Error executing query %s: %s",
                     query_str,
                     redact_credentials(str(err)),
                 )
-                session.rollback()
                 raise
 
     try:
@@ -112,12 +112,12 @@ def _execute_and_convert_query() -> list[JsonValueType]:
                         await session.execute(generate_lambda_stmt(query_str))
                     )
                 except SQLAlchemyError as err:
+                    await session.rollback()
                     _LOGGER.debug(
                         "Error executing query %s: %s",
                         query_str,
                         redact_credentials(str(err)),
                     )
-                    await session.rollback()
                     raise
         elif use_database_executor:
             result = await get_instance(call.hass).async_add_executor_job(

homeassistant/components/sql/util.py

@@ -30,6 +30,7 @@
 
 _LOGGER = logging.getLogger(__name__)
 
+_SQL_SELECT = sqlalchemy.text("SELECT 1;")
 _SQL_LAMBDA_CACHE: LRUCache = LRUCache(1000)
 
 
@@ -107,11 +108,8 @@ 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
+    elif sessmaker := await _async_validate_and_get_session_maker_for_db_url(
+        hass, db_url
     ):
         sql_data.session_makers_by_db_url[db_url] = sessmaker
     else:
@@ -210,57 +208,46 @@ async def _shutdown_db_engines(_: Event) -> None:
 
 
 async def _async_validate_and_get_session_maker_for_db_url(
-    db_url: str,
-) -> async_scoped_session[AsyncSession] | None:
+    hass: HomeAssistant, db_url: str
+) -> async_scoped_session[AsyncSession] | scoped_session[Session] | None:
     """Validate the db_url and return a async session maker."""
     try:
-        maker = async_scoped_session(
-            async_sessionmaker(
-                bind=create_async_engine(db_url, future=True), future=True
-            ),
-            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.
-    """
-    try:
-        maker = scoped_session(
-            sessionmaker(
-                bind=sqlalchemy.create_engine(db_url, future=True), future=True
+        if "+aiomysql" in db_url or "+aiosqlite" in db_url or "+asyncpg" in db_url:
+            maker = async_scoped_session(
+                async_sessionmaker(
+                    bind=create_async_engine(db_url, future=True), future=True
+                ),
+                scopefunc=asyncio.current_task,
             )
-        )
-        # Run a dummy query just to test the db_url
-        with maker() as session:
-            session.execute(sqlalchemy.text("SELECT 1;"))
+            # Run a dummy query just to test the db_url
+            async with maker() as session:
+                await session.execute(_SQL_SELECT)
+            return maker
+
+        def _get_session_maker_for_db_url() -> scoped_session[Session] | None:
+            """Validate the db_url and return a session maker.
+
+            This does I/O and should be run in the executor.
+            """
+            maker = scoped_session(
+                sessionmaker(
+                    bind=sqlalchemy.create_engine(db_url, future=True), future=True
+                )
+            )
+            # Run a dummy query just to test the db_url
+            with maker() as session:
+                session.execute(_SQL_SELECT)
+            return maker
 
+        return await hass.async_add_executor_job(_get_session_maker_for_db_url)
     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
+
+    return None
 
 
 def generate_lambda_stmt(query: str) -> StatementLambdaElement:

tests/components/sql/test_sensor.py

@@ -325,6 +325,9 @@ def execute(self, query: Any) -> None:
             """Execute the query."""
             raise SQLAlchemyError("sqlite://homeassistant:[email protected]")
 
+        def rollback(self) -> None:
+            pass
+
     with patch(
         "homeassistant.components.sql.util.scoped_session",
         return_value=MockSession,

tests/components/sql/test_services.py

@@ -230,7 +230,7 @@ async def test_query_service_invalid_db_url(hass: HomeAssistant) -> None:
 
     with (
         patch(
-            "homeassistant.components.sql.util._validate_and_get_session_maker_for_db_url",
+            "homeassistant.components.sql.util._async_validate_and_get_session_maker_for_db_url",
             return_value=None,
         ),
         pytest.raises(