Assert environment variables upon starting

admin_apps/app.py

@@ -1,13 +1,41 @@
-import os
-
 import streamlit as st
 
 from admin_apps.shared_utils import GeneratorAppScreen
+from semantic_model_generator.snowflake_utils.env_vars import (
+    SNOWFLAKE_ACCOUNT_LOCATOR,
+    SNOWFLAKE_HOST,
+    SNOWFLAKE_USER,
+    assert_required_env_vars,
+)
 
 # set_page_config must be run as the first Streamlit command on the page, before any other streamlit imports.
 st.set_page_config(layout="wide", page_icon="💬", page_title="Semantic Model Generator")
 
 
+@st.experimental_dialog(title="Setup")
+def env_setup_popup(missing_env_vars: list[str]) -> None:
+    """
+    Renders a dialog box to prompt the user to set the required environment variables.
+    Args:
+        missing_env_vars: A list of missing environment variables.
+    """
+    formatted_missing_env_vars = "\n".join(f"- **{s}**" for s in missing_env_vars)
+    st.markdown(
+        f"""Oops! It looks like the following required environment variables are missing: \n{formatted_missing_env_vars}\n\n
+Please follow the [setup instructions](https://github.com/Snowflake-Labs/semantic-model-generator?tab=readme-ov-file#setup) to properly configure your environment. Restart this app after you've set the required environment variables."""
+    )
+    st.stop()
+
+
+def verify_environment_setup() -> None:
+    """
+    Ensures that the correct environment variables are set before proceeding.
+    """
+    missing_env_vars = assert_required_env_vars()
+    if missing_env_vars:
+        env_setup_popup(missing_env_vars)
+
+
 if __name__ == "__main__":
     from admin_apps.journeys import builder, iteration
 
@@ -43,10 +71,12 @@ def onboarding_dialog() -> None:
             ):
                 iteration.show()
 
+    verify_environment_setup()
+
     # Populating common state between builder and iteration apps.
-    st.session_state["account_name"] = os.environ.get("SNOWFLAKE_ACCOUNT_LOCATOR")
-    st.session_state["host_name"] = os.environ.get("SNOWFLAKE_HOST")
-    st.session_state["user_name"] = os.environ.get("SNOWFLAKE_USER")
+    st.session_state["account_name"] = SNOWFLAKE_ACCOUNT_LOCATOR
+    st.session_state["host_name"] = SNOWFLAKE_HOST
+    st.session_state["user_name"] = SNOWFLAKE_USER
 
     # When the app first loads, show the onboarding screen.
     if "page" not in st.session_state:

admin_apps/journeys/iteration.py

@@ -1,5 +1,4 @@
 import json
-import os
 import time
 from typing import Any, Dict, List, Optional
 
@@ -33,6 +32,11 @@
     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.snowflake_utils.snowflake_connector import (
     SnowflakeConnector,
 )
@@ -487,13 +491,9 @@ def set_up_requirements() -> None:
     # Otherwise, we should collect the prebuilt YAML location from the user so that we can download it.
     with st.form("download_yaml_requirements"):
         st.markdown(
-            "Before we get started, let's make sure we have everything set up. If you'd like to populate these values by default, please follow the [environment variable setup instructions](https://github.com/Snowflake-Labs/semantic-model-generator/blob/main/README.md#setup)."
-        )
-        account_name = st.text_input(
-            "Account", value=os.environ.get("SNOWFLAKE_ACCOUNT_LOCATOR")
+            "Fill in the Snowflake stage details to download your existing YAML file."
         )
-        host_name = st.text_input("Host", value=os.environ.get("SNOWFLAKE_HOST"))
-        user_name = st.text_input("User", value=os.environ.get("SNOWFLAKE_USER"))
+        # TODO: Make these dropdown selectors by fetching all dbs/schemas similar to table approach?
         stage_database = st.text_input("Stage database", value="")
         stage_schema = st.text_input("Stage schema", value="")
         stage_name = st.text_input("Stage name", value="")
@@ -504,9 +504,9 @@ def set_up_requirements() -> None:
                 stage_schema=stage_schema,
                 stage_name=stage_name,
             )
-            st.session_state["account_name"] = account_name
-            st.session_state["host_name"] = host_name
-            st.session_state["user_name"] = user_name
+            st.session_state["account_name"] = SNOWFLAKE_ACCOUNT_LOCATOR
+            st.session_state["host_name"] = SNOWFLAKE_HOST
+            st.session_state["user_name"] = SNOWFLAKE_USER
             st.session_state["file_name"] = file_name
             st.session_state["page"] = GeneratorAppScreen.ITERATION
             st.rerun()

semantic_model_generator/snowflake_utils/env_vars.py

@@ -5,5 +5,36 @@
 SNOWFLAKE_WAREHOUSE = os.getenv("SNOWFLAKE_WAREHOUSE")
 SNOWFLAKE_USER = os.getenv("SNOWFLAKE_USER")
 SNOWFLAKE_PASSWORD = os.getenv("SNOWFLAKE_PASSWORD")
-SNOWFLAKE_HOST = os.getenv("SNOWFLAKE_HOST")
+SNOWFLAKE_HOST = os.getenv("SNOWFLAKE_HOST")  # optional per README docs
 SNOWFLAKE_AUTHENTICATOR = os.getenv("SNOWFLAKE_AUTHENTICATOR")
+SNOWFLAKE_ACCOUNT_LOCATOR = os.getenv("SNOWFLAKE_ACCOUNT_LOCATOR")
+
+
+def assert_required_env_vars() -> list[str]:
+    """
+    Ensures that the required environment variables are set before proceeding.
+    Returns: list of missing required environment variables
+
+    """
+
+    missing_env_vars = []
+    if not SNOWFLAKE_ROLE:
+        missing_env_vars.append("SNOWFLAKE_ROLE")
+    if not SNOWFLAKE_WAREHOUSE:
+        missing_env_vars.append("SNOWFLAKE_WAREHOUSE")
+    if not SNOWFLAKE_USER:
+        missing_env_vars.append("SNOWFLAKE_USER")
+    if not SNOWFLAKE_ACCOUNT_LOCATOR:
+        missing_env_vars.append("SNOWFLAKE_ACCOUNT_LOCATOR")
+    if not SNOWFLAKE_PASSWORD and not SNOWFLAKE_AUTHENTICATOR:
+        missing_env_vars.append("SNOWFLAKE_PASSWORD/SNOWFLAKE_AUTHENTICATOR")
+
+    # Assert that SNOWFLAKE_PASSWORD is required unless the user is using the externalbrowser authenticator
+    if (
+        SNOWFLAKE_AUTHENTICATOR
+        and SNOWFLAKE_AUTHENTICATOR.lower() != "externalbrowser"
+        and not SNOWFLAKE_PASSWORD
+    ):
+        missing_env_vars.append("SNOWFLAKE_PASSWORD")
+
+    return missing_env_vars