Update tools/sql_tool.py

tools/sql_tool.py

@@ -1,7 +1,12 @@
+# space/tools/sql_tool.py
 import os
 import re
+import json
+import shutil
+import glob
 import pandas as pd
 from typing import Optional
+
 from utils.config import AppConfig
 from utils.tracing import Tracer
 
@@ -21,23 +26,19 @@ class SQLTool:
                 raise RuntimeError("Missing GCP_SERVICE_ACCOUNT_JSON secret")
 
             # Accept full JSON string from Space Secret
-            if key_json.strip().startswith("{"):
-                import json
-                info = json.loads(key_json)
-            else:
-                info = {}
-
+            info = json.loads(key_json) if key_json.strip().startswith("{") else {}
             creds = service_account.Credentials.from_service_account_info(info)
             self.client = bigquery.Client(credentials=creds, project=cfg.gcp_project)
 
         elif self.backend == "motherduck":
-            import duckdb, shutil, glob
-            import os
+            import duckdb
 
+            # ---- Enforce supported DuckDB version for MotherDuck extension ----
             if not duckdb.__version__.startswith("1.3.2"):
                 raise RuntimeError(
-                "MotherDuck currently supports DuckDB 1.3.2. "
-                "Pin duckdb==1.3.2 in requirements.txt and redeploy."
+                    f"Incompatible DuckDB version {duckdb.__version__}. "
+                    "MotherDuck currently supports DuckDB 1.3.2. "
+                    "Pin duckdb==1.3.2 in requirements.txt and redeploy."
                 )
 
             token = self.cfg.motherduck_token or os.getenv("MOTHERDUCK_TOKEN")
@@ -45,17 +46,18 @@ class SQLTool:
             if not token:
                 raise RuntimeError("Missing MOTHERDUCK_TOKEN")
 
+            # ---- Clean stale extension caches compiled for other DuckDB versions ----
             try:
                 ext_root = os.path.expanduser("~/.duckdb/extensions")
                 for p in glob.glob(os.path.join(ext_root, "*")):
-                    if "1.3.2" not in p:  # keep current version caches, remove others
+                    if "1.3.2" not in p:  # keep only current version caches
                         shutil.rmtree(p, ignore_errors=True)
             except Exception:
+                # best-effort cleanup; proceed even if it fails
+                pass
 
-            # Plain DuckDB connection
-            self.client = duckdb.connect()
-
-            # Ensure MotherDuck extension is available and loaded
+            # ---- Connect & load MotherDuck extension ----
+            self.client = duckdb.connect()  # in-memory connection; we'll ATTACH MotherDuck
             self.client.execute("INSTALL motherduck;")
             self.client.execute("LOAD motherduck;")
 
@@ -70,18 +72,21 @@ class SQLTool:
         """
         Minimal NL2SQL heuristic; replace with your own mapping or LLM prompt.
         Expect users to include table names. Example:
-        "avg metric by month from analytics.events"
+          "avg metric by month from analytics.events"
         """
         m = message.lower()
 
-        # Very basic template example (edit to your tables/columns)
+        # Very basic template example (edit table/columns to your schema)
         if "avg" in m and " by " in m:
+            # DuckDB uses DATE_TRUNC('month', col); BigQuery uses DATE_TRUNC(col, MONTH).
+            # This generic SQL should work in DuckDB/MotherDuck; adapt if using BigQuery.
             return (
                 "-- Example template; edit me\n"
                 "SELECT DATE_TRUNC('month', date_col) AS month, "
                 "AVG(metric) AS avg_metric "
                 "FROM analytics.table "
-                "GROUP BY 1 ORDER BY 1;"
+                "GROUP BY 1 "
+                "ORDER BY 1;"
             )
 
         # Pass-through if the user typed SQL explicitly
@@ -93,12 +98,15 @@ class SQLTool:
 
     def run(self, message: str) -> pd.DataFrame:
         sql = self._nl_to_sql(message)
-        self.tracer.trace_event("sql_query", {"sql": sql, "backend": self.backend})
+        try:
+            self.tracer.trace_event("sql_query", {"sql": sql, "backend": self.backend})
+        except Exception:
+            pass
 
         if self.backend == "bigquery":
             df = self.client.query(sql).to_dataframe()
         else:
-            # DuckDB (MotherDuck): fetch_df returns a pandas DataFrame
+            # DuckDB (MotherDuck)
             df = self.client.execute(sql).fetch_df()
 
         return df