Create a Domino Job to automate weather data updates

import snowflake.connector
import os
import pandas as pd
import numpy as np
import requests
import shutil
import tarfile
import json

noaa_url = "https://www.ncei.noaa.gov/pub/data/ghcn/daily/superghcnd/"

# Make sure to point at a dataset directory as the file sizes are occasionally large
download_path = '/domino/datasets/local/Snowflake-Quickstart/downloads/'

# Save the countries you need for later operations
country_prefixes = ['AU', 'BE', 'FR', 'GM', 'IT', 'NL', 'PO', 'SP', 'SZ', 'UK']

# How many rows to process at once, to avoid running out of memory
chunksize = 10 ** 6

# The name you will use for the downloaded diff files
download_file_name = 'cur_diff.tar.gz'

# Columns used in your dataframe and temporary database table
column_names = ['station_id', 'data_date', 'element', 'element_value', 'mflag', 'qflag', 'sflag', 'obs_time', 'operation']

def update_weather_data():
    try:
        conn = initialize_connection()
    except:
        print("ERROR: Failed to connect to database. Exiting.")
        return 1
    else:
        print("Connected successfully to Snowflake...")

        update_start_date = get_update_start_date(conn)
        diff_df = get_diff_file_urls(update_start_date)

        if len(diff_df) == 0:
            print("Nothing to update. Exiting.")
            return 0

        result_list = perform_updates(diff_df, conn)

        with open('dominostats.json', 'w') as f:
            f.write(json.dumps({"Deletes": result_list[0], "Updates": result_list[1], "Inserts": result_list[2]}))
    finally:
        conn.close()

    return

#-----------------------------------------------------
# Returns: connection object
#-----------------------------------------------------

def initialize_connection():
    conn = snowflake.connector.connect(
        user = os.environ['username'],
        password = os.environ['password'],
        account = os.environ['accountname'],
        warehouse = os.environ['warehousename'],
        schema = os.environ['schemaname'],
        database = os.environ['dbname'])

    return conn

#----------------------------------------------------------------
# Argument(s): The connection object
# Returns: The date of the newest data record in the database
#----------------------------------------------------------------

def get_update_start_date(conn):
    cursor = conn.cursor()
    # Get the most recent date for which we have data
    cursor.execute("SELECT DATA_DATE FROM STATION_DATA ORDER BY DATA_DATE DESC LIMIT 1;")

    df = cursor.fetch_pandas_all()
    cursor.close()
    last_update_date = df['DATA_DATE'][0]

    print(f"Most recent record in database is from {last_update_date}")

    return last_update_date

#------------------------------------------------------------------
# Argument(s): The date of the most recent record in the database
# Returns: Dataframe with list of files to read for the data update
#------------------------------------------------------------------

def get_diff_file_urls(last_update_date):
    print("Checking updates from NOAA...")
    page = pd.read_html(noaa_url)
    df = page[0]
    df['Last modified'] = pd.to_datetime(df['Last modified'])
    df = df.astype({'Name':'string'})
    # Keep diff files only
    diff_df = df[df['Name'].str.contains('diff')].sort_values(by=['Last modified'], ascending=False)

    # Limit to only the diff files that are relevant to you - files newer than your last update date
    diff_df=diff_df[diff_df['Last modified'].dt.date > last_update_date]

    # Sort the files so the oldest file is first
    diff_df = diff_df.sort_values('Last modified')

    # Remove files with gigabyte-sized updates (normally an NOAA error)
    diff_df = diff_df[diff_df['Size'].str.find('G') == -1]

    print(f"Need to retrieve {len(diff_df)} update files")
    return diff_df

#------------------------------------------------------------------
# Argument(s):
# 1. Dataframe containing diff file URLs in oldest to newest order
# 2. Database connection object
# Returns: List containing number of updates
#------------------------------------------------------------------

