add more usecases and add error string for add_quotes

Canner · Feb 20, 2025 · 466eed2 · 466eed2
1 parent 455b54e
commit 466eed2
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 6 deletions.
diff --git a/wren-ai-service/demo/utils.py b/wren-ai-service/demo/utils.py
@@ -35,7 +35,9 @@ def add_quotes(sql: str) -> Tuple[str, bool]:
     try:
         quoted_sql = sqlglot.transpile(sql, read="trino", identify=True)[0]
         return quoted_sql, True
-    except Exception:
+    except Exception as e:
+        print(f"Error in adding quotes to SQL: {sql}")
+        print(f"Error: {e}")
         return sql, False
 
 

diff --git a/wren-ai-service/eval/utils.py b/wren-ai-service/eval/utils.py
@@ -24,7 +24,9 @@ def add_quotes(sql: str) -> Tuple[str, bool]:
     try:
         quoted_sql = sqlglot.transpile(sql, read="trino", identify=True)[0]
         return quoted_sql, True
-    except Exception:
+    except Exception as e:
+        print(f"Error in adding quotes to SQL: {sql}")
+        print(f"Error: {e}")
         return sql, False
 
 

diff --git a/wren-ai-service/tests/pytest/test_usecases.py b/wren-ai-service/tests/pytest/test_usecases.py
@@ -161,21 +161,22 @@ def str_presenter(dumper, data):
         "stripe": "bigquery",
         "ecommerce": "duckdb",
         "hr": "duckdb",
+        "facebook_marketing": "bigquery",
+        "google_ads": "bigquery",
     }
+    usecases = list(usecase_to_dataset_type.keys())
 
     parser = argparse.ArgumentParser()
     parser.add_argument(
         "--usecases",
         type=str,
         nargs="+",
         default=["all"],
-        choices=["all", "hubspot", "ga4", "ecommerce", "hr", "woocommerce", "stripe"],
+        choices=["all"] + usecases,
     )
     args = parser.parse_args()
 
-    if "all" in args.usecases:
-        usecases = ["hubspot", "ga4", "ecommerce", "hr", "woocommerce", "stripe"]
-    else:
+    if "all" not in args.usecases:
         usecases = args.usecases
 
     url = "http://localhost:5556"