Merge util._async_validate_and_get_session_maker_for_db_url

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

Author: davidrapan

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(