UI for editing join keys

admin_apps/journeys/builder.py

@@ -121,6 +121,7 @@ def table_selector_dialog() -> None:
         help="Checking this box will enable generation of experimental features in the semantic model. If enabling this setting, please ensure that you have the proper parameters set on your Snowflake account. Some features (e.g. joins) are currently in Private Preview and available only to select accounts. Reach out to your account team for access.",
     )
     allow_joins = False
+    st.session_state["experimental_features"] = False
     if experimental_features:
         allow_joins = True
         st.session_state["experimental_features"] = True

admin_apps/journeys/iteration.py

@@ -8,8 +8,10 @@
 import streamlit as st
 from snowflake.connector import ProgrammingError, SnowflakeConnection
 from streamlit.delta_generator import DeltaGenerator
+from streamlit_extras.row import row
 from streamlit_monaco import st_monaco
 
+from admin_apps.journeys.joins import joins_dialog
 from admin_apps.shared_utils import (
     GeneratorAppScreen,
     SnowflakeStage,
@@ -447,67 +449,75 @@ def yaml_editor(yaml_str: str) -> None:
         language="yaml",
     )
 
-    button_container = st.container()
+    button_container = row(5, vertical_align="center")
     status_container_title = "**Edit**"
     status_container = st.empty()
 
