diff --git a/admin_apps/partner/looker.py b/admin_apps/partner/looker.py index fdcf0f68..f15b5fad 100644 --- a/admin_apps/partner/looker.py +++ b/admin_apps/partner/looker.py @@ -27,6 +27,7 @@ input_semantic_file_name, run_generate_model_str_from_snowflake, set_sit_query_tag, + set_table_comment, ) from semantic_model_generator.data_processing.proto_utils import proto_to_dict @@ -208,6 +209,9 @@ def set_looker_semantic() -> None: help="Specifies the name of the warehouse that provides the compute resources for refreshing the dynamic table.", ) st.divider() + st.write( + "Please fill out the following fields to start building your semantic model." + ) col1, col2 = st.columns(2) with col1: model_name = input_semantic_file_name() @@ -231,14 +235,13 @@ def set_looker_semantic() -> None: with st.spinner("Saving Explore as a Snowflake table..."): # Cortex model generation reqires fully-qualified table name - full_tablename = f"{st.session_state['looker_target_schema']}.{st.session_state['looker_target_table_name']}" + full_tablename = f"{st.session_state['looker_target_schema']}.{st.session_state['looker_target_table_name']}".upper() looker_columns = render_looker_explore_as_table( conn=get_snowflake_connection(), model_name=st.session_state["looker_model_name"].lower(), explore_name=st.session_state["looker_explore_name"].lower(), - snowflake_context=st.session_state["looker_target_schema"], - table_name=st.session_state["looker_target_table_name"].upper(), + table_name=full_tablename, optional_db=st.session_state["looker_connection_db"], fields=None, # TO DO - Add support for field selection dynamic=dynamic_table, @@ -255,13 +258,23 @@ def set_looker_semantic() -> None: sample_values, [full_tablename], ) - st.session_state["partner_setup"] = True + + # Tag looker setup with SIT query tag set_sit_query_tag( get_snowflake_connection(), vendor="looker", action="setup_complete", ) + # Tag table with SIT query tag + # We set table comment after generating semantic file to avoid using comment in description generation + set_table_comment( + get_snowflake_connection(), + full_tablename, + get_sit_query_tag(vendor="looker", action="materialize"), + table_type="DYNAMIC" if dynamic_table else None, + ) + # Take user to iteration using generated semantic file. # User can merge Looker metadata with generated semantic file on iteration page. st.session_state["page"] = GeneratorAppScreen.ITERATION @@ -373,7 +386,6 @@ def create_query_id( def create_explore_ctas( query_string: str, - snowflake_context: str, table_name: str, column_list: list[str], dynamic: Optional[bool] = False, @@ -385,8 +397,12 @@ def create_explore_ctas( Augments Looker Explore SQL with CTAS prefix to create a materialized view. Args: query_string (str): SQL select statement to render Looker Explore. - snowflake_context (str): Database_name.Schema_name for Snowflake - table_name (str): Name of table to create in snowflake_context context. + table_name (str): Fully-qualified name of table to create in Snowflake. + column_list (list[str]): List of column names to create in Snowflake table. + dynamic (bool): Flag to create dynamic table. Default is False. + target_lag (int): Target lag for dynamic table. Default is 20. + target_lag_unit (str): Target lag unit for dynamic table. Default is "minutes". + warehouse (str): Snowflake warehouse for dynamic table. Default is None. Returns: str CTAS statement """ @@ -396,16 +412,14 @@ def create_explore_ctas( filtered_lines = [ line for line in lines if not line.strip().startswith(("LIMIT", "FETCH")) ] + # full_table_name = f"{snowflake_context}.{table_name}" filtered_query_string = "\n".join(filtered_lines) columns = ", ".join(column_list) - comment = ( - f"COMMENT = '{get_sit_query_tag(vendor = 'looker', action = 'materialize')}'" - ) if dynamic: if not [x for x in (target_lag, target_lag_unit, warehouse) if x is None]: ctas_clause = f""" - CREATE OR REPLACE DYNAMIC TABLE {snowflake_context}.{table_name}({columns}) + CREATE OR REPLACE DYNAMIC TABLE {table_name}({columns}) TARGET_LAG = '{target_lag} {target_lag_unit}' WAREHOUSE = {warehouse} """ @@ -416,15 +430,11 @@ def create_explore_ctas( Standard materialization will be used. """ ) - ctas_clause = ( - f"CREATE OR REPLACE TABLE {snowflake_context}.{table_name}({columns}) " - ) + ctas_clause = f"CREATE OR REPLACE TABLE {table_name}({columns}) " else: - ctas_clause = ( - f"CREATE OR REPLACE TABLE {snowflake_context}.{table_name}({columns}) " - ) + ctas_clause = f"CREATE OR REPLACE TABLE {table_name}({columns}) " - ctas_clause += comment + " AS\n" + filtered_query_string + ctas_clause += " AS\n" + filtered_query_string return ctas_clause @@ -469,7 +479,6 @@ def render_looker_explore_as_table( conn: SnowflakeConnection, model_name: str, explore_name: str, - snowflake_context: str, table_name: str, optional_db: Optional[str] = None, fields: Optional[str] = None, @@ -484,10 +493,13 @@ def render_looker_explore_as_table( conn: SnowflakeConnection to run the query model_name (str): Looker model name. Should be lowercase. explore_name (str): Looker explore name. Should be lowercase. - snowflake_context (str): Database_name.Schema_name for Snowflake - table_name: The fully-qualified name of the table. + table_name: Fully-qualified name of table to create in Snowflake. fields (list): List-like str of fields to extract from the Explore. Default is None. Example: "id, name, description, fields", + dynamic (bool): Flag to create dynamic table. Default is False. + target_lag (int): Target lag for dynamic table. Default is 20. + target_lag_unit (str): Target lag unit for dynamic table. Default is "minutes". + warehouse (str): Snowflake warehouse for dynamic table. Default is None. Returns: dict[str, dict[str, str]] column metadata """ @@ -515,7 +527,6 @@ def render_looker_explore_as_table( return None ctas = create_explore_ctas( explore_sql, - snowflake_context, table_name, clean_columns, dynamic, @@ -526,7 +537,7 @@ def render_looker_explore_as_table( # Create materialized equivalent of Explore # Looker sources don't require explicit database qualification but instead use connection database implicitly. - # Set optional_db if user provides it. Need to reset back to original afterwards. + # Set optional_db if user provides it. May need to reset back to original afterwards. if optional_db: current_db = conn.cursor().execute("SELECT CURRENT_DATABASE();").fetchone()[0] # type: ignore conn.cursor().execute(f"USE DATABASE {optional_db};") diff --git a/admin_apps/shared_utils.py b/admin_apps/shared_utils.py index 92a0cbb2..95631f10 100644 --- a/admin_apps/shared_utils.py +++ b/admin_apps/shared_utils.py @@ -912,6 +912,22 @@ def set_sit_query_tag( conn.cursor().execute("alter session set query_tag=''") +def set_table_comment( + conn: SnowflakeConnection, + tablename: str, + comment: str, + table_type: Optional[str] = None, +) -> None: + """ + Sets comment on provided table. + Returns: None + """ + if table_type is None: + table_type = "" + query = f"ALTER {table_type} TABLE {tablename} SET COMMENT = '{comment}'" + conn.cursor().execute(query) + + def render_image(image_file: str, size: tuple[int, int]) -> None: """ Renders image in streamlit app with custom width x height by pixel.