doczyai-pipelines/archive/streamlit/multipage/sf_conn.py

# 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