Fix error for when column does not exist in query to SQL (#91166)

Co-authored-by: J. Nick Koston <nick@koston.org>

CODEOWNERS

@@ -1138,8 +1138,8 @@ build.json @home-assistant/supervisor
 /homeassistant/components/splunk/ @Bre77
 /homeassistant/components/spotify/ @frenck
 /tests/components/spotify/ @frenck
-/homeassistant/components/sql/ @dgomes @gjohansson-ST
-/tests/components/sql/ @dgomes @gjohansson-ST
+/homeassistant/components/sql/ @dgomes @gjohansson-ST @dougiteixeira
+/tests/components/sql/ @dgomes @gjohansson-ST @dougiteixeira
 /homeassistant/components/squeezebox/ @rajlaud
 /tests/components/squeezebox/ @rajlaud
 /homeassistant/components/srp_energy/ @briglx

homeassistant/components/sql/config_flow.py

@@ -6,7 +6,7 @@ from typing import Any
 
 import sqlalchemy
 from sqlalchemy.engine import Result
-from sqlalchemy.exc import SQLAlchemyError
+from sqlalchemy.exc import NoSuchColumnError, SQLAlchemyError
 from sqlalchemy.orm import Session, scoped_session, sessionmaker
 import voluptuous as vol
 
@@ -69,9 +69,17 @@ def validate_query(db_url: str, query: str, column: str) -> bool:
         _LOGGER.debug("Execution error %s", error)
         if sess:
             sess.close()
+            engine.dispose()
         raise ValueError(error) from error
 
     for res in result.mappings():
+        if column not in res:
+            _LOGGER.debug("Column `%s` is not returned by the query", column)
+            if sess:
+                sess.close()
+                engine.dispose()
+            raise NoSuchColumnError(f"Column {column} is not returned by the query.")
+
         data = res[column]
         _LOGGER.debug("Return value from query: %s", data)
 
@@ -100,6 +108,7 @@ class SQLConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
     ) -> FlowResult:
         """Handle the user step."""
         errors = {}
+        description_placeholders = {}
 
         if user_input is not None:
             db_url = user_input.get(CONF_DB_URL)
@@ -116,6 +125,9 @@ class SQLConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
                 await self.hass.async_add_executor_job(
                     validate_query, db_url_for_validation, query, column
                 )
+            except NoSuchColumnError:
+                errors["column"] = "column_invalid"
+                description_placeholders = {"column": column}
             except SQLAlchemyError:
                 errors["db_url"] = "db_url_invalid"
             except ValueError:
@@ -143,6 +155,7 @@ class SQLConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
             step_id="user",
             data_schema=self.add_suggested_values_to_schema(CONFIG_SCHEMA, user_input),
             errors=errors,
+            description_placeholders=description_placeholders,
         )
 
 
@@ -158,6 +171,7 @@ class SQLOptionsFlowHandler(config_entries.OptionsFlow):
     ) -> FlowResult:
         """Manage SQL options."""
         errors = {}
+        description_placeholders = {}
 
         if user_input is not None:
             db_url = user_input.get(CONF_DB_URL)
@@ -171,6 +185,9 @@ class SQLOptionsFlowHandler(config_entries.OptionsFlow):
                 await self.hass.async_add_executor_job(
                     validate_query, db_url_for_validation, query, column
                 )
+            except NoSuchColumnError:
+                errors["column"] = "column_invalid"
+                description_placeholders = {"column": column}
             except SQLAlchemyError:
                 errors["db_url"] = "db_url_invalid"
             except ValueError:
@@ -193,4 +210,5 @@ class SQLOptionsFlowHandler(config_entries.OptionsFlow):
                 OPTIONS_SCHEMA, user_input or self.entry.options
             ),
             errors=errors,
+            description_placeholders=description_placeholders,
         )

homeassistant/components/sql/manifest.json

