Lambda Function for Running Snowflake Queries

snowflake_lambda.py

# UPDATE THE FOLLOWING VARIABLES IN THE CODE

# SNOWFLAKE_ACCOUNT: the account number of Snowflake (can be found in the 
# SNOWFLAKE_USER: Snowflake username
# SNOWFLAKE_PASSWORD: Snowflake password
# SNOWFLAKE_QUERY: the Snowflake query to run 
# S3_BUCKET: name of S3 bucket to store results
# S3_OUTPUT_KEY: name of the file to store with the results

import boto3
import csv
import snowflake.connector


SNOWFLAKE_ACCOUNT = "<snowflake_account>"
SNOWFLAKE_USER = "<snowflake_user>"
SNOWFLAKE_PASSWORD = "<snowflake_pwd>"

# simple query to return a value from Snowflake
SNOWFLAKE_QUERY = "<snowflake_query>"

# S3 details
S3_BUCKET = "<S3_BUCKET>" # e.g. "snowflake-query-bucket"
S3_KEY = "<S3_KEY>" # e.g. "output.csv"

def lambda_handler(event, context):
    
    con = snowflake.connector.connect(
        account=SNOWFLAKE_ACCOUNT,
        user=SNOWFLAKE_USER,
        password=SNOWFLAKE_PASSWORD
    )
    
    # run the test query to check if snowflake returns a result
    cur = con.cursor()
    try:
        cur.execute(SNOWFLAKE_QUERY)
        
        rows = cur.fetchall()
        with open(f'/tmp/{S3_OUTPUT_KEY}', 'w', newline='') as f:
            w = csv.writer(f)
            w.writerows(rows)

        #upload the data into s3
        s3 = boto3.client("s3")
        s3.upload_file(f'/tmp/{S3_OUTPUT_KEY}', S3_BUCKET, S3_OUTPUT_KEY)
        
        cur.close()
        return ("Query Execution Successful")
            
    except Exception as e:
        print(e)
        cur.close()
        return ("Query Execution Failed")