Kschmaus port eval logic

app.py

@@ -8,12 +8,12 @@
 from app_utils.shared_utils import (  # noqa: E402
     GeneratorAppScreen,
     get_snowflake_connection,
-    set_sit_query_tag,
     set_account_name,
     set_host_name,
-    set_user_name,
-    set_streamlit_location,
+    set_sit_query_tag,
     set_snowpark_session,
+    set_streamlit_location,
+    set_user_name,
 )
 from semantic_model_generator.snowflake_utils.env_vars import (  # noqa: E402
     SNOWFLAKE_ACCOUNT_LOCATOR,
@@ -22,7 +22,7 @@
 )
 
 
-@st.experimental_dialog(title="Connection Error")
+@st.dialog(title="Connection Error")
 def failed_connection_popup() -> None:
     """
     Renders a dialog box detailing that the credentials provided could not be used to connect to Snowflake.

app_utils/chat.py

@@ -1,6 +1,6 @@
 import json
 import re
-from typing import Dict, Any
+from typing import Any, Dict
 
 import requests
 import streamlit as st
@@ -9,7 +9,7 @@
 API_ENDPOINT = "https://{HOST}/api/v2/cortex/analyst/message"
 
 
-@st.cache_data(ttl=60, show_spinner=False)
+@st.cache_data(ttl=300, show_spinner=False)
 def send_message(
     _conn: SnowflakeConnection, semantic_model: str, messages: list[dict[str, str]]
 ) -> Dict[str, Any]:

app_utils/shared_utils.py

@@ -2,20 +2,19 @@
 
 import json
 import os
-import time
 import tempfile
+import time
 from dataclasses import dataclass
-from datetime import datetime
 from enum import Enum
 from io import StringIO
-from typing import Any, Optional, List, Union, Dict, Tuple
+from typing import Any, Dict, List, Optional, Tuple, Union
 
 import pandas as pd
 import streamlit as st
-from snowflake.snowpark import Session
 from PIL import Image
 from snowflake.connector import ProgrammingError
 from snowflake.connector.connection import SnowflakeConnection
+from snowflake.snowpark import Session
 
 from semantic_model_generator.data_processing.proto_utils import (
     proto_to_yaml,
@@ -27,22 +26,18 @@
 )
 from semantic_model_generator.protos import semantic_model_pb2
 from semantic_model_generator.protos.semantic_model_pb2 import Dimension, Table
+from semantic_model_generator.snowflake_utils.env_vars import (  # noqa: E402
+    assert_required_env_vars,
+)
 from semantic_model_generator.snowflake_utils.snowflake_connector import (
     SnowflakeConnector,
     fetch_databases,
     fetch_schemas_in_database,
+    fetch_stages_in_schema,
+    fetch_table_schema,
     fetch_tables_views_in_schema,
     fetch_warehouses,
-    fetch_stages_in_schema,
     fetch_yaml_names_in_stage,
-    fetch_columns_names_in_table,
-)
-
-from semantic_model_generator.snowflake_utils.env_vars import (  # noqa: E402
-    SNOWFLAKE_ACCOUNT_LOCATOR,
-    SNOWFLAKE_HOST,
-    SNOWFLAKE_USER,
-    assert_required_env_vars,
 )
 
 SNOWFLAKE_ACCOUNT = os.environ.get("SNOWFLAKE_ACCOUNT_LOCATOR", "")
@@ -79,7 +74,7 @@ def set_streamlit_location() -> bool:
     return sis
 
 
-@st.experimental_dialog(title="Setup")
+@st.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 connection setup.
@@ -105,6 +100,7 @@ def get_snowflake_connection() -> SnowflakeConnection:
     if st.session_state["sis"]:
         # Import SiS-required modules
         import sys
+
         from snowflake.snowpark.context import get_active_session
 
         # Non-Anaconda supported packages must be added to path to import from stage
@@ -127,7 +123,7 @@ def get_snowflake_connection() -> SnowflakeConnection:
                 return get_connector().open_connection(db_name="")
 
 
-@st.cache_resource(show_spinner=False)
+# @st.cache_resource(show_spinner=False)
 def set_snowpark_session(_conn: Optional[SnowflakeConnection] = None) -> None:
     """
     Creates a snowpark for python session.
@@ -203,22 +199,18 @@ def get_available_stages(schema: str) -> List[str]:
 
 
 @st.cache_resource(show_spinner=False)
