From 13db07f336acdf7d4ff346014edd4958a7e9b433 Mon Sep 17 00:00:00 2001
From: ChihYu Yeh <chihyu.jimmy.yeh@gmail.com>
Date: Sun, 30 Jun 2024 22:59:25 +0800
Subject: [PATCH] add wren-engine as new engine and remove force_deploy in
 Makefile

---
 wren-ai-service/Makefile                     |  3 +--
 wren-ai-service/src/providers/engine/wren.py | 20 ++++++++++++++++++++
 2 files changed, 21 insertions(+), 2 deletions(-)

diff --git a/wren-ai-service/Makefile b/wren-ai-service/Makefile
index fbaf2a4183..b62a918365 100644
--- a/wren-ai-service/Makefile
+++ b/wren-ai-service/Makefile
@@ -12,8 +12,7 @@ dev-down:
 
 ## wren-ai-service related ##
 start:
-	poetry run python -m src.__main__ & \
-	make force_deploy
+	poetry run python -m src.__main__
 
 force_deploy:
 	while ! nc -z localhost 5556; do \
diff --git a/wren-ai-service/src/providers/engine/wren.py b/wren-ai-service/src/providers/engine/wren.py
index 499eae67d4..718559f096 100644
--- a/wren-ai-service/src/providers/engine/wren.py
+++ b/wren-ai-service/src/providers/engine/wren.py
@@ -71,3 +71,23 @@ async def dry_run_sql(
             res = await response.text()
 
             return False, res
+
+
+@provider("wren-engine")
+class WrenEngine(Engine):
+    def __init__(self, endpoint: str = os.getenv("WREN_ENGINE_ENDPOINT")):
+        self._endpoint = endpoint
+
+    async def dry_run_sql(
+        self,
+        sql: str,
+        session: aiohttp.ClientSession,
+    ) -> Tuple[bool, Optional[Dict[str, Any]]]:
+        async with session.get(
+            f"{self._endpoint}/v1/mdl/dry-run",
+            json={"sql": remove_limit_statement(add_quotes(sql)), "limit": 1},
+        ) as response:
+            if response.status == 200:
+                return True, None
+            res = await response.text()
+            return False, res