Local Testing Demo
"""
This module provides utilities for local development and testing.
"""

from pathlib import Path
from os import environ

import configparser
import toml


def get_env_var_config() -> dict:
    """
    Returns a dictionary of the connection parameters using the SnowSQL CLI
    environment variables.
    """
    try:
        return {
            "user": environ["SNOWSQL_USER"],
            "password": environ["SNOWSQL_PWD"],
            "account": environ["SNOWSQL_ACCOUNT"],
            # "role": environ["SNOWSQL_ROLE"],
            "warehouse": environ["SNOWSQL_WAREHOUSE"],
            "database": environ["SNOWSQL_DATABASE"],
            "schema": environ["SNOWSQL_SCHEMA"],
        }
    except KeyError as exc:
        raise KeyError(
            "ERROR: Environment variable for Snowflake Connection not found. "
            + "Please set the SNOWSQL_* environment variables"
        ) from exc


def get_dev_config(
    environment: str = "dev",
    app_config_path: Path = Path.cwd().joinpath("app.toml"),
) -> dict:
    """
    Returns a dictionary of the connection parameters using the app.toml
    in the project root.
    """
    try:
        app_config = toml.load(app_config_path)
        config = configparser.ConfigParser(inline_comment_prefixes="#")
        config.read(app_config["snowsql_config_path"])
        session_config = config["connections." + app_config["snowsql_connection_name"]]
        session_config_dict = {
            k.replace("name", ""): v.strip('"') for k, v in session_config.items()
        }
        session_config_dict.update(app_config.get(environment))  # type: ignore
        return session_config_dict
    except Exception as exc:
        raise EnvironmentError(
            "Error creating snowpark session - be sure you've logged into "
            "the SnowCLI and have a valid app.toml file",
        ) from exc

project/local.py
sfc-gh-jfreeberg/local-testing-ff
from snowflake.snowpark.dataframe import DataFrame
from snowflake.snowpark.session import Session
from project.transformers import *
from project.local import get_env_var_config

def create_fact_tables(sess: Session, source_table) -> int:
    """
    This job applies the transformations in transformers.py to the built-in Citibike dataset
    and saves two tables, month_summary and bike_summary, under CITIBIKE.PUBLIC. 
    Returns the total number of rows created in both tables
    """

    SOURCE_DB = 'CITIBIKE'
    SOURCE_SCHEMA = 'PUBLIC'

    df: DataFrame = sess.table([SOURCE_DB, SOURCE_SCHEMA, source_table])
    
    df2 = add_rider_age(df)
    month_facts = calc_month_facts(df2)
    bike_facts = calc_bike_facts(df2)

    month_facts.write.save_as_table([SOURCE_DB, SOURCE_SCHEMA, 'month_facts'], table_type='', mode='overwrite')
    bike_facts.write.save_as_table([SOURCE_DB, SOURCE_SCHEMA, 'bike_facts'], table_type='', mode='overwrite')

    return month_facts.count() + bike_facts.count()

if __name__ == '__main__':
    print('Creating session')
    session = Session.builder.configs(get_env_var_config()).create()

    print('Running job...')
    rows = create_fact_tables(session, 'PUBLIC')

    print('Job complete. Number of rows created:')
    print(rows)


project/sproc.py
sfc-gh-jfreeberg/local-testing-ff
from snowflake.snowpark.dataframe import DataFrame, col
from snowflake.snowpark.functions import monthname, avg, cast, count
import datetime


def add_rider_age(df: DataFrame) -> DataFrame:
    current_year = datetime.date.today().year
    return df.with_column('RIDER_AGE', current_year - df['BIRTH_YEAR'])
    

def calc_month_facts(df: DataFrame) -> DataFrame:
    """
    Group by month and return summary stats
    """
    return df.with_column('month', monthname('STARTTIME'))\
        .group_by('month')\
        .agg(
            count('*').alias('COUNT'), 
            avg('TRIPDURATION').alias('AVG_TRIPDURATION'), 
            avg('RIDER_AGE').alias('AVG_RIDER_AGE'))


