Upload app.py

app.py

@@ -6,6 +6,23 @@ import time
 import pandas as pd
 import labelbox
 
+@st.cache_data(show_spinner=True)
+def fetch_databases(cluster_id, formatted_title, databricks_api_key):
+    query = "SHOW DATABASES;"
+    return execute_databricks_query(query, cluster_id, formatted_title, databricks_api_key)
+
+# Cached function to fetch tables
+@st.cache_data(show_spinner=True)
+def fetch_tables(selected_database, cluster_id, formatted_title, databricks_api_key):
+    query = f"SHOW TABLES IN {selected_database};"
+    return execute_databricks_query(query, cluster_id, formatted_title, databricks_api_key)
+
+# Cached function to fetch columns
+@st.cache_data(show_spinner=True)
+def fetch_columns(selected_database, selected_table, cluster_id, formatted_title, databricks_api_key):
+    query = f"SHOW COLUMNS IN {selected_database}.{selected_table};"
+    return execute_databricks_query(query, cluster_id, formatted_title, databricks_api_key)
+
 def validate_dataset_name(name):
     """Validate the dataset name."""
     # Check length
@@ -193,7 +210,7 @@ if new_dataset_name or selected_dataset_name:
     formatted_title = re.sub(r'/$', '', formatted_title)  # Remove trailing slash if present
 
     if formatted_title:
-        st.subheader("Select and existing cluster or make a new one", divider='grey', help="Jobs in preview mode will use all purpose compute clusters to help you itersate faster. Jobs in production mode will use job clusters to reduce DBUs consumed.")
+        st.subheader("Select an existing cluster", divider='grey', help="Jobs will use job clusters to reduce DBUs consumed.")
         DOMAIN = f"https://{formatted_title}"
         TOKEN = f"Bearer {databricks_api_key}"
 
@@ -217,23 +234,17 @@ if new_dataset_name or selected_dataset_name:
             }
 
             # Display dropdown with cluster names
-            make_cluster = st.toggle('Make me a new cluster', value=False)
-            if make_cluster:
-                #make a cluster
-                st.write("Making a new cluster")
-            else:
-                if cluster_dict:
-                    selected_cluster_name = st.selectbox(
-                        'Select a cluster to run on',
-                        list(cluster_dict.keys()),
-                        key='unique_key_for_cluster_selectbox',
-                        index=None,
-                        placeholder="Select a cluster..",
-                    )
-                    if selected_cluster_name:
-                        cluster_id = cluster_dict[selected_cluster_name]
-                else:
-                    st.write("No UI or API-based compute clusters found.")
+            
+            if cluster_dict:
+                selected_cluster_name = st.selectbox(
+                    'Select a cluster to run on',
+                    list(cluster_dict.keys()),
+                    key='unique_key_for_cluster_selectbox',
+                    index=None,
+                    placeholder="Select a cluster..",
+                )
+                if selected_cluster_name:
+                    cluster_id = cluster_dict[selected_cluster_name]
 
         except requests.RequestException as e:
             st.write(f"Error communicating with Databricks API: {str(e)}")
@@ -298,7 +309,7 @@ if new_dataset_name or selected_dataset_name:
             st.subheader("Run Frequency", divider='grey')
 
             # Dropdown to select frequency
-            freq_options = ["1 minute", "1 hour", "1 day", "1 week", "1 month"]
+            freq_options = ["1 day", "1 week", "1 month"]
             selected_freq = st.selectbox("Select frequency:", freq_options, placeholder="Select frequency..")
 
             day_of_week = None
@@ -306,12 +317,7 @@ if new_dataset_name or selected_dataset_name:
 
             # If the frequency is hourly, daily, weekly, or monthly, ask for a specific time
             if selected_freq != "1 minute":
-                col1, col2 = st.columns(2)
-                with col1:
-                    hour = st.selectbox("Hour:", list(range(0, 24)))
-                with col2:
-                    minute = st.selectbox("Minute:", list(range(0, 60)))
-                
+               
                 if selected_freq == "1 week":
                     days_options = ["MON", "TUE", "WED", "THU", "FRI", "SAT", "SUN"]
                     day_of_week = st.selectbox("Select day of the week:", days_options)
