diff --git a/streamlit/interface_0.py b/streamlit/interface_0.py index 5d44ce6..593d686 100644 --- a/streamlit/interface_0.py +++ b/streamlit/interface_0.py @@ -11,6 +11,7 @@ import util import requests from sf_conn import get_secret, save_to_sf from io import StringIO, BytesIO +from sf_conn import get_client_names, insert_upload_logs create_batch_url = 'https://lfksus2t62.execute-api.us-east-2.amazonaws.com/dev/create-batch' @@ -49,6 +50,11 @@ if user_mail in user_list: bucket_list = ['doczy-ai-client-1', 'Delaware First Health, Inc.', 'Community Health Choice, Inc','CareSource Network Partners LLC', 'HealthNet of Cali', 'Oklahoma Complete Health, Inc', 'HealthFirst', 'Molina Healthcare of TX', 'AvMed', 'Arizona Care1st', 'WellCare New Jersey'] + + # This is the list of client fetched from Snowflake + # TODO: Need to update the streamlit code to use the client names from this list + # And use the s3 paths to save the objects for the respective client + client_names, s3_paths = get_client_names() client_row = st.columns([0.1, 0.8]) with client_row[0]: @@ -91,6 +97,10 @@ if user_mail in user_list: # 'config_interface/'+str(uploaded_file.name)) s3_client.put_object(Bucket=client, Body=stringio.getvalue(), Key= batch_id+'/'+landing_zone+'/'+str(uploaded_file.name)) + + # TODO: Test this insert function with snowflake + # insert_upload_logs(client, batch_id, str(uploaded_file.name), user_mail) + file_names.append(str(uploaded_file.name)) st.write(f"{batch_id} created") st.write(f"Files uploaded to s3://{client}/{batch_id}/{landing_zone}") diff --git a/streamlit/sf_conn.py b/streamlit/sf_conn.py index 7a728a4..cbd0545 100644 --- a/streamlit/sf_conn.py +++ b/streamlit/sf_conn.py @@ -8,6 +8,8 @@ from botocore.exceptions import ClientError import http.client import base64 import ast +import snowflake.connector + def get_secret(): @@ -35,7 +37,8 @@ def get_secret(): # get_secret() - +# TODO: This function needs to be changed to accept Kwargs +# The function name should be more generic and cofnigurable def save_to_sf(dag_name, conf_1, conf_2, training_results_file, attempt_logs_file): mwaa_env_name = 'doczy-dev-infra-mwaa' @@ -80,3 +83,61 @@ def save_to_sf(dag_name, conf_1, conf_2, training_results_file, attempt_logs_fil # save_to_sf("2024-03-13T17-36_prompt_results.csv", "2024-03-14T23-19_history.csv") + +# Create snowflake connection with the secrets fetched from secrets manager for the schema passed as an argument +def get_snowflake_conn(schema): + secret = get_secret() + secret_dict = eval(secret) + conn = snowflake.connector.connect( + user=secret_dict['user'], + password=secret_dict['password'], + account=secret_dict['account_alias'], + warehouse=secret_dict['warehouse'], + database=secret_dict['database'], + schema=schema + ) + return conn + +def get_client_names(): + """ + Input: None + Output: client_names, s3_paths lists + """ + try: + # Get conn from snowflake_conn function for STG schema + conn = get_snowflake_conn('STG') + cursor = conn.cursor() + + # Query to get client names and their s3_paths + query = "SELECT DISTINCT client_name, s3_bucket_path FROM STG.CLIENT_CONFIG" + cursor.execute(query) + # Create 2 lists from the query results + client_names = [] + s3_paths = [] + + for row in cursor: + client_names.append(row[0]) + s3_paths.append(row[1].replace('s3://','').split('/')[0]) # Extracting the bucket name from the s3 path + cursor.close() + conn.close() + return client_names, s3_paths + except Exception as e: + print(f"Error while fetching client names from snowflake: {e}") + return None, None + + +def insert_upload_logs(batch_id, client_name, file_name, upload_datetime, upload_user): + """ + Input: batch_id, client_name, file_name, upload_datetime, upload_user + Output: status of the insert query + """ + try: + conn = get_snowflake_conn('STG') + cursor = conn.cursor() + query = f"INSERT INTO STG.UPLOAD_LOGS (BATCH_ID, CLIENT_NAME, FILE_NAME, UPLOAD_DATETIME, UPLOAD_USER) VALUES ('{batch_id}', '{client_name}', '{file_name}', '{upload_datetime}', '{upload_user}')" + cursor.execute(query) + cursor.close() + conn.close() + return 'Log inserted successfully' + except Exception as e: + return e