-    with button_container:
-        (one, two, three, four) = st.columns(4)
-        if one.button("Validate", use_container_width=True, help=VALIDATE_HELP):
-            # Validate new content
-            try:
-                validate(
-                    content,
-                    snowflake_account=st.session_state.account_name,
-                    conn=get_snowflake_connection(),
-                )
-                st.session_state["validated"] = True
-                update_container(
-                    status_container, "success", prefix=status_container_title
-                )
-                st.session_state.semantic_model = yaml_to_semantic_model(content)
-                st.session_state.last_saved_yaml = content
-            except Exception as e:
-                st.session_state["validated"] = False
-                update_container(
-                    status_container, "failed", prefix=status_container_title
-                )
-                exception_as_dialog(e)
-
-            # Rerun the app if validation was successful.
-            # We shouldn't rerun if validation failed as the error popup would immediately dismiss.
-            # This must be done outside of the try/except because the generic Exception handling is catching the
-            # exception that st.rerun() properly raises to halt execution.
-            # This is fixed in later versions of Streamlit, but other refactors to the code are required to upgrade.
-            if st.session_state["validated"]:
-                st.rerun()
-
-        if content:
-            two.download_button(
-                label="Download",
-                data=content,
-                file_name="semantic_model.yaml",
-                mime="text/yaml",
-                use_container_width=True,
-                help=DOWNLOAD_HELP,
+    def validate_and_update_session_state() -> None:
+        # Validate new content
+        try:
+            validate(
+                content,
+                snowflake_account=st.session_state.account_name,
+                conn=get_snowflake_connection(),
             )
+            st.session_state["validated"] = True
+            update_container(status_container, "success", prefix=status_container_title)
+            st.session_state.semantic_model = yaml_to_semantic_model(content)
+            st.session_state.last_saved_yaml = content
+        except Exception as e:
+            st.session_state["validated"] = False
+            update_container(status_container, "failed", prefix=status_container_title)
+            exception_as_dialog(e)
+
+    if button_container.button(
+        "Validate", use_container_width=True, help=VALIDATE_HELP
+    ):
+        validate_and_update_session_state()
+
+        # Rerun the app if validation was successful.
+        # We shouldn't rerun if validation failed as the error popup would immediately dismiss.
+        # This must be done outside of the try/except because the generic Exception handling is catching the
+        # exception that st.rerun() properly raises to halt execution.
+        # This is fixed in later versions of Streamlit, but other refactors to the code are required to upgrade.
+        if st.session_state["validated"]:
+            st.rerun()
+
+    if content:
+        button_container.download_button(
+            label="Download",
+            data=content,
+            file_name="semantic_model.yaml",
+            mime="text/yaml",
+            use_container_width=True,
+            help=DOWNLOAD_HELP,
+        )
 
-        if three.button(
-            "Upload",
+    if button_container.button(
+        "Upload",
+        use_container_width=True,
+        help=UPLOAD_HELP,
+    ):
+        upload_dialog(content)
+    if st.session_state.get("partner_setup", False):
+        from admin_apps.partner.partner_utils import integrate_partner_semantics
+
+        if button_container.button(
+            "Integrate Partner",
             use_container_width=True,
-            help=UPLOAD_HELP,
+            help=PARTNER_SEMANTIC_HELP,
+            disabled=not st.session_state["validated"],
         ):
-            upload_dialog(content)
-        if st.session_state.get("partner_setup", False):
-            from admin_apps.partner.partner_utils import integrate_partner_semantics
-
-            if four.button(
-                "Integrate Partner",
-                use_container_width=True,
-                help=PARTNER_SEMANTIC_HELP,
-                disabled=not st.session_state["validated"],
-            ):
-                integrate_partner_semantics()
+            integrate_partner_semantics()
+
+    if st.session_state.experimental_features:
+        if button_container.button(
+            "Join Editor",
+            use_container_width=True,
+        ):
+            with st.spinner("Validating your model..."):
+                validate_and_update_session_state()
+            joins_dialog()
 
     # Render the validation state (success=True, failed=False, editing=None) in the editor.
     if st.session_state.validated:
@@ -621,6 +631,11 @@ def set_up_requirements() -> None:
 
     file_name = st.selectbox("File name", options=available_files, index=None)
 
+    experimental_features = st.checkbox(
+        "Enable experimental features (optional)",
+        help="Checking this box will enable generation of experimental features in the semantic model. If enabling this setting, please ensure that you have the proper parameters set on your Snowflake account. Some features (e.g. joins) are currently in Private Preview and available only to select accounts. Reach out to your account team for access.",
+    )
+
     if st.button(
         "Submit",
         disabled=not st.session_state["selected_iteration_database"]
@@ -638,6 +653,7 @@ def set_up_requirements() -> None:
         st.session_state["user_name"] = SNOWFLAKE_USER
         st.session_state["file_name"] = file_name
         st.session_state["page"] = GeneratorAppScreen.ITERATION
+        st.session_state["experimental_features"] = experimental_features
         st.rerun()
 
 

admin_apps/journeys/joins.py

@@ -0,0 +1,220 @@
+import copy
+from typing import Optional
+
+import streamlit as st
+from streamlit_extras.row import row
+
+from semantic_model_generator.protos import semantic_model_pb2
+
+
+SUPPORTED_JOIN_TYPES = [
+    join_type
+    for join_type in semantic_model_pb2.JoinType.values()
+    if join_type != semantic_model_pb2.JoinType.join_type_unknown
+]
+SUPPORTED_RELATIONSHIP_TYPES = [
+    relationship_type
+    for relationship_type in semantic_model_pb2.RelationshipType.values()
+    if relationship_type
+    != semantic_model_pb2.RelationshipType.relationship_type_unknown
+]
+
+
+def relationship_builder(
+    relationship: semantic_model_pb2.Relationship, key: Optional[int] = 0
+) -> None:
+    """
+    Renders a UI for building/editing a semantic model relationship.
+    Args:
+        relationship: The relationship object to edit.
+
+    Returns:
+
+    """
+    with st.expander(
+        relationship.name or f"{relationship.left_table} ↔️ {relationship.right_table}",
+        expanded=True,
+    ):
+        relationship.name = st.text_input(
+            "Name", value=relationship.name, key=f"name_{key}"
+        )
+        # Logic to preselect the tables in the dropdown based on what's in the semantic model.
+        try:
+            default_left_table = [
+                table.name for table in st.session_state.semantic_model.tables
+            ].index(relationship.left_table)
+            default_right_table = [
+                table.name for table in st.session_state.semantic_model.tables
+            ].index(relationship.right_table)
+        except ValueError:
+            default_left_table = 0
+            default_right_table = 0
+        relationship.left_table = st.selectbox(
+            "Left Table",
+            options=[table.name for table in st.session_state.semantic_model.tables],
+            index=default_left_table,
+            key=f"left_table_{key}",
+        )
+
+        relationship.right_table = st.selectbox(
+            "Right Table",
+            options=[table.name for table in st.session_state.semantic_model.tables],
+            index=default_right_table,
+            key=f"right_table_{key}",
+        )
+
+        relationship.join_type = st.radio(  # type: ignore
+            "Join Type",
+            options=SUPPORTED_JOIN_TYPES,
+            format_func=lambda join_type: semantic_model_pb2.JoinType.Name(join_type),
+            index=SUPPORTED_JOIN_TYPES.index(relationship.join_type),
+            key=f"join_type_{key}",
+        )
+
+        relationship.relationship_type = st.radio(  # type: ignore
+            "Relationship Type",
+            options=SUPPORTED_RELATIONSHIP_TYPES,
+            format_func=lambda relationship_type: semantic_model_pb2.RelationshipType.Name(
+                relationship_type
+            ),
+            index=SUPPORTED_RELATIONSHIP_TYPES.index(relationship.relationship_type),
+            key=f"relationship_type_{key}",
+        )
+
+        st.divider()
+        # Builder section for the relationship's columns.
+        for col_idx, join_cols in enumerate(relationship.relationship_columns):
+            # Grabbing references to the exact Table objects that the relationship is pointing to.
+            # This allows us to pull the columns.
+            left_table_object = next(
+                (
+                    table
+                    for table in st.session_state.semantic_model.tables
+                    if table.name == relationship.left_table
+                )
+            )
+            right_table_object = next(
+                (
+                    table
+                    for table in st.session_state.semantic_model.tables
+                    if table.name == relationship.right_table
+                )
+            )
+
+            try:
+                left_columns = []
+                left_columns.extend(left_table_object.columns)
+                left_columns.extend(left_table_object.dimensions)
+                left_columns.extend(left_table_object.time_dimensions)
+                left_columns.extend(left_table_object.measures)
+
+                right_columns = []
+                right_columns.extend(right_table_object.columns)
+                right_columns.extend(right_table_object.dimensions)
+                right_columns.extend(right_table_object.time_dimensions)
+                right_columns.extend(right_table_object.measures)
+
+                default_left_col = [col.name for col in left_columns].index(
+                    join_cols.left_column
+                )
+                default_right_col = [col.name for col in right_columns].index(
+                    join_cols.right_column
+                )
+            except ValueError:
+                default_left_col = 0
+                default_right_col = 0
+
+            join_cols.left_column = st.selectbox(
+                "Left Column",
+                options=[col.name for col in left_columns],
+                index=default_left_col,
+                key=f"left_col_{key}_{col_idx}",
+            )
+            join_cols.right_column = st.selectbox(
+                "Right Column",
+                options=[col.name for col in right_columns],
+                index=default_right_col,
+                key=f"right_col_{key}_{col_idx}",
+            )
+
+            if st.button("Delete join key", key=f"delete_join_key_{key}_{col_idx}"):
+                relationship.relationship_columns.pop(col_idx)
+                st.rerun(scope="fragment")
+
+            st.divider()
+
+        join_editor_row = row(2, vertical_align="center")
+        if join_editor_row.button(
+            "Add new join key",
+            key=f"add_join_keys_{key}",
+            use_container_width=True,
+            type="primary",
+        ):
+            relationship.relationship_columns.append(
+                semantic_model_pb2.RelationKey(
+                    left_column="",
+                    right_column="",
+                )
+            )
+            st.rerun(scope="fragment")
+
+        if join_editor_row.button(
+            "🗑️ Delete join path",
+            key=f"delete_join_path_{key}",
+            use_container_width=True,
+        ):
+            st.session_state.builder_joins.pop(key)
+            st.rerun(scope="fragment")
+
+
+@st.dialog("Join Builder", width="large")
+def joins_dialog() -> None:
+
+    if "builder_joins" not in st.session_state:
+        # Making a copy of the original relationships list so we can modify freely without affecting the original.
+        st.session_state.builder_joins = st.session_state.semantic_model.relationships[
+            :
+        ]
+
+    for idx, relationship in enumerate(st.session_state.builder_joins):
+        relationship_builder(relationship, idx)
+
+    # If the user clicks "Add join", add a new join to the relationships list
+    if st.button("Add new join path", use_container_width=True):
+        st.session_state.builder_joins.append(
+            semantic_model_pb2.Relationship(
+                left_table="",
+                right_table="",
+                join_type=semantic_model_pb2.JoinType.inner,
+                relationship_type=semantic_model_pb2.RelationshipType.one_to_one,
+                relationship_columns=[],
+            )
+        )
+        st.rerun(scope="fragment")
+
+    # If the user clicks "Save", save the relationships list to the session state
+    if st.button("Save to semantic model", use_container_width=True, type="primary"):
+        # Quickly validate that all of the user's joins have the required fields.
+        for relationship in st.session_state.builder_joins:
+            if not relationship.left_table or not relationship.right_table:
+                st.error("Please fill out left and right tables for all join paths.")
+                return
+
+            if not relationship.name:
+                st.error(
+                    f"The join path between {relationship.left_table} and {relationship.right_table} is missing a name."
+                )
+                return
+
+            if not relationship.relationship_columns:
+                st.error(
+                    f"The join path between {relationship.left_table} and {relationship.right_table} is missing joinable columns."
+                )
+                return
+
+        del st.session_state.semantic_model.relationships[:]
+        st.session_state.semantic_model.relationships.extend(
+            st.session_state.builder_joins
+        )
+        st.session_state.validated = None
+        st.rerun()