@@ -319,12 +325,49 @@ if new_dataset_name or selected_dataset_name:
                 elif selected_freq == "1 month":
                     day_of_month = st.selectbox("Select day of the month:", list(range(1, 32)))
 
+                col1, col2 = st.columns(2)    
+                with col1:
+                    hour = st.selectbox("Hour:", list(range(0, 24)))
+                with col2:
+                    minute = st.selectbox("Minute:", list(range(0, 60)))
+
             else:
                 hour, minute = 0, 0
 
             # Generate the cron expression
             frequency = generate_cron_expression(selected_freq, hour, minute, day_of_week, day_of_month)
 
+            # Assumed DBU consumption rate for a 32GB, 4-core node per hour
+            X = 1  # Replace this with the actual rate from Databricks' pricing or documentation
+
+            # Calculate DBU consumption for a single run
+            min_dbu_single_run = (X/6) * (1 + 10)  # Assuming maximum scaling to 10 workers
+            max_dbu_single_run = (2*X/3) * (1 + 10)
+
+            # Estimate monthly DBU consumption based on frequency
+            frequency = st.selectbox("Select Run Frequency:", ["Daily", "Weekly", "Monthly"])
+            if frequency == "Daily":
+                min_dbu_monthly = 30 * min_dbu_single_run
+                max_dbu_monthly = 30 * max_dbu_single_run
+            elif frequency == "Weekly":
+                min_dbu_monthly = 4 * min_dbu_single_run
+                max_dbu_monthly = 4 * max_dbu_single_run
+            else:  # Monthly
+                min_dbu_monthly = min_dbu_single_run
+                max_dbu_monthly = max_dbu_single_run
+
+            # Calculate runs per month
+            if selected_freq == "1 day":
+                runs_per_month = 30
+            elif selected_freq == "1 week":
+                runs_per_month = 4
+            else:  # "1 month"
+                runs_per_month = 1
+
+            # Calculate estimated DBU consumption per month
+            min_dbu_monthly = runs_per_month * min_dbu_single_run
+            max_dbu_monthly = runs_per_month * max_dbu_single_run
+
             def generate_human_readable_message(freq, hour=0, minute=0, day_of_week=None, day_of_month=None):
                 """
                 Generate a human-readable message for the scheduling.
@@ -349,37 +392,32 @@ if new_dataset_name or selected_dataset_name:
             # Generate the human-readable message
             readable_msg = generate_human_readable_message(selected_freq, hour, minute, day_of_week, day_of_month)
 
+            # Main code block
             if frequency:
                 st.success(readable_msg, icon="📅")
+                # Display the estimated DBU consumption to the user
+                st.warning(f"Estimated DBU Consumption:\n- For a single run: {min_dbu_single_run:.2f} to {max_dbu_single_run:.2f} DBUs\n- Monthly (based on {runs_per_month} runs): {min_dbu_monthly:.2f} to {max_dbu_monthly:.2f} DBUs")
 
+                # Disclaimer
+                st.info("Disclaimer: This is only an estimation. Always monitor the job in Databricks to assess actual DBU consumption.")
+                
                 st.subheader("Select a table", divider="grey")
 
-                with st.spinner('Querying Databricks...'):
-                    query = "SHOW DATABASES;"
-                    result_data = execute_databricks_query(query, cluster_id, formatted_title, databricks_api_key)
-
-                    # Extract the databaseName values from the DataFrame
-                    database_names = result_data['databaseName'].tolist()
-
-                # Create a dropdown with the database names
+                # Fetching databases
+                result_data = fetch_databases(cluster_id, formatted_title, databricks_api_key)
+                database_names = result_data['databaseName'].tolist()
                 selected_database = st.selectbox("Select a Database:", database_names, index=None, placeholder="Select a database..")
 
                 if selected_database:
-                    with st.spinner('Querying Databricks...'):
-                        query = f"SHOW TABLES IN {selected_database};"
-                        result_data = execute_databricks_query(query, cluster_id, formatted_title, databricks_api_key)
-
-                        # Extract the tableName values from the DataFrame
-                        table_names = result_data['tableName'].tolist()
-
-                    # Create a dropdown with the database names
+                    # Fetching tables
+                    result_data = fetch_tables(selected_database, cluster_id, formatted_title, databricks_api_key)
+                    table_names = result_data['tableName'].tolist()
                     selected_table = st.selectbox("Select a Table:", table_names, index=None, placeholder="Select a table..")
 
                     if selected_table:
-                        with st.spinner('Querying Databricks...'):
-                            query = f"SHOW COLUMNS IN {selected_database}.{selected_table};"
-                            result_data = execute_databricks_query(query, cluster_id, formatted_title, databricks_api_key)
-                            column_names = result_data['col_name'].tolist()
+                        # Fetching columns
+                        result_data = fetch_columns(selected_database, selected_table, cluster_id, formatted_title, databricks_api_key)
+                        column_names = result_data['col_name'].tolist()
 
                         st.subheader("Map table schema to Labelbox schema", divider="grey")
                         # Your existing code to handle schema mapping...
@@ -388,11 +426,8 @@ if new_dataset_name or selected_dataset_name:
                         with st.spinner('Fetching first 5 rows of the selected table...'):
                             query = f"SELECT * FROM {selected_database}.{selected_table} LIMIT 5;"
                             table_sample_data = execute_databricks_query(query, cluster_id, formatted_title, databricks_api_key)
+                            st.write(table_sample_data)
 
-                        # Display the sample data in the Streamlit UI
-                        st.write(table_sample_data)
-
-                        
                         # Define two columns for side-by-side selectboxes
                         col1, col2 = st.columns(2)
 
@@ -423,75 +458,86 @@ if new_dataset_name or selected_dataset_name:
                             sample_row_data_value = result_sample[selected_row_data].iloc[0]
 
                             # Validate the extracted value
-                            if is_valid_url_or_uri(sample_row_data_value):
-                                st.success(f"Sample URI/URL from selected row data column: {sample_row_data_value}", icon="✅")
-                                dataset_id = create_new_dataset_labelbox(new_dataset_name) if create_new_dataset else dataset_id
-                                # Mode
-                                mode = "preview" if is_preview else "production"
-
-                                # Databricks instance and API key
-                                databricks_instance = formatted_title
-                                databricks_api_key = databricks_api_key
-
-                                # Dataset ID and New Dataset
-                                new_dataset = 1 if create_new_dataset else 0
-                                dataset_id = dataset_id 
-
-                                # Table Path
-                                table_path = f"{selected_database}.{selected_table}"
-                                # Frequency
-                                frequency = frequency
-
-                                # Cluster ID and New Cluster
-                                new_cluster = 1 if make_cluster else 0
-                                cluster_id = cluster_id if not make_cluster else ""
-
-                                # Schema Map
-                                row_data_input = selected_row_data
-                                global_key_input = selected_global_key
-                                schema_map_dict = {'row_data': row_data_input}
-                                if global_key_input:
-                                    schema_map_dict['global_key'] = global_key_input
-
-                                # Convert the dict to a stringified JSON
-                                schema_map_str = json.dumps(schema_map_dict)
-                                
-                        
-                                data = {
-                                    "mode": mode,
-                                    "databricks_instance": databricks_instance,
-                                    "databricks_api_key": databricks_api_key,
-                                    "new_dataset": new_dataset,
-                                    "dataset_id": dataset_id,
-                                    "table_path": table_path,
-                                    "labelbox_api_key": labelbox_api_key,
-                                    "frequency": frequency,
-                                    "new_cluster": new_cluster,
-                                    "cluster_id": cluster_id,
-                                    "schema_map": schema_map_str
-                                }
-                                
-
-                                if st.button("Deploy Pipeline!", type="primary"):
-                                    # Ensure all fields are filled out
-                                    required_fields = [
-                                        mode, databricks_instance, databricks_api_key, new_dataset, dataset_id,
-                                        table_path, labelbox_api_key, frequency, new_cluster, cluster_id, schema_map_str
-                                    ]
-
-
-                                    # Sending a POST request to the Flask app endpoint
-                                    with st.spinner("Deploying pipeline..."):
-                                        response = requests.post("https://us-central1-dbt-prod.cloudfunctions.net/deploy-databricks-pipeline", json=data)
-
-                                    # Check if request was successful
-                                    if response.status_code == 200:
-                                        # Display the response using Streamlit
-                                        st.balloons()
-                                        st.success("Pipeline deployed successfully!", icon="🚀")
-                                        st.json(response.json())
-                                    else:
-                                        st.error(f"Failed to deploy pipeline. Response: {response.text}", icon="🚫")
-
-                            else:
-                                st.error(f"row_data '{sample_row_data_value}' is not a valid URI or URL. Please select a different column.", icon="🚫")
+                            dataset_id = create_new_dataset_labelbox(new_dataset_name) if create_new_dataset else dataset_id
+                            # Mode
+                            mode = "preview" if is_preview else "production"
+
+                            # Databricks instance and API key
+                            databricks_instance = formatted_title
+                            databricks_api_key = databricks_api_key
+
+                            # Dataset ID and New Dataset
+                            new_dataset = 1 if create_new_dataset else 0
+                            dataset_id = dataset_id 
+
+                            # Table Path
+                            table_path = f"{selected_database}.{selected_table}"
+                            # Frequency
+                            frequency = frequency
+
+                            # Schema Map
+                            row_data_input = selected_row_data
+                            global_key_input = selected_global_key
+                            schema_map_dict = {'row_data': row_data_input}
+                            if global_key_input:
+                                schema_map_dict['global_key'] = global_key_input
+
+                            # Convert the dict to a stringified JSON
+                            schema_map_str = json.dumps(schema_map_dict)
+                            
+                    
+                            data = {
+                                "mode": mode,
+                                "databricks_instance": databricks_instance,
+                                "databricks_api_key": databricks_api_key,
+                                "new_dataset": new_dataset,
+                                "dataset_id": dataset_id,
+                                "table_path": table_path,
+                                "labelbox_api_key": labelbox_api_key,
+                                "frequency": frequency,
+                                "new_cluster": 0,
+                                "cluster_id": cluster_id,
+                                "schema_map": schema_map_str
+                            }
+                            
+
+                            if st.button("Deploy Pipeline!", type="primary"):
+                                # Ensure all fields are filled out
+                                required_fields = [
+                                    mode, databricks_instance, databricks_api_key, new_dataset, dataset_id,
+                                    table_path, labelbox_api_key, frequency, cluster_id, schema_map_str
+                                ]
+
+
+                                # Sending a POST request to the Flask app endpoint
+                                with st.spinner("Deploying pipeline..."):
+                                    response = requests.post("https://us-central1-dbt-prod.cloudfunctions.net/deploy-databricks-pipeline", json=data)
+
+                                # Check if request was successful
+                                if response.status_code == 200:
+                                    # Display the response using Streamlit
+                                    st.balloons()
+                                    response = response.json()
+                                    # Extract the job_id
+                                    job_id = response['message'].split('job_id":')[1].split('}')[0]
+                                    from urllib.parse import urlparse, parse_qs
+
+                                    # Parse the Databricks instance URL to extract the organization ID
+                                    parsed_url = urlparse(formatted_title)
+                                    query_params = parse_qs(parsed_url.query)
+                                    organization_id = query_params.get("o", [""])[0]
+
+                                    # Generate the Databricks Job URL
+                                    job_url = f"http://{formatted_title}/?o={organization_id}#job/{job_id}"
+                                    st.success(f"Pipeline deployed successfully! [{job_url}]({job_url}) 🚀")
+                                else:
+                                    st.error(f"Failed to deploy pipeline. Response: {response.text}", icon="🚫")
+
+st.markdown("""
+<style>
+/* Add a large bottom padding to the main content */
+.main .block-container {
+    padding-bottom: 1000px;  /* Adjust this value as needed */
+}
+</style>
+""", unsafe_allow_html=True)