Skip to content

Commit

Permalink
add wren-engine as new engine and remove force_deploy in Makefile
Browse files Browse the repository at this point in the history
  • Loading branch information
cyyeh committed Jul 1, 2024
1 parent de803c7 commit 13db07f
Show file tree
Hide file tree
Showing 2 changed files with 21 additions and 2 deletions.
3 changes: 1 addition & 2 deletions wren-ai-service/Makefile
Original file line number Diff line number Diff line change
Expand Up @@ -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 \
Expand Down
20 changes: 20 additions & 0 deletions wren-ai-service/src/providers/engine/wren.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

0 comments on commit 13db07f

Please sign in to comment.