[DH-5756] Raise an error when sql is not valid (#476)

Dataherald · Apr 26, 2024 · fbd96ea · fbd96ea
1 parent 66688d1
commit fbd96ea
Showing 1 changed file with 10 additions and 1 deletion.
diff --git a/dataherald/api/types/requests.py b/dataherald/api/types/requests.py
@@ -1,4 +1,5 @@
-from pydantic import BaseModel
+from pydantic import BaseModel, validator
+from sql_metadata import Parser
 
 from dataherald.types import LLMConfig
 
@@ -17,6 +18,14 @@ class SQLGenerationRequest(BaseModel):
     sql: str | None
     metadata: dict | None
 
+    @validator("sql")
+    def validate_model_name(cls, v: str | None):
+        try:
+            Parser(v).tables  # noqa: B018
+        except Exception as e:
+            raise ValueError(f"SQL {v} is malformed. Please check the syntax.") from e
+        return v
+
 
 class StreamSQLGenerationRequest(BaseModel):
     finetuning_id: str | None