Mayank Aamseek
83 lines
2.7 KiB
Python
83 lines
2.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
|
|
|
|
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()
|
|
|
|
|
|
def save_to_sf(training_results_file, attempt_logs_file):
|
|
|
|
mwaa_env_name = 'doczy-dev-infra-mwaa'
|
|
dag_name = 'load_training_results'
|
|
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 = f"""{{
|
|
"training_results_file_name": "{training_results_file}",
|
|
"attempt_logs_file_name": "{attempt_logs_file}"
|
|
}}"""
|
|
|
|
|
|
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 base64.b64decode(mydata['stdout'])
|
|
|
|
|
|
# save_to_sf("2024-03-13T17-36_prompt_results.csv", "2024-03-14T23-19_history.csv")
|
|
|