AARETE\agupta
142 lines
4.7 KiB
Python
142 lines
4.7 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
|