def calc_bike_facts(df: DataFrame) -> DataFrame:
    """
    Group by bike ID and return summary statistics
    """
    return df.group_by('BIKEID')\
        .agg(
            count('*').alias('COUNT'), 
            avg('TRIPDURATION').alias('AVG_TRIPDURATION'), 
            avg('RIDER_AGE').alias('AVG_RIDER_AGE'))

project/transformers.py
sfc-gh-jfreeberg/local-testing-ff
"""
The Conda environment.yml will use this to create an editable installation
of this project.
"""

from setuptools import setup, find_packages

setup(
    name="Example Snowpark Python project",
    version="0.1.0",
    packages=find_packages()
)

setup.py
sfc-gh-jfreeberg/local-testing-ff
import pytest
from snowflake.snowpark.session import Session
from snowflake.snowpark.mock.connection import MockServerConnection
from project.local import get_env_var_config

def pytest_addoption(parser):
    parser.addoption("--snowflake-session", action="store", default="live")


@pytest.fixture(scope='module')
def session(request) -> Session:
    if request.config.getoption('--snowflake-session') == 'local':
        return Session(MockServerConnection())
    else:
        return Session.builder.configs(get_env_var_config()).create()

test/conftest.py
sfc-gh-jfreeberg/local-testing-ff
from snowflake.snowpark.mock.functions import patch
from snowflake.snowpark.functions import monthname
from snowflake.snowpark.mock.snowflake_data_type import ColumnEmulator
from snowflake.snowpark.types import TimestampType
import datetime
import calendar

@patch(monthname)
def patch_to_timestamp(column: ColumnEmulator) -> ColumnEmulator:
    ret_column = ColumnEmulator(data=[
        calendar.month_abbr[datetime.datetime.strptime(row, '%Y-%m-%d %H:%M:%S.%f %z').month]
        for row in column])
    ret_column.sf_type = TimestampType()
    return ret_column
test/patches.py
sfc-gh-jfreeberg/local-testing-ff
from snowflake.snowpark.session import Session
from snowflake.snowpark.types import StructType, StructField, IntegerType, StringType, DecimalType, FloatType

from project.sproc import create_fact_tables

def test_sproc(pytestconfig, session: Session):
    from patches import patch_to_timestamp  # patch for built-in function in transformer

    # Set up source table
    tbl = session.create_dataframe(
        data=[
            [1983, '2018-03-01 09:47:00.000 +0000', 551, 30958],
            [1988, '2018-03-01 09:47:01.000 +0000', 242, 19278],
            [1992, '2018-03-01 09:47:01.000 +0000', 768, 18461],
            [1980, '2018-03-01 09:47:03.000 +0000', 690, 15533],
            [1991, '2018-03-01 09:47:03.000 +0000', 490, 32449],
            [1959, '2018-03-01 09:47:04.000 +0000', 457, 29411],
            [1971, '2018-03-01 09:47:08.000 +0000', 279, 28015],
            [1964, '2018-03-01 09:47:09.000 +0000', 546, 15148],
            [1983, '2018-03-01 09:47:11.000 +0000', 358, 16967],
            [1985, '2018-03-01 09:47:12.000 +0000', 848, 20644],
            [1984, '2018-03-01 09:47:14.000 +0000', 295, 16365]
        ],
        schema=['BIRTH_YEAR', 'STARTTIME', 'TRIPDURATION',	'BIKEID'],
    )

    tbl.write.mode('overwrite').save_as_table(['CITIBIKE', 'PUBLIC', 'TRIPS_TEST'], mode='overwrite')

    # Expected values
    n_rows_expected = 12 
    bike_facts_expected = session.create_dataframe(
        data=[
            [30958, 1, 551.0, 40.0], 
            [19278, 1, 242.0, 35.0], 
            [18461, 1, 768.0, 31.0],
            [15533, 1, 690.0, 43.0], 
            [32449, 1, 490.0, 32.0], 
            [29411, 1, 457.0, 64.0], 
            [28015, 1, 279.0, 52.0], 
            [15148, 1, 546.0, 59.0], 
            [16967, 1, 358.0, 40.0], 
            [20644, 1, 848.0, 38.0], 
            [16365, 1, 295.0, 39.0]
        ],
        schema=StructType([
            StructField("BIKEID", IntegerType()), 
            StructField("COUNT", IntegerType()), 
            StructField("AVG_TRIPDURATION", FloatType()), 
            StructField("AVG_RIDER_AGE", FloatType())
        ])
    ).collect()

    month_facts_expected = session.create_dataframe(
        data=[['Mar', 11, 502.18182, 43.00000]],
        schema=StructType([
            StructField("MONTH", StringType()), 
            StructField("COUNT", IntegerType()), 
            StructField("AVG_TRIPDURATION", DecimalType()), 
            StructField("AVG_RIDER_AGE", DecimalType())
        ])
    ).collect()

    # Call sproc, get actual values
    n_rows_actual = create_fact_tables(session, 'TRIPS_TEST')
    bike_facts_actual = session.table(['CITIBIKE', 'PUBLIC', 'bike_facts']).collect() 
    month_facts_actual = session.table(['CITIBIKE', 'PUBLIC', 'month_facts']).collect() 

    # Comparisons
    assert n_rows_expected == n_rows_actual
    assert bike_facts_expected == bike_facts_actual


