Snowflake Integration

Run a Snowflake query as a Flyte task

To submit parameterized SQL, batch inserts, or a query whose result is a DataFrame, install the plugin and construct a Snowflake task with a SnowflakeConfig:

pip install flyteplugins-snowflake

import pandas as pd
from flyteplugins.snowflake.task import Snowflake, SnowflakeConfig

config = SnowflakeConfig(
    account="myorg-myaccount",
    user="flyte_user",
    database="ANALYTICS",
    schema="PUBLIC",
    warehouse="COMPUTE_WH",
)

lookup = Snowflake(
    name="lookup_user",
    query_template="SELECT * FROM users WHERE id = %(user_id)s",
    plugin_config=config,
    inputs={"user_id": int},
    output_dataframe_type=pd.DataFrame,
    snowflake_private_key="snowflake-pk",
)

SnowflakeConfig is a dataclass containing the required connection metadata: account, database, schema, warehouse, and user. Its optional connection_kwargs is copied into the serialized task configuration and can carry additional Snowflake Connector parameters.

Importing flyteplugins.snowflake exports Snowflake, SnowflakeConfig, and SnowflakeConnector. The import also registers the pandas-to-Snowflake and Snowflake-to-pandas DataFrame handlers. Application code normally creates Snowflake tasks rather than instantiating SnowflakeConnector; the connector is registered for the snowflake task type at import time.

Configure authentication with Flyte secrets

For key-pair authentication, pass the Flyte secret-store key names—not the private-key contents—to snowflake_private_key and, when needed, snowflake_private_key_passphrase:

from flyteplugins.snowflake.task import Snowflake, SnowflakeConfig

config = SnowflakeConfig(
    account="myorg-myaccount",
    user="flyte_user",
    database="ANALYTICS",
    schema="PUBLIC",
    warehouse="COMPUTE_WH",
)

task = Snowflake(
    name="my_task",
    query_template="SELECT 1",
    plugin_config=config,
    snowflake_private_key="private-key",
    snowflake_private_key_passphrase="passphrase",
)

Snowflake.custom_config() serializes the secret references under secrets. Snowflake._to_env_var() creates each environment-variable name by joining secret_group and the key with an underscore (when a group is present), replacing hyphens with underscores, and uppercasing the result. For example, the task above serializes PRIVATE_KEY and PASSPHRASE; with secret_group="snowflake", the names become SNOWFLAKE_PRIVATE_KEY and SNOWFLAKE_PASSPHRASE.

The shipped example creates a TaskEnvironment containing the Snowflake tasks, installs flyteplugins-snowflake in its image, and maps the Flyte secret key snowflake to the environment variable SNOWFLAKE_PRIVATE_KEY. Configure the equivalent flyte.TaskEnvironment secret mapping for a DataFrame-producing task so that the decoder's fixed environment-variable name is available at runtime.

At connector execution time, the resolved secret values are supplied to SnowflakeConnector.create, get, and delete. _get_snowflake_connection() converts PEM private-key text to private-key bytes with _crypto.get_private_key() and passes the result as private_key to snowflake.connector.connect. The optional passphrase is used while parsing the key.

The DataFrame decoder has a separate, fixed environment-variable convention: its _get_connection() reads SNOWFLAKE_PRIVATE_KEY and SNOWFLAKE_PRIVATE_KEY_PASSPHRASE. If a task returns a DataFrame, map the Flyte secrets to those names when configuring the runtime, as in the shipped example.

For password, OAuth, role, authenticator, session parameters, or other Snowflake Connector options, put the values in connection_kwargs instead:

from flyteplugins.snowflake.task import SnowflakeConfig

config = SnowflakeConfig(
    account="myorg-myaccount",
    user="flyte_user",
    database="ANALYTICS",
    schema="PUBLIC",
    warehouse="COMPUTE_WH",
    connection_kwargs={"password": "...", "role": "ADMIN"},
)

These entries are forwarded as keyword arguments to snowflake.connector.connect. SnowflakeJobMetadata also persists connection_kwargs, so the connector can use the same connection parameters during status polling and cancellation.

