Update tools/sql_tool.py

tools/sql_tool.py

@@ -7,12 +7,16 @@ from utils.config import AppConfig
 from utils.tracing import Tracer
 
 
+RESERVED_MD_WORKSPACE_NAMES = {"", "workspace", "default"}  # treat these as workspace/no-DB context
+
+
 class SQLTool:
     def __init__(self, cfg: AppConfig, tracer: Tracer):
         self.cfg = cfg
         self.tracer = tracer
         self.backend = cfg.sql_backend  # "bigquery" or "motherduck"
 
+        # ---------------- BIGQUERY BACKEND ----------------
         if self.backend == "bigquery":
             from google.cloud import bigquery
             from google.oauth2 import service_account
@@ -25,6 +29,7 @@ class SQLTool:
             creds = service_account.Credentials.from_service_account_info(info)
             self.client = bigquery.Client(credentials=creds, project=cfg.gcp_project)
 
+        # ---------------- MOTHERDUCK BACKEND ----------------
         elif self.backend == "motherduck":
             import duckdb
 
@@ -41,22 +46,18 @@ class SQLTool:
             if not token:
                 raise RuntimeError("Missing MOTHERDUCK_TOKEN")
 
-            # Primary path: connect directly to the database
-            # Correct formats: "md:" (workspace) or "md:<dbname>" (specific DB)
-            try:
-                uri = f"md:{db_name}?motherduck_token={token}" if db_name and db_name != "workspace" else f"md:?motherduck_token={token}"
-                self.client = duckdb.connect(uri)
-                # If we connected to workspace explicitly, set DB context if provided
-                if db_name and db_name != "workspace":
-                    # Ensure we are actually in the right DB context
-                    self._ensure_db_context(db_name, allow_create)
-            except Exception as e:
-                # Fallback: connect to workspace, then create/use DB if needed
+            # Workspace vs concrete DB handling
+            if db_name in RESERVED_MD_WORKSPACE_NAMES:
+                # Connect to workspace; caller should fully-qualify tables if needed
                 self.client = duckdb.connect(f"md:?motherduck_token={token}")
-                if not db_name or db_name == "workspace":
-                    # Using workspace only, caller must fully-qualify schema.table in queries
-                    pass
-                else:
+                # No USE/CREATE in workspace mode
+            else:
+                # Try direct connection to the DB (preferred)
+                try:
+                    self.client = duckdb.connect(f"md:{db_name}?motherduck_token={token}")
+                except Exception:
+                    # Fallback: connect to workspace, then USE/CREATE the DB if permitted
+                    self.client = duckdb.connect(f"md:?motherduck_token={token}")
                     self._ensure_db_context(db_name, allow_create)
 
         else:
@@ -66,35 +67,40 @@ class SQLTool:
     def _ensure_db_context(self, db_name: str, allow_create: bool):
         """
         Try to USE the target DB; if it doesn't exist and allow_create=True, create it and USE it.
+        Skips reserved workspace names.
         """
-        # DuckDB/MotherDuck: USE <db_name>;
+        if db_name in RESERVED_MD_WORKSPACE_NAMES:
+            # No-op for workspace/default
+            return
+
+        # Attempt USE first
         try:
             self.client.execute(f"USE {self._quote_ident(db_name)};")
+            return
         except Exception as use_err:
             if not allow_create:
                 raise RuntimeError(
                     f"Database '{db_name}' not found and ALLOW_CREATE_DB is false. "
                     f"Original error: {use_err}"
                 )
-            # Attempt to create then USE
-            try:
-                self.client.execute(f"CREATE DATABASE {self._quote_ident(db_name)};")
-                self.client.execute(f"USE {self._quote_ident(db_name)};")
-            except Exception as create_err:
-                raise RuntimeError(
-                    f"Could not create or use database '{db_name}'. "
-                    f"Original errors: USE: {use_err} | CREATE: {create_err}"
-                )
+
+        # Attempt CREATE then USE
+        try:
+            # CREATE DATABASE <name>; is supported on MotherDuck for valid names (not 'default')
+            self.client.execute(f"CREATE DATABASE {self._quote_ident(db_name)};")
+            self.client.execute(f"USE {self._quote_ident(db_name)};")
+        except Exception as create_err:
+            raise RuntimeError(
+                f"Could not create or use database '{db_name}'. "
+                f"Original errors: CREATE: {create_err}"
+            )
 
     @staticmethod
     def _quote_ident(name: str) -> str:
         """
-        Very light identifier quoting. Adjust if you allow special chars.
+        Very light identifier quoting. Replace non [a-zA-Z0-9_] with underscore.
         """
-        if not name:
-            return name
-        # basic guard; you can tighten rules for your org naming conventions
-        safe = re.sub(r"[^a-zA-Z0-9_]", "_", name)
+        safe = re.sub(r"[^a-zA-Z0-9_]", "_", (name or ""))
         return safe
 
     # ----- NL → SQL heuristic (toy example; edit to your schema) -----
@@ -113,7 +119,7 @@ class SQLTool:
             )
 
         # If user typed SQL already, run it as-is
-        if re.match(r"^\s*select ", m):
+        if re.match(r"^\\s*select ", m):
             return message
 
         # Fallback