Update tools/sql_tool.py

tools/sql_tool.py

@@ -2,11 +2,7 @@
 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
 
@@ -20,7 +16,6 @@ class SQLTool:
         if self.backend == "bigquery":
             from google.cloud import bigquery
             from google.oauth2 import service_account
-
             key_json = os.getenv("GCP_SERVICE_ACCOUNT_JSON")
             if not key_json:
                 raise RuntimeError("Missing GCP_SERVICE_ACCOUNT_JSON secret")
@@ -32,54 +27,42 @@ class SQLTool:
 
         elif self.backend == "motherduck":
             import duckdb
-
-            # ---- Enforce supported DuckDB version for MotherDuck extension ----
+            # MotherDuck currently supports DuckDB 1.3.2 broadly across hosts
             if not duckdb.__version__.startswith("1.3.2"):
                 raise RuntimeError(
                     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")
-            db_name = self.cfg.motherduck_db or "default"
+            db_name = (self.cfg.motherduck_db or "workspace").strip()
             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 only current version caches
-                        shutil.rmtree(p, ignore_errors=True)
-            except Exception:
-                # best-effort cleanup; proceed even if it fails
-                pass
-
-            # ---- 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;")
-
-            # Attach the remote MotherDuck database and use it
-            self.client.execute(f"SET motherduck_token='{token}';")
-            self.client.execute(f"ATTACH 'md:/{db_name}' AS md;")
-            self.client.execute("USE md;")  # subsequent queries run against 'md' by default
+            # Easiest, correct way: connect directly to MotherDuck database.
+            # This will auto-download/load the extension; no manual INSTALL/LOAD/ATTACH needed.
+            # Valid URIs include:
+            #   "md:"                  -> connects to workspace (all DBs)
+            #   f"md:{db_name}"        -> connects to a specific DB
+            #   f"md:{db_name}?motherduck_token=..." -> with token in URI
+            uri = f"md:{db_name}?motherduck_token={token}"
+            self.client = duckdb.connect(uri)
+
+            # Optional: set a default database context (USE) if you connected to 'md:' (workspace)
+            # if db_name in ("", "workspace"):
+            #     self.client.execute("USE your_database;")
         else:
             raise RuntimeError("Unknown SQL backend")
 
     def _nl_to_sql(self, message: str) -> str:
         """
         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"
+        Edit table/column names to your schema.
         """
         m = message.lower()
 
-        # Very basic template example (edit table/columns to your schema)
+        # Simple example (DuckDB/MotherDuck DATE_TRUNC flavor)
         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, "
@@ -93,7 +76,6 @@ class SQLTool:
         if re.match(r"^\s*select ", m):
             return message
 
-        # Fallback
         return "SELECT * FROM analytics.table LIMIT 100;"
 
     def run(self, message: str) -> pd.DataFrame:
@@ -104,9 +86,7 @@ class SQLTool:
             pass
 
         if self.backend == "bigquery":
-            df = self.client.query(sql).to_dataframe()
+            return self.client.query(sql).to_dataframe()
         else:
             # DuckDB (MotherDuck)
-            df = self.client.execute(sql).fetch_df()
-
-        return df
+            return self.client.execute(sql).fetch_df()