def perform_updates(diff_df, conn):
    cursor = conn.cursor()

    delete_count = 0
    update_count = 0
    insert_count = 0

    # Iterate over all diff files
    for index, row in diff_df.iterrows():

        # Define the size of data chunks for us to read at a time
        chunksize = 10 ** 6
        oldest_file = getattr(row, 'Name')
        oldest_file_url = noaa_url + oldest_file

        print(f"Now dowloading {oldest_file}...")

        # Download the file
        response = requests.get(oldest_file_url, stream=True)
        if response.status_code == 200:
            with open(download_path+download_file_name, 'wb') as f:
                f.write(response.raw.read())
        else:
            print(f"Failed to download {oldest_file}. SKIPPING!")
            continue

        print("Done. Extracting...")

        # Extract the file
        file = tarfile.open(download_path+download_file_name)
        file.extractall(download_path)
        file.close()

        print("Done. Processing...")

        # Get the directory name by removing the .tar.gz
        dir_name = oldest_file.split('.', 1)[0]

        # Perform updates
        delete_count += do_delete(cursor, dir_name)
        update_count += do_update(cursor, dir_name)
        insert_count += do_insert(cursor, dir_name)

        # Clean up
        shutil.rmtree(download_path+dir_name)

    cursor.close()

    result_list = [delete_count, update_count, insert_count]

    return result_list

#------------------------------------------------------------------
# Argument(s):
# 1. Database connection cursor
# 2. File path
# Returns: Number of rows deleted
#------------------------------------------------------------------

def do_delete(cursor, dir_name):

    print(">>> Starting delete operation <<<")

    chunkcount = 0
    west_eu_df = pd.DataFrame(columns = column_names)

    for del_df in pd.read_csv(download_path+dir_name+'/delete.csv', names = column_names,
                         dtype={'station_id': 'str', 'data_date': 'str', 'element': 'str', 'element_value': 'int',
                                'mflag': 'str', 'qflag':'str', 'sflag':'str', 'obs_time': 'str'},
                         chunksize = chunksize):
        chunkcount = chunkcount + 1
        print(f"Reading chunk {chunkcount} of delete data")
        del_df['data_date'] = pd.to_datetime(del_df['data_date'], yearfirst=True)

        # Filter the file for Western European stations
        for cur_country in country_prefixes:
            country_df = del_df[del_df['station_id'].str.startswith(cur_country)]
            west_eu_df = pd.concat([west_eu_df, country_df], ignore_index=True)

    print(f"Read data in {chunkcount} chunks. Deleting {len(west_eu_df)} records")

    # Output file
    west_eu_df.to_csv(download_path+dir_name+'/diff-data.csv', index=False, sep="|")

    sfQuery = """CREATE OR REPLACE TEMPORARY TABLE TEMP_DIFF_DATA
                (station_id TEXT,
                data_date DATE,
                element TEXT,
                element_value INT,
                mflag char(1),
                qflag char(1),
                sflag char(1),
                obs_time TIME,
                operation char(1))"""
    cursor.execute(sfQuery)
    print("table created")
    sfQuery = 'PUT file://' + download_path+dir_name+'/diff-data.csv @%TEMP_DIFF_DATA'
    cursor.execute(sfQuery)
    print("file uploaded")
    cursor.execute("COPY INTO TEMP_DIFF_DATA file_format = (type = csv field_delimiter = '|' skip_header = 1)")
    print("data copied from file to table")

    cursor.execute("BEGIN TRANSACTION")

    sfQuery = """DELETE FROM STATION_DATA sd USING TEMP_DIFF_DATA tdd
                 WHERE sd.station_id = tdd.station_id
                 AND sd.element = tdd.element
                 AND sd.data_date = tdd.data_date""";
    try:
        print("Executing deletion...")
        cursor.execute(sfQuery)
    except:
        cursor.execute("ROLLBACK")
        print("Transaction failed")
        return 0
    else:
        cursor.execute("COMMIT")
        print("Delete transaction successful")

    os.remove(download_path+dir_name+'/diff-data.csv')

    return len(west_eu_df)

#------------------------------------------------------------------
# Argument(s):
# 1. Database connection cursor
# 2. File path
# Returns: Number of rows updated
#------------------------------------------------------------------

