Rahul Dubey

21st March, 2024

125

Streamlining Snowflake Data Loading: A Guide to DAGs and Zip File Management in GCP Composer

In this comprehensive tutorial, we'll delve into the intricacies of managing Zip files during data loading into Snowflake using Airflow within GCP Composer. Follow along with our step-by-step guide, complete with code snippets, to optimize your ETL processes seamlessly.

Case Study:
Picture yourself in an organization transitioning its data platform from BigQuery to Snowflake. With all data operations now centralized in Snowflake, your task is to construct a data pipeline for Attribution Analytics. The attribution data arrives in GCS buckets as Zip files from partners, requiring careful handling.

What You'll Need:
Before we begin, ensure you have the following prerequisites:

Google Cloud Platform Account
Snowflake Account
GCP Composer Environment with basic Snowflake integration set up
Code Editor
Python Installed

Step 1: Setting Up the Environment
Ensure that your GCP Composer environment is configured with basic integration with Snowflake. This includes establishing integration objects and external stages to facilitate direct loading from GCS buckets.

Step 2: Defining the DAG Script
Begin by defining the Directed Acyclic Graph (DAG) script using the DAG object from the airflow package. This script will serve as the foundation for orchestrating the data pipeline.

# Import necessary packages
from datetime import datetime, timedelta
from airflow import DAG
from airflow.operators.python_operator import PythonOperator
from airflow.contrib.operators.snowflake_operator import SnowflakeOperator
import pandas as pd
from google.cloud import storage
import zipfile
import io

# Define default arguments for the DAG
default_args = {
    'owner': 'ORGANIZATION',
    'start_date': datetime(2023, 2, 19),
    'email': ['[email protected]'],
    'email_on_failure': True,
    'email_on_retry': False
}

# Create the DAG object
dag = DAG('SFLK_ZIP_LOAD', description='This DAG loads ZIP files into Snowflake', max_active_runs=1, catchup=False, default_args=default_args)

Step 3: Defining Custom Tasks
Define custom tasks using Operators. For our pipeline, we'll have three main tasks: Truncate Task, Unzipping Task, and Loading Task.

# Define Truncate Task
TRUNC_QUERY = '''DELETE FROM <DATABASE_NAME.SCHEMA_NAME.TABLE_NAME> WHERE <DATE_FIELD> = CURRENT_DATE'''

trunc_task = SnowflakeOperator(
               task_id='TRUNCATE_TASK',
               sql=[TRUNC_QUERY],
               snowflake_conn_id='<connection_id>',
               database='<DATABASE_NAME>',
               schema='<SCHEMA_NAME>',
               warehouse='<DATAWAREHOUSE_NAME>',
               role='<ROLE_NAME>',
               dag=dag)

# Define Unzipping Task
def unzip_file_in_gcs(**context):
    # Define GCS Client parameters
    bucket_name = '<BUCKET_NAME>'
    file_name = '<FILE_NAME>.zip'

    # Connect to the GCS bucket
    storage_client = storage.Client()
    bucket = storage_client.bucket(bucket_name)
    blob = bucket.blob(file_name)

    # Download the zip file to memory
    zip_file_content = blob.download_as_string()

    # Unzip the file
    zip_file = zipfile.ZipFile(io.BytesIO(zip_file_content))
    zip_file.extractall(path='/home/airflow/gcs/data/temp/')

    # Upload each file in the zip to the GCS bucket
    with open('/home/airflow/gcs/data/temp/<FILE_NAME>.csv', 'rb') as f:
        file_content = f.read()
        new_blob = bucket.blob('<FILE_NAME>.csv')
        new_blob.upload_from_string(file_content)

unzip_task = PythonOperator(
        task_id="UNZIP_TASK",
        python_callable=unzip_file_in_gcs,
        provide_context=True,
        dag=dag
    )

# Define Loading Task
LOAD_QUERY = '''COPY INTO <DATABASE_NAME>.<SCHEMA_NAME>.<TABLE_NAME> FROM @<DATABASE_NAME>.<SCHEMA_NAME>.<STAGE_NAME>/<FILE_NAME>.csv
file_format = (format_name = <DATABASE_NAME>.<SCHEMA_NAME>.FF_CSV)'''

load_task = SnowflakeOperator(
        task_id='LOAD_TASK',
        sql=[LOAD_QUERY],
        snowflake_conn_id='<connection_id>',
        database='<DATABASE_NAME>',
        schema='<SCHEMA_NAME>',
        warehouse='<DATAWAREHOUSE_NAME>',
        role='<ROLE_NAME>',
        dag=dag)

Step 4: Task Flow
Link all tasks together to form the data pipeline.

# Define task flow
with dag:
  trunc_task >> unzip_task >> load_task

Step 5: Upload and Run the DAG Script
Upload the DAG script to the GCS bucket associated with the Composer environment. The Apache Airflow Web UI will automatically detect the new DAG and commence execution.

Conclusion:
In this tutorial, we've walked through the process of efficiently handling Zip files stored in a GCS bucket using Apache Airflow within GCP Composer, ultimately loading data into Snowflake tables. By following our guide and leveraging DAGs effectively, you can streamline your ETL processes with ease.

Rahul Dubey

Related Tutorial

Discover the Power of Go Programming: Boost Performance Without Sacrificing Readability.

Popular Tools for Developers

Popular Tools for Project Manager

Popular Tools for Youtubers

Other Tools for Daily Life

About us

Find Online Tools

Quick links

Social links

Other