@@ -1,7 +1,7 @@
 {
   "domain": "sql",
   "name": "SQL",
-  "codeowners": ["@dgomes", "@gjohansson-ST"],
+  "codeowners": ["@dgomes", "@gjohansson-ST", "@dougiteixeira"],
   "config_flow": true,
   "documentation": "https://www.home-assistant.io/integrations/sql",
   "iot_class": "local_polling",

homeassistant/components/sql/strings.json

@@ -5,7 +5,8 @@
     },
     "error": {
       "db_url_invalid": "Database URL invalid",
-      "query_invalid": "SQL Query invalid"
+      "query_invalid": "SQL Query invalid",
+      "column_invalid": "The column `{column}` is not returned by the query"
     },
     "step": {
       "user": {
@@ -51,7 +52,8 @@
     },
     "error": {
       "db_url_invalid": "[%key:component::sql::config::error::db_url_invalid%]",
-      "query_invalid": "[%key:component::sql::config::error::query_invalid%]"
+      "query_invalid": "[%key:component::sql::config::error::query_invalid%]",
+      "column_invalid": "[%key:component::sql::config::error::column_invalid%]"
     }
   },
   "issues": {

tests/components/sql/__init__.py

@@ -42,6 +42,19 @@ ENTRY_CONFIG_INVALID_QUERY_OPT = {
     CONF_UNIT_OF_MEASUREMENT: "MiB",
 }
 
+ENTRY_CONFIG_INVALID_COLUMN_NAME = {
+    CONF_NAME: "Get Value",
+    CONF_QUERY: "SELECT 5 as value",
+    CONF_COLUMN_NAME: "size",
+    CONF_UNIT_OF_MEASUREMENT: "MiB",
+}
+
+ENTRY_CONFIG_INVALID_COLUMN_NAME_OPT = {
+    CONF_QUERY: "SELECT 5 as value",
+    CONF_COLUMN_NAME: "size",
+    CONF_UNIT_OF_MEASUREMENT: "MiB",
+}
+
 ENTRY_CONFIG_NO_RESULTS = {
     CONF_NAME: "Get Value",
     CONF_QUERY: "SELECT kalle as value from no_table;",

tests/components/sql/test_config_flow.py

@@ -13,6 +13,8 @@ from homeassistant.data_entry_flow import FlowResultType
 
 from . import (
     ENTRY_CONFIG,
+    ENTRY_CONFIG_INVALID_COLUMN_NAME,
+    ENTRY_CONFIG_INVALID_COLUMN_NAME_OPT,
     ENTRY_CONFIG_INVALID_QUERY,
     ENTRY_CONFIG_INVALID_QUERY_OPT,
     ENTRY_CONFIG_NO_RESULTS,
@@ -120,6 +122,43 @@ async def test_flow_fails_invalid_query(
     }
 
 
+async def test_flow_fails_invalid_column_name(
+    recorder_mock: Recorder, hass: HomeAssistant
+) -> None:
+    """Test config flow fails invalid column name."""
+    result4 = await hass.config_entries.flow.async_init(
+        DOMAIN, context={"source": config_entries.SOURCE_USER}
+    )
+
+    assert result4["type"] == FlowResultType.FORM
+    assert result4["step_id"] == "user"
+
+    result5 = await hass.config_entries.flow.async_configure(
+        result4["flow_id"],
+        user_input=ENTRY_CONFIG_INVALID_COLUMN_NAME,
+    )
+
+    assert result5["type"] == FlowResultType.FORM
+    assert result5["errors"] == {
+        "column": "column_invalid",
+    }
+
+    result5 = await hass.config_entries.flow.async_configure(
+        result4["flow_id"],
+        user_input=ENTRY_CONFIG,
+    )
+
+    assert result5["type"] == FlowResultType.CREATE_ENTRY
+    assert result5["title"] == "Get Value"
+    assert result5["options"] == {
+        "name": "Get Value",
+        "query": "SELECT 5 as value",
+        "column": "value",
+        "unit_of_measurement": "MiB",
+        "value_template": None,
+    }
+
+
 async def test_options_flow(recorder_mock: Recorder, hass: HomeAssistant) -> None:
     """Test options config flow."""
     entry = MockConfigEntry(
@@ -318,6 +357,60 @@ async def test_options_flow_fails_invalid_query(
     }
 
 
+async def test_options_flow_fails_invalid_column_name(
+    recorder_mock: Recorder, hass: HomeAssistant
+) -> None:
+    """Test options flow fails invalid column name."""
+    entry = MockConfigEntry(
+        domain=DOMAIN,
+        data={},
+        options={
+            "name": "Get Value",
+            "query": "SELECT 5 as value",
+            "column": "value",
+            "unit_of_measurement": "MiB",
+            "value_template": None,
+        },
+    )
+    entry.add_to_hass(hass)
+
+    with patch(
+        "homeassistant.components.sql.async_setup_entry",
+        return_value=True,
+    ):
+        assert await hass.config_entries.async_setup(entry.entry_id)
+        await hass.async_block_till_done()
+
+    result = await hass.config_entries.options.async_init(entry.entry_id)
+
+    result2 = await hass.config_entries.options.async_configure(
+        result["flow_id"],
+        user_input=ENTRY_CONFIG_INVALID_COLUMN_NAME_OPT,
+    )
+
+    assert result2["type"] == FlowResultType.FORM
+    assert result2["errors"] == {
+        "column": "column_invalid",
+    }
+
+    result4 = await hass.config_entries.options.async_configure(
+        result["flow_id"],
+        user_input={
+            "query": "SELECT 5 as value",
+            "column": "value",
+            "unit_of_measurement": "MiB",
+        },
+    )
+
+    assert result4["type"] == FlowResultType.CREATE_ENTRY
+    assert result4["data"] == {
+        "name": "Get Value",
+        "query": "SELECT 5 as value",
+        "column": "value",
+        "unit_of_measurement": "MiB",
+    }
+
+
 async def test_options_flow_db_url_empty(
     recorder_mock: Recorder, hass: HomeAssistant
 ) -> None: