Dropdown selectors for iteration flow (#121)

Replaces the free text entry fields in the iteration flow with
dropdowns, so that the user can click to select their file instead of
having to type everything in manually. This fully completes #94 .

## Testing
Go through the iteration flow and verify functionality of the dropdown
selectors. Proceed to use the app and upload your YAML to a stage.

Additionally, go through the builder flow, and verify that the step
where you select the Snowflake stage at the end works with the new
dropdown selectors.

admin_apps/journeys/iteration.py

@@ -6,10 +6,11 @@
 import requests
 import sqlglot
 import streamlit as st
-from snowflake.connector import SnowflakeConnection
+from snowflake.connector import ProgrammingError, SnowflakeConnection
 from streamlit.delta_generator import DeltaGenerator
 from streamlit_monaco import st_monaco
 
+from admin_apps.journeys.builder import get_available_databases, get_available_schemas
 from admin_apps.shared_utils import (
     GeneratorAppScreen,
     SnowflakeStage,
@@ -38,6 +39,8 @@
     SNOWFLAKE_USER,
 )
 from semantic_model_generator.snowflake_utils.snowflake_connector import (
+    fetch_stages_in_schema,
+    fetch_yaml_names_in_stage,
     set_database,
     set_schema,
 )
@@ -368,19 +371,25 @@ def upload_handler(file_name: str) -> None:
     else:
         # If coming from the builder flow, we need to ask the user for the exact stage path to upload to.
         st.markdown("Please enter the destination of your YAML file.")
-        with st.form("upload_form"):
-            stage_database = st.text_input("Stage database", value="")
-            stage_schema = st.text_input("Stage schema", value="")
-            stage_name = st.text_input("Stage name", value="")
-            new_name = st.text_input("File name (omit .yaml suffix)", value="")
+        stage_selector_container()
+        new_name = st.text_input("File name (omit .yaml suffix)", value="")
+
+        if st.button("Submit Upload"):
+            if (
+                not st.session_state["selected_iteration_database"]
+                or not st.session_state["selected_iteration_schema"]
+                or not st.session_state["selected_iteration_stage"]
+                or not new_name
+            ):
+                st.error("Please fill in all fields.")
+                return
 
-            if st.form_submit_button("Submit Upload"):
-                st.session_state["snowflake_stage"] = SnowflakeStage(
-                    stage_database=stage_database,
-                    stage_schema=stage_schema,
-                    stage_name=stage_name,
-                )
-                upload_handler(new_name)
+            st.session_state["snowflake_stage"] = SnowflakeStage(
+                stage_database=st.session_state["selected_iteration_database"],
+                stage_schema=st.session_state["selected_iteration_schema"],
+                stage_name=st.session_state["selected_iteration_stage"],
+            )
+            upload_handler(new_name)
 
 
 def update_container(
@@ -480,33 +489,125 @@ def yaml_editor(yaml_str: str) -> None:
         update_container(status_container, "editing", prefix=status_container_title)
 
 
+@st.cache_resource(show_spinner=False)
+def get_available_stages(schema: str) -> List[str]:
+    """
+    Fetches the available stages from the Snowflake account.
+
+    Returns:
+        List[str]: A list of available stages.
+    """
+    return fetch_stages_in_schema(get_snowflake_connection(), schema)
+
+
+@st.cache_resource(show_spinner=False)
+def get_yamls_from_stage(stage: str) -> List[str]:
+    """
+    Fetches the YAML files from the specified stage.
+
+    Args:
+        stage (str): The name of the stage to fetch the YAML files from.
+
+    Returns:
+        List[str]: A list of YAML files in the specified stage.
+    """
+    return fetch_yaml_names_in_stage(get_snowflake_connection(), stage)
+
+
+def stage_selector_container() -> None:
+    """
+    Common component that encapsulates db/schema/stage selection for the admin app.
+    When a db/schema/stage is selected, it is saved to the session state for reading elsewhere.
+    Returns: None
+    """
+    available_schemas = []
+    available_stages = []
+
+    # First, retrieve all databases that the user has access to.
+    stage_database = st.selectbox(
+        "Stage database",
+        options=get_available_databases(),
+        index=None,
+        key="selected_iteration_database",
+    )
+    if stage_database:
+        # When a valid database is selected, fetch the available schemas in that database.
+        try:
+            set_database(get_snowflake_connection(), stage_database)
+            available_schemas = get_available_schemas(stage_database)
+        except (ValueError, ProgrammingError):
+            st.error("Insufficient permissions to read from the selected database.")
+            st.stop()
+
+    stage_schema = st.selectbox(
+        "Stage schema",
+        options=available_schemas,
+        index=None,
+        key="selected_iteration_schema",
+    )
+    if stage_schema:
+        # When a valid schema is selected, fetch the available stages in that schema.
+        try:
+            set_schema(get_snowflake_connection(), stage_schema)
+            available_stages = get_available_stages(stage_schema)
+        except (ValueError, ProgrammingError):
+            st.error("Insufficient permissions to read from the selected schema.")
+            st.stop()
+
+    st.selectbox(
+        "Stage name",
+        options=available_stages,
+        index=None,
+        key="selected_iteration_stage",
+    )
+
+
 @st.experimental_dialog("Welcome to the Iteration app! 💬", width="large")
 def set_up_requirements() -> None:
     """
     Collects existing YAML location from the user so that we can download it.
     """
-    # 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(
-            "Fill in the Snowflake stage details to download your existing YAML file."
-        )
-        # 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="")
-        file_name = st.text_input("File name", value="<your_file>.yaml")
-        if st.form_submit_button("Submit"):
-            st.session_state["snowflake_stage"] = SnowflakeStage(
-                stage_database=stage_database,
-                stage_schema=stage_schema,
-                stage_name=stage_name,
+    st.markdown(
+        "Fill in the Snowflake stage details to download your existing YAML file."
+    )
+
+    stage_selector_container()
+
+    # Based on the currently selected stage, show a dropdown of YAML files for the user to pick from.
+    available_files = []
+    if (
+        "selected_iteration_stage" in st.session_state
+        and st.session_state["selected_iteration_stage"]
+    ):
+        # When a valid stage is selected, fetch the available YAML files in that stage.
+        try:
+            available_files = get_yamls_from_stage(
+                st.session_state["selected_iteration_stage"]
             )
-            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()
+        except (ValueError, ProgrammingError):
+            st.error("Insufficient permissions to read from the selected stage.")
+            st.stop()
+
+    file_name = st.selectbox("File name", options=available_files, index=None)
+
+    if st.button(
+        "Submit",
+        disabled=not st.session_state["selected_iteration_database"]
+        or not st.session_state["selected_iteration_schema"]
+        or not st.session_state["selected_iteration_stage"]
+        or not file_name,
+    ):
+        st.session_state["snowflake_stage"] = SnowflakeStage(
+            stage_database=st.session_state["selected_iteration_database"],
+            stage_schema=st.session_state["selected_iteration_schema"],
+            stage_name=st.session_state["selected_iteration_stage"],
+        )
+        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()
 
 
 SAVE_HELP = """Save changes to the active semantic model in this app. This is

semantic_model_generator/snowflake_utils/snowflake_connector.py

@@ -305,6 +305,43 @@ def fetch_tables_views_in_schema(
     return results
 
 
+def fetch_stages_in_schema(conn: SnowflakeConnection, schema_name: str) -> list[str]:
+    """
+    Fetches all stages that the current user has access to in the current schema
+    Args:
+        conn: SnowflakeConnection to run the query
+        schema_name: The name of the schema to connect to.
+
+    Returns: a list of fully qualified stage names
+    """
+
+    query = f"show stages in schema {schema_name};"
+    cursor = conn.cursor()
+    cursor.execute(query)
+    stages = cursor.fetchall()
+
+    return [f"{result[2]}.{result[3]}.{result[1]}" for result in stages]
+
+
+def fetch_yaml_names_in_stage(conn: SnowflakeConnection, stage_name: str) -> list[str]:
+    """
+    Fetches all yaml files that the current user has access to in the current stage
+    Args:
+        conn: SnowflakeConnection to run the query
+        stage_name: The fully qualified name of the stage to connect to.
+
+    Returns: a list of yaml file names
+    """
+
+    query = f"list @{stage_name} pattern='.*\\.yaml';"
+    cursor = conn.cursor()
+    cursor.execute(query)
+    yaml_files = cursor.fetchall()
+
+    # The file name is prefixed with "@{stage_name}/", so we need to remove that prefix.
+    return [result[0].split("/")[-1] for result in yaml_files]
+
+
 def get_valid_schemas_tables_columns_df(
     conn: SnowflakeConnection,
     table_schema: Optional[str] = None,