update

wren-ai-service/Justfile

@@ -67,5 +67,5 @@ prepare-wren-engine:
 use-wren-ui-as-engine:
 	poetry run python -m src.force_update_config
 
-test-sql mdl_path="./tests/data/usecases/hubspot/mdl.json" dataset_type="bigquery":
-	poetry run python demo/test_sql.py --mdl-path "{{mdl_path}}" --dataset-type "{{dataset_type}}"
+test-sql mdl_path="" dataset_type="bigquery" dataset="":
+	poetry run python demo/test_sql.py --mdl-path "{{mdl_path}}" --dataset-type "{{dataset_type}}" --dataset "{{dataset}}"

wren-ai-service/demo/test_sql.py

@@ -1,7 +1,7 @@
 import argparse
 import json
 
-from utils import get_data_from_wren_engine
+from utils import get_data_from_wren_engine, rerun_wren_engine
 
 
 def main():
@@ -10,7 +10,10 @@ def main():
     )
 
     parser.add_argument(
-        "--mdl-path", type=str, required=True, help="Path to MDL JSON file"
+        "--mdl-path",
+        type=str,
+        required=True,
+        help="Path to MDL JSON file",
     )
 
     parser.add_argument(
@@ -21,6 +24,14 @@ def main():
         help="Dataset type (default: bigquery)",
     )
 
+    parser.add_argument(
+        "--dataset",
+        type=str,
+        default="ecommerce",
+        choices=["ecommerce", "hr", ""],
+        help="Dataset (default: ecommerce)",
+    )
+
     args = parser.parse_args()
 
     # Load MDL JSON file
@@ -34,6 +45,8 @@ def main():
         print(f"Error: Invalid JSON in MDL file {args.mdl_path}")
         return
 
+    rerun_wren_engine(mdl_json, args.dataset_type, args.dataset)
+
     # Execute query
     print("Enter SQL query (end with semicolon on a new line to execute, 'q' to quit):")
     lines = []

wren-ai-service/demo/utils.py

@@ -64,7 +64,7 @@ def _update_wren_engine_configs(configs: list[dict]):
     assert response.status_code == 200
 
 
-def rerun_wren_engine(mdl_json: Dict, dataset_type: str, dataset: str):
+def rerun_wren_engine(mdl_json: Dict, dataset_type: str, dataset: Optional[str] = None):
     assert dataset_type in DATA_SOURCES
 
     SOURCE = dataset_type