Use scalar parameters

Declare each named SQL parameter in inputs with its Python type, and use the matching %(name)s placeholder in query_template:

import pandas as pd
from flyteplugins.snowflake.task import Snowflake

lookup = Snowflake(
    name="lookup_user",
    query_template="SELECT * FROM users WHERE id = %(user_id)s",
    plugin_config=config,
    inputs={"user_id": int},
    output_dataframe_type=pd.DataFrame,
    snowflake_private_key="snowflake-pk",
)

The task constructor builds a NativeInterface from inputs and declares a results output only when output_dataframe_type is not None. It normalizes whitespace in the SQL template, and sql() emits the normalized statement as an ANSI tasks_pb2.Sql value. During submission, scalar inputs are passed unchanged to Snowflake's cursor.execute_async(query, inputs); the placeholder names therefore need to match the input dictionary keys.

Insert multiple rows with batch expansion

Set batch=True, declare list inputs, and put one named-placeholder row inside exactly one VALUES (...) clause:

from flyteplugins.snowflake.task import Snowflake, SnowflakeConfig

sf_config = SnowflakeConfig(
    user="KEVIN",
    account="PWGJLTH-XKB21544",
    database="FLYTE",
    schema="PUBLIC",
    warehouse="COMPUTE_WH",
)

snowflake_insert_task = Snowflake(
    name="snowflake",
    inputs={"id": list[int], "name": list[str], "age": list[int]},
    plugin_config=sf_config,
    query_template=(
        "INSERT INTO FLYTE.PUBLIC.TEST (ID, NAME, AGE) "
        "VALUES (%(id)s, %(name)s, %(age)s);"
    ),
    snowflake_private_key="snowflake",
    batch=True,
)

When the task receives id=[1, 2] and name=["Alice", "Bob"], _expand_batch_query() produces the corresponding numbered placeholders and flattened parameter dictionary:

from flyteplugins.snowflake.connector import _expand_batch_query

query = "INSERT INTO t (id, name) VALUES (%(id)s, %(name)s)"
inputs = {"id": [1, 2], "name": ["Alice", "Bob"]}

expanded, params = _expand_batch_query(query, inputs)

assert expanded == (
    "INSERT INTO t (id, name) VALUES "
    "(%(id_0)s, %(name_0)s), (%(id_1)s, %(name_1)s)"
)
assert params == {"id_0": 1, "name_0": "Alice", "id_1": 2, "name_1": "Bob"}

The connector applies this transformation only when both serialized batch is true and inputs is truthy. Every input value must be a non-empty list of the same length. The query must contain exactly one case-insensitive VALUES (...) clause, and its placeholders must match the input keys. Unequal lengths, empty lists, no VALUES clause, or multiple VALUES clauses raise ValueError. A batch task invoked without inputs follows the ordinary execute_async(query, inputs) path rather than creating an empty batch.

Understand the connector lifecycle

Flyte dispatches the serialized task type snowflake to SnowflakeConnector. Its lifecycle is asynchronous, while blocking Snowflake Connector operations run in thread executors.

1. Create. SnowflakeConnector.create() reads the task custom configuration, requires account, user, database, and warehouse, and defaults a missing serialized schema to PUBLIC. It obtains a cached connection, calls cursor.execute_async(), reads cursor.sfqid, closes the cursor, and returns SnowflakeJobMetadata.
2. Get. SnowflakeConnector.get() reconnects from that metadata and calls get_query_status_throw_if_error(query_id). The Snowflake status name is passed to convert_to_flyte_phase. Status errors become a failed Flyte Resource containing the exception text.
3. Delete. SnowflakeConnector.delete() executes SELECT SYSTEM$CANCEL_QUERY('{query_id}'), then closes both cursor and connection.

The metadata object carries the account, user, database, schema, warehouse, query ID, output flag, and connection kwargs between these calls. A task has an output when its serialized interface contains at least one output variable; Snowflake uses the output name results.

