Decouple the Rewriter and Connector in Ibis Server

ibis-server/README.md

@@ -78,5 +78,17 @@ Run the server
 just run
 ```
 
+### Run the testing
+- Preapre the Wren Engine server (see [Wren Engine Example](../example/README.md)
+- Run pytest with `WREN_ENGINE_ENDPOINT` env. (The default port of Wren Engine is 8080)
+  - run all test
+  ```
+  WREN_ENGINE_ENDPOINT=http://localhost:8080 just test
+  ```
+  - run specific test using [pytest args](https://docs.pytest.org/en/latest/example/markers.html#using-k-expr-to-select-tests-based-on-their-name)
+  ```
+  WREN_ENGINE_ENDPOINT=http://localhost:8080 just test -k 'postgres'
+  ```
+
 ## Contributing
 Please see [CONTRIBUTING.md](docs/CONTRIBUTING.md) for more information.

ibis-server/app/model/connector.py

@@ -1,6 +1,5 @@
 import pandas as pd
 
-from app.mdl.rewriter import Rewriter
 from app.model import ConnectionInfo
 from app.model.data_source import DataSource
 
@@ -17,19 +16,11 @@ def __init__(
         self.manifest_str = manifest_str
 
     def query(self, sql: str, limit: int) -> pd.DataFrame:
-        rewritten_sql = Rewriter(self.manifest_str, self.data_source).rewrite(sql)
-        return (
-            self.connection.sql(
-                rewritten_sql,
-            )
-            .limit(limit)
-            .to_pandas()
-        )
+        return self.connection.sql(sql).limit(limit).to_pandas()
 
     def dry_run(self, sql: str) -> None:
         try:
-            rewritten_sql = Rewriter(self.manifest_str, self.data_source).rewrite(sql)
-            self.connection.sql(rewritten_sql)
+            self.connection.sql(sql)
         except Exception as e:
             raise QueryDryRunError(f"Exception: {type(e)}, message: {e!s}")
 

ibis-server/app/model/validator.py

@@ -1,13 +1,15 @@
 from __future__ import annotations
 
+from app.mdl.rewriter import Rewriter
 from app.model.connector import Connector
 
 rules = ["column_is_valid"]
 
 
 class Validator:
-    def __init__(self, connector: Connector):
+    def __init__(self, connector: Connector, rewriter: Rewriter):
         self.connector = connector
+        self.rewriter = rewriter
 
     def validate(self, rule: str, parameters: dict[str, str]):
         if rule not in rules:
@@ -28,9 +30,9 @@ def _validate_column_is_valid(self, parameters: dict[str, str]):
             raise MissingRequiredParameterError("columnName")
 
         try:
-            self.connector.dry_run(
-                f'SELECT "{column_name}" FROM "{model_name}" LIMIT 1'
-            )
+            sql = f'SELECT "{column_name}" FROM "{model_name}" LIMIT 1'
+            rewritten_sql = self.rewriter.rewrite(sql)
+            self.connector.dry_run(rewritten_sql)
         except Exception as e:
             raise ValidationError(f"Exception: {type(e)}, message: {e!s}")
 

ibis-server/app/routers/v2/connector.py

@@ -29,19 +29,23 @@ def query(
     dry_run: Annotated[bool, Query(alias="dryRun")] = False,
     limit: int | None = None,
 ) -> Response:
+    rewritten_sql = Rewriter(dto.manifest_str, data_source).rewrite(dto.sql)
     connector = Connector(data_source, dto.connection_info, dto.manifest_str)
     if dry_run:
-        connector.dry_run(dto.sql)
+        connector.dry_run(rewritten_sql)
         return Response(status_code=204)
     return JSONResponse(
-        to_json(connector.query(dto.sql, limit=limit), dto.column_dtypes)
+        to_json(connector.query(rewritten_sql, limit=limit), dto.column_dtypes)
     )
 
 
 @router.post("/{data_source}/validate/{rule_name}")
 @log_dto
 def validate(data_source: DataSource, rule_name: str, dto: ValidateDTO) -> Response:
-    validator = Validator(Connector(data_source, dto.connection_info, dto.manifest_str))
+    validator = Validator(
+        Connector(data_source, dto.connection_info, dto.manifest_str),
+        Rewriter(dto.manifest_str, data_source),
+    )
     validator.validate(rule_name, dto.parameters)
     return Response(status_code=204)
 

ibis-server/justfile

@@ -31,8 +31,9 @@ run:
 dev:
     poetry run fastapi dev --port {{port}}
 
-test:
-    poetry run pytest
+# run the pytest with given args
+test *ARGS:
+    poetry run pytest {{ ARGS }}
 
 docker-build:
 # alias for `docker-build`