Snowflake Query

This example shows how to use a Flyte SnowflakeTask to execute a query.

from flytekit import kwtypes, task, workflow
from flytekit.types.schema import FlyteSchema
from flytekitplugins.snowflake import SnowflakeConfig, SnowflakeTask

This is the world’s simplest query. Note that in order for registration to work properly, you’ll need to give your Snowflake task a name that’s unique across your project/domain for your Flyte installation.

snowflake_task_no_io = SnowflakeTask(
    name="sql.snowflake.no_io",
    inputs={},
    query_template="SELECT 1",
    output_schema_type=None,
    task_config=SnowflakeConfig(account="ha63105.us-central1.gcp", database="SNOWFLAKE_SAMPLE_DATA",
                                schema="TPCH_SF1000", warehouse="COMPUTE_WH"),
)


@workflow
def no_io_wf():
    return snowflake_task_no_io()

Of course, in real world applications we are usually more interested in using Snowflake to query a dataset. In this case we use SNOWFLAKE_SAMPLE_DATA which is default dataset in snowflake service. here The data is formatted according to this schema:

C_CUSTKEY (int)

C_NAME (string)

C_ADDRESS (string)

C_NATIONKEY (int)

C_PHONE (string)

C_ACCTBAL (float)

C_MKTSEGMENT (string)

C_COMMENT (string)

Let’s look out how we can parameterize our query to filter results for a specific country, provided as a user input specifying a nation key.

snowflake_task_templatized_query = SnowflakeTask(
    name="sql.snowflake.w_io",
    # Define inputs as well as their types that can be used to customize the query.
    inputs=kwtypes(nation_key=int),
    task_config=SnowflakeConfig(account="ha63105.us-central1.gcp", database="SNOWFLAKE_SAMPLE_DATA",
                                schema="TPCH_SF1000", warehouse="COMPUTE_WH"),
    query_template="SELECT * from CUSTOMER where C_NATIONKEY =  {{ .inputs.nation_key }} limit 100",
)


@workflow
def full_snowflake_wf(nation_key: int):
    return snowflake_task_templatized_query(nation_key=nation_key)

Check query result on snowflake console: https://<account>.snowflakecomputing.com/console#/monitoring/queries/detail

For example, https://ha63105.us-central1.gcp.snowflakecomputing.com/console#/monitoring/queries/detail