Each get() result includes a ready Snowflake Dashboard dashboard log link. For account="myorg-myaccount" and query_id="query-abc", _construct_query_link() produces:

https://app.snowflake.com/myorg/myaccount/#/compute/history/queries/query-abc/detail

An account without a hyphen is used as the account locator in the URL. The link helper splits a hyphenated account at the first hyphen.

Return query results as a DataFrame

Declare the output type when a SELECT should be consumed as a DataFrame:

import pandas as pd
from flyteplugins.snowflake.task import Snowflake

snowflake_select_task = Snowflake(
    name="snowflake_select",
    output_dataframe_type=pd.DataFrame,
    plugin_config=sf_config,
    query_template="SELECT * FROM FLYTE.PUBLIC.TEST;",
    snowflake_private_key="snowflake",
)

The connector does not fetch rows in get(). After a successful query with a declared output, it returns results as a Flyte DataFrame whose URI has this layout:

snowflake://{user}/{account}/{warehouse}/{database}/{schema}/{query_id}

The registered SnowflakeToPandasDecodingHandler later decodes that URI. Its _read_from_sf() parses the connection metadata and query ID, calls get_results_from_sfqid(query_id), and returns fetch_pandas_all().

The shipped end-to-end pattern awaits the batch insert, awaits the DataFrame-producing SELECT, and passes the resulting pandas DataFrame to a downstream task:

@env.task
async def downstream_task(df: pd.DataFrame) -> float:
    mean_age = df["AGE"].mean()
    return mean_age.item()


@env.task
async def main(ids: list[int], names: list[str], ages: list[int]) -> float:
    await snowflake_insert_task(id=ids, name=names, age=ages)
    df = await snowflake_select_task()
    return await downstream_task(df)

Configuration and operational details

Setting	Behavior
SnowflakeConfig.account	Required task connection metadata and input to dashboard-link construction.
SnowflakeConfig.database, schema, warehouse, user	Required fields on SnowflakeConfig; the connector additionally defaults a missing serialized schema to PUBLIC.
SnowflakeConfig.connection_kwargs	Optional Snowflake Connector keyword arguments, persisted in job metadata for later lifecycle calls.
Snowflake.output_dataframe_type	When set, declares the results output and causes successful get() to return a result URI.
Snowflake.batch	Defaults to False; when true, enables the single-VALUES list expansion path.
Snowflake.secret_group	Prefixes generated secret environment-variable names; hyphens become underscores and names are uppercased.
snowflake_private_key / snowflake_private_key_passphrase	Flyte secret-store key names for PEM key contents and its optional passphrase.

_get_snowflake_connection() is decorated with async_lru.alru_cache; its connection arguments, including private-key content and passphrase, participate in the cache key. It creates the blocking Snowflake connection in an executor. create() and get() use the cached connection without explicitly closing it; delete() closes the connection after cancellation.

Troubleshooting

• Missing connection metadata: create() raises ValueError("Missing Snowflake account. Set it through task configuration.") when the account is absent. Missing user, database, or warehouse raises ValueError("User, database and warehouse must be specified in the task configuration.").
• Batch validation failures: Check that all inputs values are non-empty lists of equal length, the SQL has one VALUES (...) clause, and every placeholder is written as %(key)s for an input key.
• No DataFrame output: output_dataframe_type controls whether results is declared. Without it, a successful connector poll has no outputs, even if the SQL itself produces rows.
• DataFrame credential lookup: The decoder reads the fixed SNOWFLAKE_PRIVATE_KEY and optional SNOWFLAKE_PRIVATE_KEY_PASSPHRASE environment variables. A task's custom secret_group does not change those names in dataframe.py.
• Status failures: Exceptions raised by get_query_status_throw_if_error() are caught by get() and returned as a failed Resource with the dashboard link. Connection/setup failures outside that status call can still propagate.
• Testing expectations: The Snowflake connector tests mock connections and cursors, including execute_async, status polling, output URI creation, and cancellation. They do not establish a live Snowflake connection, so actual authentication, server status behavior, and SQL execution still require validation against Snowflake.