Added dynamic client names + s3 path to UI0 and added upload logs function. NEEDS TESTING

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}")

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