Auto ingest data to Snowflake with Snowpark Python API and RSA key based authentication

# loosely based on https://gist.github.com/vepetkov/a24b501e180cc5842b786ae680035239

import os
import snowflake.connector
from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives.asymmetric import rsa
from cryptography.hazmat.primitives.asymmetric import dsa
from cryptography.hazmat.primitives import serialization

# Params
whName = "COMPUTE_WH"
dbName = "DB_SNOWPARK"
schemaName = "BANK_MARKETING"
stgName = "BANK_STG"
#filePath = "./bank-marketing-de/bank_marketing.csv"
fileArray = ["./bank-marketing-de/bank_marketing.csv",
             "./bank-marketing-de/campaign.csv",
             "./bank-marketing-de/client.csv",
             "./bank-marketing-de/economics.csv",
             ]

# rsa key based authentication
def snowpark_key_auth():
    with open("rsa_key.p8", "rb") as key:
        p_key= serialization.load_pem_private_key(
        key.read(),
        password=None,
        backend=default_backend()
        )

    pkb = p_key.private_bytes(
        encoding=serialization.Encoding.DER,
        format=serialization.PrivateFormat.PKCS8,
        encryption_algorithm=serialization.NoEncryption())
 
    return pkb

# connect to snowpark
ctx = snowflake.connector.connect(
        account='yfnsuct-cc62406',
        user='API_SNOWPARK',
        private_key=snowpark_key_auth(),
        warehouse=whName,
        database=dbName,
        schema=schemaName
        )

cs = ctx.cursor()

# Step-1
for filePath in fileArray:
    # Upload each local file to the internal table stage
    uploadStmt = f'put file://{filePath} @{stgName} auto_compress=false;'
    cs.execute(f"use schema {schemaName};") # or set the Db & Scheme for the whole connection directly
    cs.execute(uploadStmt)

# Check the uploads
cs.execute(f"list @{stgName};")
uploadedFiles = cs.fetchall()
# check upload status
print(uploadedFiles)

cs.close()
ctx.close()

# based on DE Simplified course code & modified

import os
from datetime import datetime
import logging
from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives.asymmetric import rsa
from cryptography.hazmat.primitives.asymmetric import dsa
from cryptography.hazmat.primitives import serialization

def logging_setup():
    # create logger
    logger = logging.getLogger('simple_example')
    logger.setLevel(logging.DEBUG)

    # create console handler and set level to debug
    ch = logging.StreamHandler()
    ch.setLevel(logging.DEBUG)

    # create formatter
    formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')

    # add formatter to ch
    ch.setFormatter(formatter)

    # add ch to logger
    logger.addHandler(ch)

    logging.basicConfig(filename='debug.log', 
                        #encoding='utf-8', 
                        level=logging.DEBUG)
    return logger, ch

logger, ch = logging_setup()

# Params
whName = "COMPUTE_WH"
dbName = "DB_SNOWPARK"
schemaName = "BANK_MARKETING"
#stgName = "BANK_STG"

# separated to debug conda env packages
from snowflake.snowpark import Session

# rsa key based authentication
def snowpark_key_auth() -> Session:
    
    # load the private rsa key
    with open("rsa_key.p8", "rb") as key:
        p_key= serialization.load_pem_private_key(
            key.read(),
            password=None,
            backend=default_backend()
        )

    # private key text
    private_key_text = p_key.private_bytes(
    encoding=serialization.Encoding.DER,
    format=serialization.PrivateFormat.PKCS8,
    encryption_algorithm=serialization.NoEncryption())
    
    connection_parameters = {
       "ACCOUNT":"yfnsuct-cc62406",
        "USER":"API_SNOWPARK",
        "private_key" : private_key_text
    }
    # creating snowflake session object
    return Session.builder.configs(connection_parameters).create()  

#Generate dynamic ddl statement based on column name and type list
def generate_ddl_statement(column_names, data_types, table_name):

    ddl_template = "CREATE TABLE IF NOT EXISTS {} (\n{});"

    columns = []
    for name, data_type in zip(column_names, data_types):
        column_definition = f"    {name} {data_type}"
        columns.append(column_definition)

    ddl_statement = ddl_template.format(table_name, ",\n".join(columns))

    print("Got the table creation ddl statement")

    return ddl_statement