def do_update(cursor, dir_name):

    print(">>> Starting update operation <<<")
    west_eu_df = pd.DataFrame(columns = column_names)
    chunkcount = 0

    with pd.read_csv(download_path+dir_name+'/update.csv', names = column_names,
                         dtype={'station_id': 'str', 'data_date': 'str', 'element': 'str', 'element_value': 'int',
                                'mflag': 'str', 'qflag':'str', 'sflag':'str', 'obs_time': 'str'},
                         chunksize = chunksize) as reader:

        for update_df in reader:

            chunkcount = chunkcount + 1
            update_df['data_date'] = pd.to_datetime(update_df['data_date'], yearfirst=True)

            for cur_country in country_prefixes:
                country_df = update_df[update_df['station_id'].str.startswith(cur_country)]
                west_eu_df = pd.concat([west_eu_df, country_df], ignore_index=True)


    print(f"Read data in {chunkcount} chunks. Updating {len(west_eu_df)} records")

    # Output file
    west_eu_df.to_csv(download_path+dir_name+'/diff-data.csv', index=False, sep="|")

    sfQuery = """CREATE OR REPLACE TEMPORARY TABLE TEMP_DIFF_DATA
                (station_id TEXT,
                data_date DATE,
                element TEXT,
                element_value INT,
                mflag char(1),
                qflag char(1),
                sflag char(1),
                obs_time TIME,
                operation char(1))"""
    cursor.execute(sfQuery)
    print("table created")
    sfQuery = 'PUT file://' + download_path+dir_name+'/diff-data.csv @%TEMP_DIFF_DATA'
    cursor.execute(sfQuery)
    print("file uploaded")
    cursor.execute("COPY INTO TEMP_DIFF_DATA file_format = (type = csv field_delimiter = '|' skip_header = 1)")
    print("data copied")

    cursor.execute("BEGIN TRANSACTION")

    sfQuery = """UPDATE STATION_DATA sd
                SET sd.element_value = tdd.element_value
                FROM TEMP_DIFF_DATA tdd
                 WHERE sd.station_id = tdd.station_id
                 AND sd.element = tdd.element
                 AND sd.data_date = tdd.data_date""";
    try:
        print("Executing update...")
        cursor.execute(sfQuery)
    except:
        cursor.execute("ROLLBACK")
        print("Transaction failed")
        return 0
    else:
        cursor.execute("COMMIT")
        print("Update transaction successful")

    os.remove(download_path+dir_name+'/diff-data.csv')

    return len(west_eu_df)

#------------------------------------------------------------------
# Argument(s):
# 1. Database connection cursor
# 2. File path
# Returns: Number of rows inserted
#------------------------------------------------------------------

def do_insert(cursor, dir_name):

    print(">>> Starting insert operation <<<")

    west_eu_df = pd.DataFrame(columns = column_names)
    chunkcount = 0

    with pd.read_csv(download_path+dir_name+'/insert.csv', names = column_names,
                         dtype={'station_id': 'str', 'data_date': 'str', 'element': 'str', 'element_value': 'int',
                                'mflag': 'str', 'qflag':'str', 'sflag':'str', 'obs_time': 'str'},
                    chunksize = chunksize) as reader:

        for insert_df in reader:
            chunkcount = chunkcount + 1
            insert_df['data_date'] = pd.to_datetime(insert_df['data_date'], yearfirst=True)

            for cur_country in country_prefixes:
                country_df = insert_df[insert_df['station_id'].str.startswith(cur_country)]
                west_eu_df = pd.concat([west_eu_df, country_df], ignore_index=True)

    print(f"Read data in {chunkcount} chunks. Inserting {len(west_eu_df)} records")

    west_eu_df.to_csv(download_path+dir_name+'/diff-data.csv', index=False, sep="|")

    sfQuery = """CREATE OR REPLACE TEMPORARY TABLE TEMP_DIFF_DATA
                (station_id TEXT,
                data_date DATE,
                element TEXT,
                element_value INT,
                mflag char(1),
                qflag char(1),
                sflag char(1),
                obs_time TIME,
                operation char(1))"""
    cursor.execute(sfQuery)
    print("table created")
    sfQuery = 'PUT file://' + download_path+dir_name+'/diff-data.csv @%TEMP_DIFF_DATA'
    cursor.execute(sfQuery)
    print("file uploaded")
    cursor.execute("COPY INTO TEMP_DIFF_DATA file_format = (type = csv field_delimiter = '|' skip_header = 1)")
    print("data copied")

    cursor.execute("BEGIN TRANSACTION")

    sfQuery = """INSERT INTO STATION_DATA
                (station_id, data_date, element, element_value, mflag, qflag, sflag, obs_time)
                SELECT station_id, data_date, element, element_value, mflag, qflag, sflag, obs_time
                FROM TEMP_DIFF_DATA""";
    try:
        print("Executing insert...")
        cursor.execute(sfQuery)
    except:
        cursor.execute("ROLLBACK")
        print("Transaction failed")
        return 0
    else:
        cursor.execute("COMMIT")
        print("Transaction successful")

    return len(west_eu_df)

if __name__ == "__main__":
    update_weather_data()