Michael McGuinness
138 lines
4.6 KiB
Python
138 lines
4.6 KiB
Python
# Use this code snippet in your app.
|
|
# If you need more information about configurations
|
|
# or implementing the sample code, visit the AWS docs:
|
|
# https://aws.amazon.com/developer/language/python/
|
|
|
|
import boto3
|
|
from botocore.exceptions import ClientError
|
|
import http.client
|
|
import base64
|
|
import ast
|
|
import snowflake.connector
|
|
import json
|
|
|
|
|
|
def get_secret():
|
|
|
|
secret_name = "doczy_dev_db_creds"
|
|
region_name = "us-east-2"
|
|
|
|
# Create a Secrets Manager client
|
|
session = boto3.session.Session()
|
|
client = session.client(service_name="secretsmanager", region_name=region_name)
|
|
|
|
try:
|
|
get_secret_value_response = client.get_secret_value(SecretId=secret_name)
|
|
except ClientError as e:
|
|
# For a list of exceptions thrown, see
|
|
# https://docs.aws.amazon.com/secretsmanager/latest/apireference/API_GetSecretValue.html
|
|
raise e
|
|
|
|
secret = get_secret_value_response["SecretString"]
|
|
return 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, **kwargs):
|
|
|
|
mwaa_env_name = "doczy-dev-infra-mwaa"
|
|
dag_name = dag_name
|
|
mwaa_cli_command = "dags trigger"
|
|
|
|
# Create the client with the specified profile
|
|
session = boto3.Session()
|
|
client = session.client("mwaa", region_name="us-east-2")
|
|
|
|
# get web token
|
|
mwaa_cli_token = client.create_cli_token(Name=mwaa_env_name)
|
|
|
|
conn = http.client.HTTPSConnection(mwaa_cli_token["WebServerHostname"])
|
|
|
|
# This section passes the payload to the MWAA CLI
|
|
# The file parameters should be added dynamically in streamlit, once the file names are passed while triggering the dag, the data will be ingested
|
|
# training_results_file = "training_results_sample.csv"
|
|
# attempt_logs_file = "attempt_logs_sample.csv"
|
|
# conf = "{\"" + "training_results_file_name" + "\":\"" + {training_results_file} + "\", \"" + "attempt_logs_file_name" + "\":\"" + {attempt_logs_file} + "\"}".format(training_results_file=training_results_file, attempt_logs_file=attempt_logs_file)
|
|
|
|
conf = json.dumps(kwargs)
|
|
|
|
payload = mwaa_cli_command + " " + dag_name + " --conf '{}'".format(conf)
|
|
headers = {
|
|
"Authorization": "Bearer " + mwaa_cli_token["CliToken"],
|
|
"Content-Type": "text/plain",
|
|
}
|
|
conn.request("POST", "/aws_mwaa/cli", payload, headers)
|
|
res = conn.getresponse()
|
|
data = res.read()
|
|
dict_str = data.decode("UTF-8")
|
|
mydata = ast.literal_eval(dict_str)
|
|
return payload
|
|
|
|
|
|
# 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"],
|
|
role=secret_dict["ROLE"],
|
|
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, bucket_name FROM STG.CLIENT_LOGS"
|
|
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]) # 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.CONTRACT_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
|