#Generate dynamic copy statement based on table name    
def generate_copy_statement(table_name,stage_name,csv_file_path,file_format):
    copy_command = f"""
    COPY INTO {table_name}
    FROM @{stage_name}/{csv_file_path}
    FILE_FORMAT = (FORMAT_NAME = '{file_format}')
    """

    print(copy_command)
    return copy_command

# call the main function  
def main(snowpark_session,named_stg,copy_file_format,infer_schema_file_format):
    
    #set the context (database, schema and warehouse)
    snowpark_session.sql("use role sysadmin").collect()
    snowpark_session.sql(f"use database {dbName}").collect()
    snowpark_session.sql(f"use schema {schemaName}").collect()
    snowpark_session.sql(f"use warehouse {whName}").collect()

    # Step-2
    # fetch all files under the named stage location
    list_sql = "list @{}".format(named_stg)
    stg_files = snowpark_session.sql(list_sql).collect()

    # Step-3
    # Iterate through the stage file entires....
    for row in stg_files:
        row_value = row.as_dict()
        stg_file_path_value = row_value.get('name')

        #split the file path and extract names
        file_path, file_name = os.path.split(stg_file_path_value)
        stg_location = "@"+file_path

        # Dynamically create infer-schema SQL statement
        # by passing stage location + file name 
        infer_schema_sql = """ \
            SELECT *
            FROM TABLE(
                INFER_SCHEMA(
                LOCATION=>'{}/',
                files => '{}',
                FILE_FORMAT=>'{}'
            )
        )
        """.format(stg_location, file_name,infer_schema_file_format)

        print(infer_schema_sql)

        # Execute the infer schema SQL statement.
        inferred_schema_rows = snowpark_session.sql(infer_schema_sql).collect()

        col_name_lst = []
        col_data_type_lst = []

        # Step-4
        # Iterate through the list & extract column-name/type
        for row in inferred_schema_rows:
            row_value = row.as_dict()
            column_name = row_value.get('COLUMN_NAME')
            column_type = row_value.get('TYPE')
            #print(file_name.split('.')[0],column_name,column_type)
            col_name_lst.append(column_name)
            col_data_type_lst.append(column_type)
            
        print("---------------")
        print(col_name_lst)
        print(col_data_type_lst)

        table_name = file_name.split('.')[0]+"_tbl"
        # Step-5
        create_ddl_stmt = generate_ddl_statement(col_name_lst,col_data_type_lst, table_name.upper())
        copy_stmt = generate_copy_statement(table_name,named_stg,file_name,copy_file_format)
        
        # Step-6
        sql_file_path = table_name+".sql"
        with open(sql_file_path, "w") as sql_file:
            sql_file.write("-- Following statement is creating table\n")
            sql_file.write(create_ddl_stmt)
            sql_file.write("\n-- Following statement is executing copy command \n")
            sql_file.write(copy_stmt)

        # Execute the ddl statement + copy statement
        # Step-7
        snowpark_session.sql(create_ddl_stmt).collect()
        snowpark_session.sql(copy_stmt).collect()

#entry point  
if __name__ == '__main__':
    #capture start time
    utc_start_time = datetime.utcnow()

    #get snowpark session
    snowpark_session = snowpark_key_auth()

    #set tag before starting the process
    snowpark_session.sql("alter session set query_tag = 'Start Snowpark Automation'").collect()

    #stage name & fime format names.
    named_stg = "BANK_STG"
    whName = "COMPUTE_WH"
    dbName = "DB_SNOWPARK"
    schemaName = "BANK_MARKETING"   
    copy_file_format = "copy_csv_ff"
    infer_schema_file_format = "infer_schema_csv_ff"

    #calling main method
    main(snowpark_session,named_stg,copy_file_format,infer_schema_file_format) 

    #record end time
    utc_end_time = datetime.utcnow()
    
    #set the query tag
    snowpark_session.sql("alter session set query_tag = 'End Snowpark Automation'").collect()

    #print overall execution time..
    print("Execution Time", utc_end_time-utc_start_time)