Use sqlglot for better checks (#24)

Author: aniket-s-kulkarni

pyproject.toml

@@ -18,7 +18,7 @@ dependencies = [
     "langchain-ollama>=0.2.3",
     "langchain-openai>=0.3.7",
     "langgraph>=0.3.12",
-    "mcp[cli]>=1.3.0",
+    "mcp>=1.3.0",
     "openai>=1.65.3",
     "pandas>=2.2.3",
     "prompt-toolkit>=3.0.50",
@@ -29,6 +29,7 @@ dependencies = [
     "pyyaml>=6.0.2",
     "requests>=2.32.3",
     "rich>=13.9.4",
+    "sqlglot>=26.23.0",
     "structlog>=25.1.0",
     "typer>=0.15.2",
     "uvicorn>=0.34.0",

src/dremioai/api/dremio/sql.py

@@ -42,7 +42,7 @@ class ArcticSource(BaseModel):
 
 class Query(BaseModel):
     sql: str = Field(..., alias="sql")
-    context: Optional[Any] = None
+    context: Optional[List[str]] = None
     references: Optional[Dict[str, ArcticSource]] = None
 
 

src/dremioai/config/settings.py

@@ -122,6 +122,7 @@ class Dremio(BaseModel):
     project_id: Optional[str] = None
     enable_experimental: Optional[bool] = False  # enable experimental tools
     oauth2: Optional[OAuth2] = None
+    allow_dml: Optional[bool] = False
     model_config = ConfigDict(validate_assignment=True)
 
     @field_serializer("raw_pat")

src/dremioai/tools/tools.py

@@ -45,6 +45,8 @@
 from dremioai.api.dremio.catalog import get_schema, get_lineage, get_descriptions
 from csv import reader
 from io import StringIO
+from sqlglot import parse_one
+from sqlglot import expressions
 
 logger = log.logger(__name__)
 
@@ -244,6 +246,31 @@ async def invoke(self) -> Dict[str, Any]:
 
 class RunSqlQuery(Tools):
     For: ClassVar[Annotated[ToolType, ToolType.FOR_SELF | ToolType.FOR_DATA_PATTERNS]]
+    _safe = [
+        expressions.Select,
+        expressions.With,
+        expressions.Union,
+    ]
+
+    @staticmethod
+    def ensure_query_allowed(s: str):
+        if settings.instance().dremio.allow_dml:
+            return
+
+        try:
+            q = parse_one(s)
+            if any(isinstance(q, t) for t in RunSqlQuery._safe):
+                return
+        except:
+            if not re.search(
+                r"\b(drop|insert|update|truncate|delete|copy into|alter|create)\b",
+                s,
+                re.IGNORECASE,
+            ):
+                return
+        raise ValueError(
+            "The query contains a DML statement. Only select queries are allowed"
+        )
 
     async def invoke(self, s: str) -> Dict[str, List[Any]]:
         """Run a SELECT sql query on the Dremio cluster and return the results.
@@ -253,11 +280,7 @@ async def invoke(self, s: str) -> Dict[str, List[Any]]:
         Args:
             s: sql query
         """
-        # TODO: graduate to a more sophisticated SQL parser and check to allow better queries
-        if re.search(r"(drop|insert|update|truncate|delete)", s, re.IGNORECASE):
-            raise ValueError(
-                "The query contains a DML statement. Only select queries are allowed"
-            )
+        RunSqlQuery.ensure_query_allowed(s)
         try:
             s = f"/* dremioai: submitter={self.__class__.__name__} */\n{s}"
             df = await sql.run_query(query=s, use_df=True)
@@ -476,6 +499,8 @@ async def invoke(self) -> Dict[str, Any]:
             "jvm_gc_pause_seconds": "Indicates how long the JVM was paused for garbage collection, and also is a rubric to know if the system is in use",
             "memory_heap_usage": "Indicates the amount of memory used by the JVM",
             "memory_heap_committed": "Indicates the amount of memory committed by the JVM",
+            "dremio_engine_executors": "Number of executors running in the Dremio engine. It correlates to dremio_engine_replica_running using engine_id label",
+            "dremio_engine_replica_running": "Number of running replicas in the Dremio engine. It correlates to dremio_engine_executors using engine_id label",
         }