Migrating from Redshift to Snowflake

Narrativ handles a pretty good amount of data, and the rate that we collect it increases every day as we add more customers. This is great- after all, collecting data is a key part of our business! However, it can be…challenging to keep acceptable performance on data systems as the size gets large.

We have been using Redshift for all of our stats data. Redshift has surprised us on multiple occasions with how well it handles some of our complex queries over terabytes of data- the implementation of window functions for one is extremely fast. On the other hand, it can be expensive. You have to manage a cluster with a fixed amount of disk space, and when the disk space gets close to filling up, performance actually suffers. We were able to offload older data to Spectrum (an external schema attachment to Redshift that lets you query data at rest on S3 — see our tool Spectrify), but that causes problems too. Now users have to remember which data is in the live set and which is in the cold set, and add unions to many of their existing queries to hit the whole data set. This also required us to spend some time manually doing maintenance to move data around every couple months.

With all of that considered, we decided to migrate to Snowflake. The biggest advantage for us was that Snowflake lets you separate storage from compute- but everything else worked pretty much the same. This way if we want to dump 100 TB of data into a single table, we can do that without having to keep track of disk space, and we can easily scale compute to the queries we are running instead of having it be linked to the size of the data. Helpfully, there is a SQLAlchemy dialect and just about every Redshift query we had worked out of the box.

The interesting part that I wanted to share is something called a “Snowpipe.” This is a built in setting in Snowflake that lets you set up automatic trickle loading from an S3 bucket directly to a Snowflake table. You set up a notification on your S3 bucket, and each time a file gets added, Snowflake automatically imports it. We took it a step further and added a helper to our Alembic migrations to make sure it worked automatically going forward. This was really great since this was already how we were loading Redshift (though we had to write the code to do it in Redshift ourselves).

First you create a table using regular Alembic constructs, and then you add some extra Snowflake specific stuff.

def setup_table(op, tablename, columns):
    op.create_table(
        tablename,
        *columns
    )

    if current_config().STATS_DB == 'snowflake':
        build_stage(op, tablename)

        # The snowpipe will watch s3 for file changes and import each new file as it arrives.
        build_snowpipe(op, tablename, column_names=[col.name for col in columns])

        # We need to know which SQS queue snowflake will be listening on for new files, and
        # then add a listeneter to that queue
        sqs_arn = get_snowpipe_arn(op, tablename)
        add_sqs_listener(sqs_arn, tablename)

After you create a table, the first thing you have to set up is a “Stage” which tells Snowflake where files are located in S3, and a file format. Our enrichment pipeline make newline delimited JSON.

def initial_setup(op):
    op.execute("""
        CREATE OR REPLACE FILE FORMAT STD_JSON
        TYPE = JSON
        COMPRESSION = 'GZIP'
        ENABLE_OCTAL = FALSE
        ALLOW_DUPLICATE = FALSE
        STRIP_OUTER_ARRAY = FALSE
        STRIP_NULL_VALUES = FALSE
        IGNORE_UTF8_ERRORS = FALSE
    """)
    
def build_stage(op, tablename):
    op.execute("""
        CREATE OR REPLACE STAGE PUBLIC.{tablename}_STAGE
        URL = 'S3://{bucket}/snowflake/{tablename}/json/'
        CREDENTIALS = (
            AWS_KEY_ID = '{key_id}' 
            AWS_SECRET_KEY='{secret_key}')
        FILE_FORMAT=(FORMAT_NAME = 'STD_JSON' TYPE=JSON)
    """.format(
        bucket=BUCKET, 
        tablename=tablename, 
        key_id=current_config().SNOWFLAKE_AWS_KEY_ID,
        secret_key=current_config().SNOWFLAKE_AWS_SECRET_KEY))

Once the stage is created, the next step is to create the Pipe. This is the code that defines the transformation- copy files from THIS stage to THIS table with THIS command.

def build_snowpipe(op, tablename, column_names):
    op.execute("""
        CREATE OR REPLACE PIPE {tablename}
        AUTO_INGEST=TRUE
        AS COPY INTO {tablename}(
            {into_columns}
        ) FROM (
            SELECT
                {from_columns}
            FROM @PUBLIC.{tablename}_STAGE
        )
        FILE_FORMAT = STD_JSON
        ON_ERROR = CONTINUE
    """.format(**{
        'into_columns': ',\n'.join(['{}'.format(col) for col in column_names]),
        # you have to escape the : character in sqlalchemy queries!
        'from_columns': ',\n'.join(['PARSE_JSON($1)\:{}'.format(col) for col in column_names]),
        'tablename': tablename,
    }))

The reason that the column names are necessary is that, by default, Snowflake will import your JSON document into an unstructured format. This is fine for some things, but we really needed the added performance of a structured table. This transforms the rows into the correct schema during the import process. Now, the final step is to setup the notification so that Snowflake knows to run this automatically-

def get_snowpipe_arn(op, tablename):
    result = op.get_bind().execute("""
        SELECT
            NOTIFICATION_CHANNEL_NAME
        FROM INFORMATION_SCHEMA.pipes
        WHERE PIPE_NAME = '{tablename}'
    """.format(tablename=tablename.upper()))
    return result.fetchone()[0]


def add_sqs_listener(sqs_arn, tablename):
    notification_config = {
        'Events': ['s3:ObjectCreated:*'],
        'Id': tablename,
        'QueueArn': sqs_arn,
        'Filter': {'Key': {'FilterRules': [{'Name': 'prefix', 'Value': 'snowflake/{}/json'.format(tablename)}]}}
    }

    client = boto3.client('s3')
    existing_configs = client.get_bucket_notification_configuration(Bucket=BUCKET).get('QueueConfigurations', [])
    try:
        index = [config['Id'] for config in existing_configs].index(tablename)
        existing_configs[index] = notification_config
    except ValueError:
        existing_configs.append(notification_config)

    client.put_bucket_notification_configuration(
        Bucket=BUCKET,
        NotificationConfiguration={'QueueConfigurations': existing_configs})

This uses boto to set up the notification to go to the SNS queue that Snowflake created for the Snowpipe in the previous step. Once this step is setup, trickle loading will now happen automatically! The last thing is that you have to keep the Pipeline correct as you add or delete columns, so we wrote that as well-

def add_column(op, tablename, column):
    op.add_column(tablename, column)
    if current_config().STATS_DB == 'snowflake':
        # replace the snowpipe query with one with the new column name
        result = op.get_bind().execute("""
            SELECT * FROM {tablename} WHERE FALSE
        """.format(tablename=tablename.upper()))
        build_snowpipe(op, tablename, result.keys())


def drop_column(op, tablename, column_name):
    if current_config().STATS_DB == 'snowflake':
        # replace the snowpipe query with one without the removed column name
        result = op.get_bind().execute("""
            SELECT * FROM {tablename} WHERE FALSE
        """.format(tablename=tablename.upper()))
        column_names_without_removed_column = filter(lambda x: x != column_name, result.keys())
        build_snowpipe(op, tablename, column_names_without_removed_column)
    op.drop_column(tablename, column_name)

Once we had this set up, simply moving our Redshift migrations over to run on our Snowflake cluster got all of our data tables importing correctly and automatically- and lets us stand down the infrastructure that we used to use to load Redshift!

We have learned some other tricks on getting data from Redshift to Snowflake that we will share in future posts. Good luck!