Configuring SQL Credentials

To securely access BigQuery and Snowflake databases in Flyte, you must configure credentials using secret references rather than hardcoding sensitive information. This guide shows how to set up BigQueryConfig and SnowflakeConfig with their respective task types.

Configuring BigQuery Credentials

To authenticate with BigQuery, you provide the name of a secret containing your Google Application Credentials JSON. Flyte uses this name to look up the actual credentials at runtime.

from flyteplugins.bigquery import BigQueryConfig, BigQueryTask
from flyte.io import DataFrame

# 1. Define the configuration
bq_config = BigQueryConfig(
    ProjectID="my-gcp-project",
    Location="US"
)

# 2. Reference the secret by name in the task
bigquery_task = BigQueryTask(
    name="bigquery_query",
    plugin_config=bq_config,
    query_template="SELECT * FROM `my-gcp-project.dataset.table` WHERE id = {{.inputs.id}}",
    inputs={"id": int},
    output_dataframe_type=DataFrame,
    # This is the NAME of the secret in your secret store
    google_application_credentials="my-gcp-secret"
)

Key Components

• BigQueryConfig: Holds metadata like ProjectID and Location. You can also pass a google.cloud.bigquery.QueryJobConfig to the QueryJobConfig attribute for advanced settings.
• google_application_credentials: This parameter in BigQueryTask must be the identifier for the secret, not the JSON content itself.

---

Configuring Snowflake Credentials

Snowflake tasks support key-pair authentication via secrets and alternative methods (like passwords) via connection arguments.

Key-Pair Authentication

When using key-pair authentication, you specify the secret keys for the private key and optional passphrase. Flyte automatically generates the expected environment variable names based on a secret_group.

from flyteplugins.snowflake import Snowflake, SnowflakeConfig
import flyte

# 1. Define connection metadata
sf_config = SnowflakeConfig(
    account="account-id",
    database="FLYTE_DB",
    schema="PUBLIC",
    warehouse="COMPUTE_WH",
    user="FLYTE_USER",
)

# 2. Configure the task with secret references
snowflake_task = Snowflake(
    name="snowflake_query",
    plugin_config=sf_config,
    query_template="SELECT * FROM MY_TABLE WHERE ID = %(id)s",
    inputs={"id": int},
    # Secret references
    snowflake_private_key="snowflake-pk",
    snowflake_private_key_passphrase="snowflake-pass",
    secret_group="snowflake_creds"
)

Secret Naming Convention

The Snowflake task generates environment variable names using the pattern {secret_group}_{key}, converted to uppercase with hyphens replaced by underscores.

In the example above:

• snowflake_private_key ("snowflake-pk") with secret_group ("snowflake_creds") expects an environment variable named: SNOWFLAKE_CREDS_SNOWFLAKE_PK.
• If secret_group is omitted, it defaults to the key name (e.g., SNOWFLAKE_PK).

Alternative Authentication (Password/Role)

For authentication methods other than key-pair (like password-based auth) or to specify a specific role, use the connection_kwargs attribute in SnowflakeConfig.

from flyteplugins.snowflake import SnowflakeConfig

sf_config = SnowflakeConfig(
    account="account-id",
    database="FLYTE_DB",
    schema="PUBLIC",
    warehouse="COMPUTE_WH",
    user="FLYTE_USER",
    connection_kwargs={
        "password": "my-password", # Note: Use Flyte Secrets for production
        "role": "MY_ROLE",
        "authenticator": "externalbrowser"
    }
)

Troubleshooting Secret Resolution

Environment Variable Mismatches

If your Snowflake task fails to authenticate, verify that the environment variables in your execution environment match the generated names. You can explicitly define these in a TaskEnvironment:

import flyte

# Explicitly mapping a secret to the expected environment variable name
env = flyte.TaskEnvironment(
    name="snowflake_env",
    secrets=[
        flyte.Secret(
            key="snowflake-pk", 
            as_env_var="SNOWFLAKE_CREDS_SNOWFLAKE_PK"
        )
    ],
)

Query Formatting

Both BigQueryTask and Snowflake automatically normalize your query_template by:

1. Replacing newlines (\n) and tabs (\t) with spaces.
2. Collapsing multiple consecutive spaces into a single space.
3. Stripping leading and trailing whitespace.

Ensure your query remains valid SQL after these transformations (e.g., avoid relying on SQL comments that require newlines).