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

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