-def validate_table_columns(table: str, columns_must_exist: Tuple[str,...]) -> bool:
-    """
-    Fetches the available stages from the Snowflake account.
-
-    Returns:
-        List[str]: A list of available stages.
-    """
-    columns_names = fetch_columns_names_in_table(get_snowflake_connection(), table)
-    for col in columns_must_exist:
-        if col not in columns_names:
+def validate_table_schema(table: str, schema: Dict[str, str]) -> bool:
+    table_schema = fetch_table_schema(get_snowflake_connection(), table)
+    if set(schema) != set(table_schema):
+        return False
+    for col_name, col_type in table_schema.items():
+        if not (schema[col_name] in col_type):
             return False
     return True
 
 
 @st.cache_resource(show_spinner=False)
-def validate_table_exist(schema: str, table_name:str) -> bool:
+def validate_table_exist(schema: str, table_name: str) -> bool:
     """
     Validate table exist in the Snowflake account.
 
@@ -522,7 +514,7 @@ def init_session_states() -> None:
         st.session_state.confirmed_edits = False
 
 
-@st.experimental_dialog("Edit Dimension")  # type: ignore[misc]
+@st.dialog("Edit Dimension")  # type: ignore[misc]
 def edit_dimension(table_name: str, dim: semantic_model_pb2.Dimension) -> None:
     """
     Renders a dialog box to edit an existing dimension.
@@ -572,7 +564,7 @@ def edit_dimension(table_name: str, dim: semantic_model_pb2.Dimension) -> None:
         st.rerun()
 
 
-@st.experimental_dialog("Add Dimension")  # type: ignore[misc]
+@st.dialog("Add Dimension")  # type: ignore[misc]
 def add_dimension(table: semantic_model_pb2.Table) -> None:
     """
     Renders a dialog box to add a new dimension.
@@ -611,7 +603,7 @@ def add_dimension(table: semantic_model_pb2.Table) -> None:
         st.rerun()
 
 
-@st.experimental_dialog("Edit Measure")  # type: ignore[misc]
+@st.dialog("Edit Measure")  # type: ignore[misc]
 def edit_measure(table_name: str, measure: semantic_model_pb2.Measure) -> None:
     """
     Renders a dialog box to edit an existing measure.
@@ -684,7 +676,7 @@ def edit_measure(table_name: str, measure: semantic_model_pb2.Measure) -> None:
         st.rerun()
 
 
-@st.experimental_dialog("Add Measure")  # type: ignore[misc]
+@st.dialog("Add Measure")  # type: ignore[misc]
 def add_measure(table: semantic_model_pb2.Table) -> None:
     """
     Renders a dialog box to add a new measure.
@@ -744,7 +736,7 @@ def add_measure(table: semantic_model_pb2.Table) -> None:
         st.rerun()
 
 
-@st.experimental_dialog("Edit Time Dimension")  # type: ignore[misc]
+@st.dialog("Edit Time Dimension")  # type: ignore[misc]
 def edit_time_dimension(
     table_name: str, tdim: semantic_model_pb2.TimeDimension
 ) -> None:
@@ -789,7 +781,7 @@ def edit_time_dimension(
         st.rerun()
 
 
-@st.experimental_dialog("Add Time Dimension")  # type: ignore[misc]
+@st.dialog("Add Time Dimension")  # type: ignore[misc]
 def add_time_dimension(table: semantic_model_pb2.Table) -> None:
     """
     Renders a dialog box to add a new time dimension.
@@ -987,7 +979,7 @@ def display_table(table_name: str) -> None:
         add_time_dimension(table)
 
 
-@st.experimental_dialog("Add Table")  # type: ignore[misc]
+@st.dialog("Add Table")  # type: ignore[misc]
 def add_new_table() -> None:
     """
     Renders a dialog box to add a new logical table.
@@ -1104,7 +1096,7 @@ def import_yaml() -> None:
             st.rerun()
 
 
-@st.experimental_dialog("Model YAML", width="large")  # type: ignore
+@st.dialog("Model YAML", width="large")  # type: ignore
 def show_yaml_in_dialog() -> None:
     yaml = proto_to_yaml(st.session_state.semantic_model)
     st.code(
@@ -1414,17 +1406,3 @@ def to_dict(self) -> dict[str, str]:
             "Schema": self.stage_schema,
             "Stage": self.stage_name,
         }
-
-
-@dataclass
-class SnowflakeTable:
-    table_database: str
-    table_schema: str
-    table_name: str
-
-    def to_dict(self) -> dict[str, str]:
-        return {
-            "Database": self.table_database,
-            "Schema": self.table_schema,
-            "Table": self.table_name,
-        }

journeys/builder.py

@@ -68,7 +68,7 @@ def update_tables() -> None:
     st.session_state["selected_tables"] = valid_selected_tables
 
 
-@st.experimental_dialog("Selecting your tables", width="large")
+@st.dialog("Selecting your tables", width="large")
 def table_selector_dialog() -> None:
     st.write(
         "Please fill out the following fields to start building your semantic model."