Address copilot review

Author: jairus-m

src/dbt_mcp/dbt_admin/run_artifacts/extractors.py

@@ -409,7 +409,7 @@ def extract_from_catalog(data: dict[str, Any]) -> dict[str, list[tuple]]:
                     col.get("name") or col_name,
                     col.get("index"),
                     col.get("type"),
-                    col.get("comment") or "",
+                    col.get("comment"),
                 )
             )
 

src/dbt_mcp/dbt_admin/run_artifacts/store.py

@@ -5,6 +5,7 @@
 """
 
 import logging
+import re
 import time
 from typing import Any
 
@@ -19,6 +20,7 @@
 from dbt_mcp.errors.artifact_search import (
     ArtifactNotLoadedError,
     ArtifactQueryError,
+    ArtifactSearchError,
     ArtifactValidationError,
 )
 
@@ -56,8 +58,16 @@ class ArtifactStore:
 
     def __init__(self) -> None:
         self.conn = duckdb.connect()
-        self.conn.execute("INSTALL fts;")
-        self.conn.execute("LOAD fts;")
+        try:
+            self.conn.execute("LOAD fts;")
+        except duckdb.Error:
+            try:
+                self.conn.execute("INSTALL fts;")
+                self.conn.execute("LOAD fts;")
+            except duckdb.Error as e:
+                raise ArtifactSearchError(
+                    f"Failed to load the DuckDB FTS extension: {e}"
+                ) from e
         self._loaded_tables: set[str] = set()
         self._tables_created: bool = False
         self._pending_index_tables: set[str] = set()
@@ -277,7 +287,7 @@ def _build_indexes(self, config: TableConfig) -> None:
             )
 
         for col in config.index_columns:
-            idx_name = f"idx_{config.table_name[:4]}_{col}"
+            idx_name = f"idx_{config.table_name}_{col}"
             self.conn.execute(
                 f"CREATE INDEX IF NOT EXISTS {idx_name} "
                 f'ON {config.table_name}("{col}");'
@@ -327,7 +337,10 @@ def describe_table(self, table_name: str) -> list[dict[str, str]]:
 
     def query(self, sql: str) -> list[dict[str, Any]]:
         """Execute a read-only SQL query. Results capped at 500 rows."""
-        tokens = sql.strip().upper().split()
+        sanitized = _strip_sql_comments(sql)
+        if ";" in sanitized:
+            raise ArtifactQueryError("Multi-statement queries are not allowed.")
+        tokens = sanitized.strip().upper().split()
         for token in tokens:
             if token in READONLY_BLOCKED:
                 raise ArtifactQueryError(
@@ -391,6 +404,13 @@ def _validate_table_name(self, table_name: str) -> None:
             )
 
 
+def _strip_sql_comments(sql: str) -> str:
+    """Remove SQL block (/* */) and line (--) comments."""
+    sql = re.sub(r"/\*.*?\*/", " ", sql, flags=re.DOTALL)
+    sql = re.sub(r"--[^\n]*", " ", sql)
+    return sql
+
+
 def _serialize(val: Any) -> Any:
     """Ensure values are JSON-serializable."""
     if val is None or isinstance(val, (str, int, float, bool)):

tests/unit/dbt_admin/run_artifacts/test_store.py

@@ -133,12 +133,14 @@ def test_mutating_keywords_are_blocked(self, store: ArtifactStore) -> None:
             with pytest.raises(ArtifactQueryError, match="Blocked keyword"):
                 store.query(f"{keyword} something")
 
-    def test_mutating_keywords_blocked_anywhere_in_query(
-        self, store: ArtifactStore
-    ) -> None:
-        with pytest.raises(ArtifactQueryError, match="Blocked keyword"):
+    def test_multi_statement_query_blocked(self, store: ArtifactStore) -> None:
+        with pytest.raises(ArtifactQueryError, match="Multi-statement"):
             store.query("SELECT 1; DROP TABLE nodes")
 
+    def test_comment_bypass_blocked(self, store: ArtifactStore) -> None:
+        with pytest.raises(ArtifactQueryError, match="Multi-statement"):
+            store.query("SELECT 1;/**/DROP TABLE nodes")
+
     def test_invalid_sql_raises(self, store: ArtifactStore) -> None:
         with pytest.raises(ArtifactQueryError, match="Query failed"):
             store.query("SELECT * FROM table_that_does_not_exist_xyz")