Decouple the Rewriter and Connector in Ibis Server (#692)

* decouple the connector and rewriter

* enhace the test script

* move testing doc to development.md

* update doc

* Update development.md

Author: goldmedal

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/docs/development.md

@@ -37,9 +37,6 @@ To start the server:
 - Execute `just dev` to start the server in development mode (auto-reloads on code changes)
 - The default port is `8000`. You can change it by running `just port=8001 run` or `just port=8001 dev`
 
-To run the tests:
-- Execute `just test`
-
 ### Environment Variables
 - `WREN_ENGINE_ENDPOINT`: The endpoint of the Wren Java engine
 - `LOG_LEVEL`: The log level of the server (default is INFO)
@@ -48,6 +45,13 @@ To run the tests:
 - Build the image: `just docker-build`
 - Run the container: `just docker-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 specific data source test using [pytest marker](https://docs.pytest.org/en/stable/example/markers.html). There are some markers for different data sources. See the list in [pyproject.toml](../pyproject.toml).
+  ```
+  WREN_ENGINE_ENDPOINT=http://localhost:8080 just test 'postgres'
+  ```
 
 ## How to add new data source
 Please see [How to Add a New Data Source](how-to-add-data-source.md) for more information.

ibis-server/justfile

@@ -31,8 +31,9 @@ run:
 dev:
     poetry run fastapi dev --port {{port}}
 
-test:
-    poetry run pytest
+# run the pytest tests for the given marker
+test MARKER:
+    poetry run pytest -m {{ MARKER }}
 
 docker-build:
 # alias for `docker-build`