test/test_sproc.py
sfc-gh-jfreeberg/local-testing-ff
from snowflake.snowpark.session import Session
from snowflake.snowpark.types import StructType, StructField, IntegerType, FloatType
from project.transformers import add_rider_age, calc_bike_facts, calc_month_facts
from snowflake.snowpark.mock.functions import patch
from snowflake.snowpark.functions import to_timestamp
from snowflake.snowpark.mock.snowflake_data_type import ColumnEmulator

def test_add_rider_age(session: Session):
    input = session.create_dataframe(
        [
            [1980], 
            [1995], 
            [2000]
        ], 
        schema=StructType([StructField("BIRTH_YEAR", IntegerType())])
    )

    expected = session.create_dataframe(
        [
            [1980, 43], 
            [1995, 28], 
            [2000, 23]
        ],
        schema=StructType([StructField("BIRTH_YEAR", IntegerType()), StructField("RIDER_AGE", IntegerType())])
    )
    
    actual = add_rider_age(input)
    assert expected.collect() == actual.collect()


def test_calc_bike_facts(session: Session):
    input = session.create_dataframe([
            [1, 10, 20],
            [1, 5, 30],
            [2, 20, 50],
            [2, 10, 60]
        ], 
        schema=StructType([
            StructField("BIKEID", IntegerType()), 
            StructField("TRIPDURATION", IntegerType()), 
            StructField("RIDER_AGE", IntegerType())
        ])
    )

    expected = session.create_dataframe([
            [1, 2, 7.5, 25.0],
            [2, 2, 15.0, 55.0],
        ], 
        schema=StructType([
            StructField("BIKEID", IntegerType()), 
            StructField("COUNT", IntegerType()), 
            StructField("AVG_TRIPDURATION", FloatType()), 
            StructField("AVG_RIDER_AGE", FloatType())
        ])
    )

    actual = calc_bike_facts(input)
    assert expected.collect() == actual.collect()


def test_calc_month_facts(session: Session):
    from patches import patch_to_timestamp

    input = session.create_dataframe(
        data=[
            ['2018-03-01 09:47:00.000 +0000', 1, 10,  15],
            ['2018-03-01 09:47:14.000 +0000', 2, 20, 12],
            ['2018-04-01 09:47:04.000 +0000', 3, 6,  30]
        ],
        schema=['STARTTIME', 'BIKE_ID', 'TRIPDURATION', 'RIDER_AGE']
    )

    expected = session.create_dataframe(
        data=[
            ['Mar', 2, 15, 13.5],
            ['Apr', 1, 6, 30.0]
        ],
        schema=['MONTH', 'COUNT', 'AVG_TRIPDURATION', 'AVG_RIDER_AGE']
    )

    actual = calc_month_facts(input)

    assert expected.collect() == actual.collect()

test/test_transformers.py
sfc-gh-jfreeberg/local-testing-ff