refactor(rewriter): integrate by experiment parameter

ibis-server/app/mdl/rewriter.py

@@ -1,5 +1,4 @@
 import importlib
-from abc import ABC, abstractmethod
 
 import httpx
 import orjson
@@ -16,26 +15,36 @@
 importlib.import_module("ibis.backends.sql.dialects")
 
 
-class Rewriter(ABC):
-    def __init__(self, manifest_str: str, data_source: DataSource = None):
+class Rewriter:
+    def __init__(
+        self,
+        manifest_str: str,
+        data_source: DataSource = None,
+        experiment=False,
+    ):
         self.manifest_str = manifest_str
         self.data_source = data_source
+        if experiment:
+            self._rewriter = EmbeddedEngineRewriter(manifest_str)
+        else:
+            self._rewriter = ExternalEngineRewriter(manifest_str)
 
-    @abstractmethod
     def rewrite(self, sql: str) -> str:
-        pass
+        planned_sql = self._rewriter.rewrite(sql)
+        logger.debug("Planned SQL: {}", planned_sql)
+        dialect_sql = self._transpile(planned_sql) if self.data_source else planned_sql
+        logger.debug("Dialect SQL: {}", dialect_sql)
+        return dialect_sql
 
-    def transpile(self, planned_sql: str) -> str:
-        dialect_sql = sqlglot.transpile(
+    def _transpile(self, planned_sql: str) -> str:
+        return sqlglot.transpile(
             planned_sql, read="trino", write=self.data_source.name
         )[0]
-        logger.debug("Dialect SQL: {}", dialect_sql)
-        return dialect_sql
 
 
-class ExternalEngineRewriter(Rewriter):
-    def __init__(self, manifest_str: str, data_source: DataSource = None):
-        super().__init__(manifest_str, data_source)
+class ExternalEngineRewriter:
+    def __init__(self, manifest_str: str):
+        self.manifest_str = manifest_str
 
     def rewrite(self, sql: str) -> str:
         try:
@@ -48,30 +57,22 @@ def rewrite(self, sql: str) -> str:
                 },
                 content=orjson.dumps({"manifestStr": self.manifest_str, "sql": sql}),
             )
-            planned_sql = r.raise_for_status().text
-            logger.debug("Planned SQL: {}", planned_sql)
-            return (
-                planned_sql if self.data_source is None else self.transpile(planned_sql)
-            )
+            return r.raise_for_status().text
         except httpx.ConnectError as e:
             raise ConnectionError(f"Can not connect to Wren Engine: {e}")
         except httpx.HTTPStatusError as e:
             raise RewriteError(e.response.text)
 
 
-class EmbeddedEngineRewriter(Rewriter):
-    def __init__(self, manifest_str: str, data_source: DataSource = None):
-        super().__init__(manifest_str, data_source)
+class EmbeddedEngineRewriter:
+    def __init__(self, manifest_str: str):
+        self.manifest_str = manifest_str
 
     def rewrite(self, sql: str) -> str:
         from wren_core import transform_sql
 
         try:
-            planned_sql = transform_sql(self.manifest_str, sql)
-            logger.debug("Planned SQL: {}", planned_sql)
-            return (
-                planned_sql if self.data_source is None else self.transpile(planned_sql)
-            )
+            return transform_sql(self.manifest_str, sql)
         except Exception as e:
             raise RewriteError(str(e))
 

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

@@ -4,7 +4,7 @@
 from fastapi.responses import JSONResponse
 
 from app.dependencies import verify_query_dto
-from app.mdl.rewriter import ExternalEngineRewriter
+from app.mdl.rewriter import Rewriter
 from app.model import (
     DryPlanDTO,
     QueryDTO,
@@ -27,9 +27,7 @@ def query(
     dry_run: Annotated[bool, Query(alias="dryRun")] = False,
     limit: int | None = None,
 ) -> Response:
-    rewritten_sql = ExternalEngineRewriter(dto.manifest_str, data_source).rewrite(
-        dto.sql
-    )
+    rewritten_sql = Rewriter(dto.manifest_str, data_source=data_source).rewrite(dto.sql)
     connector = Connector(data_source, dto.connection_info, dto.manifest_str)
     if dry_run:
         connector.dry_run(rewritten_sql)
@@ -41,7 +39,7 @@ def query(
 def validate(data_source: DataSource, rule_name: str, dto: ValidateDTO) -> Response:
     validator = Validator(
         Connector(data_source, dto.connection_info, dto.manifest_str),
-        ExternalEngineRewriter(dto.manifest_str, data_source),
+        Rewriter(dto.manifest_str, data_source=data_source),
     )
     validator.validate(rule_name, dto.parameters)
     return Response(status_code=204)
@@ -61,9 +59,9 @@ def get_constraints(data_source: DataSource, dto: MetadataDTO) -> list[Constrain
 
 @router.post("/dry-plan")
 def dry_plan(dto: DryPlanDTO) -> str:
-    return ExternalEngineRewriter(dto.manifest_str).rewrite(dto.sql)
+    return Rewriter(dto.manifest_str).rewrite(dto.sql)
 
 
 @router.post("/{data_source}/dry-plan")
 def dry_plan_for_data_source(data_source: DataSource, dto: DryPlanDTO) -> str:
-    return ExternalEngineRewriter(dto.manifest_str, data_source).rewrite(dto.sql)
+    return Rewriter(dto.manifest_str, data_source=data_source).rewrite(dto.sql)

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

@@ -4,7 +4,7 @@
 from fastapi.responses import JSONResponse
 
 from app.dependencies import verify_query_dto
-from app.mdl.rewriter import EmbeddedEngineRewriter
+from app.mdl.rewriter import Rewriter
 from app.model import (
     DryPlanDTO,
     QueryDTO,
@@ -25,9 +25,9 @@ def query(
     dry_run: Annotated[bool, Query(alias="dryRun")] = False,
     limit: int | None = None,
 ) -> Response:
-    rewritten_sql = EmbeddedEngineRewriter(dto.manifest_str, data_source).rewrite(
-        dto.sql
-    )
+    rewritten_sql = Rewriter(
+        dto.manifest_str, data_source=data_source, experiment=True
+    ).rewrite(dto.sql)
     connector = Connector(data_source, dto.connection_info, dto.manifest_str)
     if dry_run:
         connector.dry_run(rewritten_sql)
@@ -37,19 +37,21 @@ def query(
 
 @router.post("/dry-plan")
 def dry_plan(dto: DryPlanDTO) -> str:
-    return EmbeddedEngineRewriter(dto.manifest_str).rewrite(dto.sql)
+    return Rewriter(dto.manifest_str, experiment=True).rewrite(dto.sql)
 
 
 @router.post("/{data_source}/dry-plan")
 def dry_plan_for_data_source(data_source: DataSource, dto: DryPlanDTO) -> str:
-    return EmbeddedEngineRewriter(dto.manifest_str, data_source).rewrite(dto.sql)
+    return Rewriter(dto.manifest_str, data_source=data_source, experiment=True).rewrite(
+        dto.sql
+    )
 
 
 @router.post("/{data_source}/validate/{rule_name}")
 def validate(data_source: DataSource, rule_name: str, dto: ValidateDTO) -> Response:
     validator = Validator(
         Connector(data_source, dto.connection_info, dto.manifest_str),
-        EmbeddedEngineRewriter(dto.manifest_str, data_source),
+        Rewriter(dto.manifest_str, data_source=data_source, experiment=True),
     )
     validator.validate(rule_name, dto.parameters)
     return Response(status_code=204)