feat: Improve query security and integrity checks (#180)

* feat(update_dashboard): Check that all columns are returned from query

* feat: Deny queries that contain potentially undesirable keywords

* `devtools::document()` (GitHub Actions)

* chore: Check query before running in `test_query()` too

* chore(py): Export UnsafeQueryError from types

* chore: Add news/changelog item

* docs(py): Document raises unsafe query error

---------

Co-authored-by: gadenbuie <[email protected]>

Author: gadenbuie

pkg-py/CHANGELOG.md

@@ -17,6 +17,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 * The tools used in a `QueryChat` chatbot are now configurable. Use the new `tools` parameter of `QueryChat()` to select either or both `"query"` or `"update"` tools. Choose `tools=["update"]` if you only want QueryChat to be able to update the dashboard (useful when you want to be 100% certain that the LLM will not see _any_ raw data). (#168)
 
+### Improvements
+
+* The update tool now requires that the SQL query returns all columns from the original data source, ensuring that the dashboard can display the complete data frame after filtering or sorting. If the query does not return all columns, an informative error message will be provided. (#180)
+
+* Obvious SQL keywords that lead to data modification (e.g., `INSERT`, `UPDATE`, `DELETE`, `DROP`, etc.) are now prohibited in queries run via the query tool or update tool, to prevent accidental data changes. If such keywords are detected, an informative error message will be provided. (#180)
+
 ## [0.3.0] - 2025-12-10
 
 ### Breaking Changes

pkg-py/src/querychat/_datasource.py

@@ -9,11 +9,17 @@
 from sqlalchemy import inspect, text
 from sqlalchemy.sql import sqltypes
 
+from ._utils import check_query
+
 if TYPE_CHECKING:
     from narwhals.stable.v1.typing import IntoFrame
     from sqlalchemy.engine import Connection, Engine
 
 
+class MissingColumnsError(ValueError):
+    """Raised when a query result is missing required columns."""
+
+
 class DataSource(ABC):
     """
     An abstract class defining the interface for data sources used by QueryChat.
@@ -70,6 +76,34 @@ def execute_query(self, query: str) -> pd.DataFrame:
         """
         ...
 
+    @abstractmethod
+    def test_query(
+        self, query: str, *, require_all_columns: bool = False
+    ) -> pd.DataFrame:
+        """
+        Test SQL query by fetching only one row.
+
+        Parameters
+        ----------
+        query
+            SQL query to test
+        require_all_columns
+            If True, validates that result includes all original table columns.
+            Additional computed columns are allowed.
+
+        Returns
+        -------
+        :
+            Query results as a pandas DataFrame with at most one row
+
+        Raises
+        ------
+        MissingColumnsError
+            If require_all_columns is True and result is missing required columns
+
+        """
+        ...
+
     @abstractmethod
     def get_data(self) -> pd.DataFrame:
         """
@@ -136,6 +170,9 @@ def __init__(self, df: IntoFrame, table_name: str):
 SET lock_configuration = true;
         """)
 
+        # Store original column names for validation
+        self._colnames = list(self._df.columns)
+
     def get_db_type(self) -> str:
         """
         Get the database type.
@@ -225,9 +262,60 @@ def execute_query(self, query: str) -> pd.DataFrame:
         :
             Query results as pandas DataFrame
 
+        Raises
+        ------
+        UnsafeQueryError
+            If the query starts with a disallowed SQL operation
+
         """
+        check_query(query)
         return self._conn.execute(query).df()
 
+    def test_query(
+        self, query: str, *, require_all_columns: bool = False
+    ) -> pd.DataFrame:
+        """
+        Test query by fetching only one row.
+
+        Parameters
+        ----------
+        query
+            SQL query to test
+        require_all_columns
+            If True, validates that result includes all original table columns
+
+        Returns
+        -------
+        :
+            Query results with at most one row
+
+        Raises
+        ------
+        UnsafeQueryError
+            If the query starts with a disallowed SQL operation
+        MissingColumnsError
+            If require_all_columns is True and result is missing required columns
+
+        """
+        check_query(query)
+        result = self._conn.execute(f"{query} LIMIT 1").df()
+
+        if require_all_columns:
+            result_columns = set(result.columns)
+            original_columns_set = set(self._colnames)
+            missing_columns = original_columns_set - result_columns
+
+            if missing_columns:
+                missing_list = ", ".join(f"'{col}'" for col in sorted(missing_columns))
+                original_list = ", ".join(f"'{col}'" for col in self._colnames)
+                raise MissingColumnsError(
+                    f"Query result missing required columns: {missing_list}. "
+                    f"The query must return all original table columns. "
+                    f"Original columns: {original_list}"
+                )
+
+        return result
+
     def get_data(self) -> pd.DataFrame:
         """
         Return the unfiltered data as a DataFrame.
@@ -283,6 +371,10 @@ def __init__(self, engine: Engine, table_name: str):
         if not inspector.has_table(table_name):
             raise ValueError(f"Table '{table_name}' not found in database")
 
+        # Store original column names for validation
+        columns_info = inspector.get_columns(table_name)
+        self._colnames = [col["name"] for col in columns_info]
+
     def get_db_type(self) -> str:
         """
         Get the database type.
@@ -441,10 +533,70 @@ def execute_query(self, query: str) -> pd.DataFrame:
         :
             Query results as pandas DataFrame
 
+        Raises
+        ------
+        UnsafeQueryError
+            If the query starts with a disallowed SQL operation
+
         """
+        check_query(query)
         with self._get_connection() as conn:
             return pd.read_sql_query(text(query), conn)
 
+    def test_query(
+        self, query: str, *, require_all_columns: bool = False
+    ) -> pd.DataFrame:
+        """
+        Test query by fetching only one row.
+
+        Parameters
+        ----------
+        query
+            SQL query to test
+        require_all_columns
+            If True, validates that result includes all original table columns
+
+        Returns
+        -------
+        :
+            Query results with at most one row
+
+        Raises
+        ------
+        UnsafeQueryError
+            If the query starts with a disallowed SQL operation
+        MissingColumnsError
+            If require_all_columns is True and result is missing required columns
+
+        """
+        check_query(query)
+        with self._get_connection() as conn:
+            # Use pandas read_sql_query with limit to get at most one row
+            limit_query = f"SELECT * FROM ({query}) AS subquery LIMIT 1"
+            try:
+                df = pd.read_sql_query(text(limit_query), conn)
+            except Exception:
+                # If LIMIT syntax doesn't work, fall back to regular read and take first row
+                df = pd.read_sql_query(text(query), conn).head(1)
+
+            if require_all_columns:
+                result_columns = set(df.columns)
+                original_columns_set = set(self._colnames)
+                missing_columns = original_columns_set - result_columns
+
+                if missing_columns:
+                    missing_list = ", ".join(
+                        f"'{col}'" for col in sorted(missing_columns)
+                    )
+                    original_list = ", ".join(f"'{col}'" for col in self._colnames)
+                    raise MissingColumnsError(
+                        f"Query result missing required columns: {missing_list}. "
+                        f"The query must return all original table columns. "
+                        f"Original columns: {original_list}"
+                    )
+
+            return df
+
     def get_data(self) -> pd.DataFrame:
         """
         Return the unfiltered data as a DataFrame.

pkg-py/src/querychat/_utils.py

@@ -1,12 +1,86 @@
 from __future__ import annotations
 
 import os
+import re
 import warnings
 from contextlib import contextmanager
 from typing import TYPE_CHECKING, Literal, Optional
 
 import narwhals.stable.v1 as nw
 
+
+class UnsafeQueryError(ValueError):
+    """Raised when a query contains an unsafe/write operation."""
+
+
+def check_query(query: str) -> None:
+    """
+    Check if a SQL query appears to be a non-read-only (write) operation.
+
+    Raises UnsafeQueryError if the query starts with a dangerous keyword.
+
+    Two categories of keywords are checked:
+
+    - Always blocked: DELETE, TRUNCATE, CREATE, DROP, ALTER, GRANT, REVOKE,
+      EXEC, EXECUTE, CALL
+    - Blocked unless QUERYCHAT_ENABLE_UPDATE_QUERIES=true: INSERT, UPDATE,
+      MERGE, REPLACE, UPSERT
+
+    Parameters
+    ----------
+    query
+        The SQL query string to check
+
+    Raises
+    ------
+    UnsafeQueryError
+        If the query starts with a disallowed keyword
+
+    """
+    # Normalize: newlines/tabs -> space, collapse multiple spaces, trim, uppercase
+    normalized = re.sub(r"[\r\n\t]+", " ", query)
+    normalized = re.sub(r" +", " ", normalized)
+    normalized = normalized.strip().upper()
+
+    # Always blocked - destructive/schema/admin operations
+    always_blocked = [
+        "DELETE",
+        "TRUNCATE",
+        "CREATE",
+        "DROP",
+        "ALTER",
+        "GRANT",
+        "REVOKE",
+        "EXEC",
+        "EXECUTE",
+        "CALL",
+    ]
+
+    # Blocked unless escape hatch enabled - data modification
+    update_keywords = ["INSERT", "UPDATE", "MERGE", "REPLACE", "UPSERT"]
+
+    # Check always-blocked keywords first
+    always_pattern = r"^(" + "|".join(always_blocked) + r")\b"
+    match = re.match(always_pattern, normalized)
+    if match:
+        raise UnsafeQueryError(
+            f"Query appears to contain a disallowed operation: {match.group(1)}. "
+            "Only SELECT queries are allowed."
+        )
+
+    # Check update keywords (can be enabled via envvar)
+    enable_updates = os.environ.get("QUERYCHAT_ENABLE_UPDATE_QUERIES", "").lower()
+    if enable_updates not in ("true", "1", "yes"):
+        update_pattern = r"^(" + "|".join(update_keywords) + r")\b"
+        match = re.match(update_pattern, normalized)
+        if match:
+            raise UnsafeQueryError(
+                f"Query appears to contain an update operation: {match.group(1)}. "
+                "Only SELECT queries are allowed. "
+                "Set QUERYCHAT_ENABLE_UPDATE_QUERIES=true to allow update queries."
+            )
+
+
 if TYPE_CHECKING:
     from narwhals.stable.v1.typing import IntoFrame
 

pkg-py/src/querychat/tools.py

@@ -75,7 +75,7 @@ def update_dashboard(query: str, title: str) -> ContentToolResult:
 
         try:
             # Test the query but don't execute it yet
-            data_source.execute_query(query)
+            data_source.test_query(query, require_all_columns=True)
 
             # Add Apply Filter button
             button_html = f"""<button

pkg-py/src/querychat/types/__init__.py

@@ -1,11 +1,19 @@
-from .._datasource import DataFrameSource, DataSource, SQLAlchemySource  # noqa: A005
+from .._datasource import (  # noqa: A005
+    DataFrameSource,
+    DataSource,
+    MissingColumnsError,
+    SQLAlchemySource,
+)
 from .._querychat_module import ServerValues
+from .._utils import UnsafeQueryError
 from ..tools import UpdateDashboardData
 
 __all__ = (
     "DataFrameSource",
     "DataSource",
+    "MissingColumnsError",
     "SQLAlchemySource",
     "ServerValues",
+    "UnsafeQueryError",
     "UpdateDashboardData",
 )