From 36b75b1290237c5e38271df372a2c4b46ef82af5 Mon Sep 17 00:00:00 2001
From: Chris Nivera <chris.nivera@snowflake.com>
Date: Mon, 21 Oct 2024 15:55:10 -0700
Subject: [PATCH] checkbox for onboarding

---
 journeys/iteration.py | 30 ++++++++++++++++++++++--------
 1 file changed, 22 insertions(+), 8 deletions(-)

diff --git a/journeys/iteration.py b/journeys/iteration.py
index ad355320..d2160171 100644
--- a/journeys/iteration.py
+++ b/journeys/iteration.py
@@ -36,11 +36,6 @@
     yaml_to_semantic_model,
 )
 from semantic_model_generator.protos import semantic_model_pb2
-from semantic_model_generator.snowflake_utils.env_vars import (
-    SNOWFLAKE_ACCOUNT_LOCATOR,
-    SNOWFLAKE_HOST,
-    SNOWFLAKE_USER,
-)
 from semantic_model_generator.validate_model import validate
 
 
@@ -234,6 +229,11 @@ def edit_verified_query(
 
             elif st.session_state.get("successful_sql", False):
                 # Moved outside the `if run:` block to ensure it's always evaluated
+                mark_as_onboarding = st.checkbox(
+                    "Mark as onboarding question",
+                    key=f"edit_onboarding_idx_{message_index}",
+                    help="Mark this question as an onboarding verified query.",
+                )
                 save = st.button(
                     "Save as verified query",
                     use_container_width=True,
@@ -243,12 +243,18 @@ def edit_verified_query(
                     sql_no_analyst_comment = user_updated_sql.replace(
                         " /* Generated by Cortex Analyst */", ""
                     )
-                    add_verified_query(question, sql_no_analyst_comment)
+                    add_verified_query(
+                        question,
+                        sql_no_analyst_comment,
+                        is_onboarding_question=mark_as_onboarding,
+                    )
                     st.session_state["editing"] = False
                     st.session_state["confirmed_edits"] = True
 
 
-def add_verified_query(question: str, sql: str) -> None:
+def add_verified_query(
+    question: str, sql: str, is_onboarding_question: bool = False
+) -> None:
     """Save verified question and SQL into an in-memory list with additional details."""
     # Verified queries follow the Snowflake definitions.
     verified_query = semantic_model_pb2.VerifiedQuery(
@@ -257,6 +263,7 @@ def add_verified_query(question: str, sql: str) -> None:
         sql=sql,
         verified_by=st.session_state["user_name"],
         verified_at=int(time.time()),
+        use_as_onboarding_question=is_onboarding_question,
     )
     st.session_state.semantic_model.verified_queries.append(verified_query)
     st.success(
@@ -307,6 +314,11 @@ def display_content(
                 df = pd.read_sql(sql, conn)
                 st.dataframe(df, hide_index=True)
 
+                mark_as_onboarding = st.checkbox(
+                    "Mark as onboarding question",
+                    key=f"onboarding_idx_{message_index}",
+                    help="Mark this question as an onboarding verified query.",
+                )
                 left, right = st.columns(2)
                 if right.button(
                     "Save as verified query",
@@ -317,7 +329,9 @@ def display_content(
                     cleaned_sql = sql_no_cte.replace(
                         " /* Generated by Cortex Analyst */", ""
                     )
-                    add_verified_query(question, cleaned_sql)
+                    add_verified_query(
+                        question, cleaned_sql, is_onboarding_question=mark_as_onboarding
+                    )
 
                 if left.button(
                     "Edit",