From 466eed2b0e754778fce464f3da9ef43dfaaac657 Mon Sep 17 00:00:00 2001
From: ChihYu Yeh <chihyu.jimmy.yeh@gmail.com>
Date: Fri, 21 Feb 2025 03:06:56 +0800
Subject: [PATCH] add more usecases and add error string for add_quotes

---
 wren-ai-service/demo/utils.py                 | 4 +++-
 wren-ai-service/eval/utils.py                 | 4 +++-
 wren-ai-service/tests/pytest/test_usecases.py | 9 +++++----
 3 files changed, 11 insertions(+), 6 deletions(-)

diff --git a/wren-ai-service/demo/utils.py b/wren-ai-service/demo/utils.py
index cc44b3d0f..194064283 100644
--- 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
index 3cc4a5085..13189a3f9 100644
--- 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
index 1ae3b3554..03a9e78f6 100644
--- a/wren-ai-service/tests/pytest/test_usecases.py
+++ b/wren-ai-service/tests/pytest/test_usecases.py
@@ -161,7 +161,10 @@ 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(
@@ -169,13 +172,11 @@ def str_presenter(dumper, data):
         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"