flyteplugins-snowflake 2.0.0__py3-none-any.whl

flyteplugins/snowflake/__init__.py

@@ -0,0 +1,58 @@
+"""
+Key features:
+
+- Parameterized SQL queries with typed inputs
+- Key-pair and password-based authentication
+- Returns query results as DataFrames
+- Automatic links to the Snowflake query dashboard in the Flyte UI
+- Query cancellation on task abort
+
+Basic usage example:
+```python
+import flyte
+from flyte.io import DataFrame
+from flyteplugins.snowflake import Snowflake, SnowflakeConfig
+
+config = SnowflakeConfig(
+    account="myorg-myaccount",
+    user="flyte_user",
+    database="ANALYTICS",
+    schema="PUBLIC",
+    warehouse="COMPUTE_WH",
+)
+
+count_users = Snowflake(
+    name="count_users",
+    query_template="SELECT COUNT(*) FROM users",
+    plugin_config=config,
+    output_dataframe_type=DataFrame,
+)
+
+flyte.TaskEnvironment.from_task("snowflake_env", count_users)
+
+if __name__ == "__main__":
+    flyte.init_from_config()
+
+    # Run locally (connector runs in-process, requires credentials and packages locally)
+    run = flyte.with_runcontext(mode="local").run(count_users)
+
+    # Run remotely (connector runs on the control plane)
+    run = flyte.with_runcontext(mode="remote").run(count_users)
+
+    print(run.url)
+```
+"""
+
+from flyte.io._dataframe.dataframe import DataFrameTransformerEngine
+
+from flyteplugins.snowflake.connector import SnowflakeConnector
+from flyteplugins.snowflake.dataframe import (
+    PandasToSnowflakeEncodingHandlers,
+    SnowflakeToPandasDecodingHandler,
+)
+from flyteplugins.snowflake.task import Snowflake, SnowflakeConfig
+
+DataFrameTransformerEngine.register(PandasToSnowflakeEncodingHandlers())
+DataFrameTransformerEngine.register(SnowflakeToPandasDecodingHandler())
+
+__all__ = ["Snowflake", "SnowflakeConfig", "SnowflakeConnector"]

flyteplugins/snowflake/connector.py

@@ -0,0 +1,404 @@
+import asyncio
+import re
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional, Tuple
+
+from async_lru import alru_cache
+from flyte import logger
+from flyte.connectors import AsyncConnector, ConnectorRegistry, Resource, ResourceMeta
+from flyte.connectors.utils import convert_to_flyte_phase
+from flyte.io import DataFrame
+from flyteidl2.core.execution_pb2 import TaskExecution, TaskLog
+from flyteidl2.core.tasks_pb2 import TaskTemplate
+from google.protobuf import json_format
+
+from snowflake.connector import SnowflakeConnection
+from snowflake.connector import connect as snowflake_connect
+
+TASK_TYPE = "snowflake"
+
+
+@dataclass
+class SnowflakeJobMetadata(ResourceMeta):
+    """
+    Metadata for a Snowflake query job.
+
+    Attributes:
+        account: Snowflake account identifier.
+        user: Snowflake user name.
+        database: Snowflake database name.
+        schema: Snowflake schema name.
+        warehouse: Snowflake warehouse name.
+        query_id: Unique identifier for the submitted query.
+        has_output: Indicates if the query produces output.
+        connection_kwargs: Additional connection parameters.
+    """
+
+    account: str
+    user: str
+    database: str
+    schema: str
+    warehouse: str
+    query_id: str
+    has_output: bool
+    connection_kwargs: Optional[Dict[str, Any]] = None
+
+
+def _get_private_key(private_key_content: str, private_key_passphrase: Optional[str] = None) -> bytes:
+    """
+    Decode the private key from the secret and return it in DER format.
+
+    Args:
+        private_key_content: The private key content in PEM format.
+        private_key_passphrase: The passphrase for the private key, if any.
+
+    Returns:
+        The private key in DER format.
+    """
+    from cryptography.hazmat.backends import default_backend
+    from cryptography.hazmat.primitives import serialization
+
+    private_key_bytes = private_key_content.strip().encode()
+    password = private_key_passphrase.encode() if private_key_passphrase else None
+
+    private_key = serialization.load_pem_private_key(
+        private_key_bytes,
+        password=password,
+        backend=default_backend(),
+    )
+
+    return private_key.private_bytes(
+        encoding=serialization.Encoding.DER,
+        format=serialization.PrivateFormat.PKCS8,
+        encryption_algorithm=serialization.NoEncryption(),
+    )
+
+
+@alru_cache
+async def _get_snowflake_connection(
+    account: str,
+    user: str,
+    database: str,
+    schema: str,
+    warehouse: str,
+    private_key_content: Optional[str] = None,
+    private_key_passphrase: Optional[str] = None,
+    **connection_kwargs,
+) -> SnowflakeConnection:
+    """
+    Create and return a Snowflake connection.
+
+    Supports private key authentication (recommended) and other auth methods via `connection_kwargs`.
+    See: https://docs.snowflake.com/en/developer-guide/python-connector/python-connector-api
+
+    Args:
+        account: Snowflake account identifier.
+        user: Snowflake user name.
+        database: Snowflake database name.
+        schema: Snowflake schema name.
+        warehouse: Snowflake warehouse name.
+        private_key_content: The private key content in PEM format, if any.
+        private_key_passphrase: The passphrase for the private key, if any.
+        **connection_kwargs: Additional connection parameters.
+
+    Returns:
+        An active Snowflake connection.
+    """
+
+    def _create_connection():
+        connection_params = {
+            "account": account,
+            "user": user,
+            "database": database,
+            "schema": schema,
+            "warehouse": warehouse,
+            **connection_kwargs,
+        }
+
+        # Add private key authentication if provided
+        if private_key_content:
+            private_key = _get_private_key(private_key_content, private_key_passphrase)
+            connection_params["private_key"] = private_key
+
+        # Let Snowflake connector validate authentication requirements
+        return snowflake_connect(**connection_params)
+
+    loop = asyncio.get_running_loop()
+    return await loop.run_in_executor(None, _create_connection)
+
+
+def _construct_query_link(account: str, query_id: str) -> str:
+    """Construct a Snowflake console link for the query.
+
+    Args:
+        account: Snowflake account identifier.
+        query_id: Unique identifier for the submitted query.
+
+    Returns:
+        URL to the Snowflake query dashboard for the given query ID.
+    """
+    if "-" in account:
+        parts = account.split("-", 1)
+        if len(parts) == 2:
+            org_name, account_name = parts
+            base_url = f"https://app.snowflake.com/{org_name}/{account_name}"
+        else:
+            base_url = f"https://app.snowflake.com/{account}"
+    else:
+        # Simple account locator
+        base_url = f"https://app.snowflake.com/{account}"
+
+    return f"{base_url}/#/compute/history/queries/{query_id}/detail"
+
+
+def _expand_batch_query(query: str, inputs: Dict[str, List[Any]]) -> Tuple[str, Dict[str, Any]]:
+    """
+    Expand a parameterized INSERT query with list inputs into a multi-row VALUES statement.
+
+    Example:
+        query = "INSERT INTO t (a, b) VALUES (%(a)s, %(b)s)"
+        inputs = {"a": [1, 2], "b": ["x", "y"]}
+
+    Produces:
+        "INSERT INTO t (a, b) VALUES (%(a_0)s, %(b_0)s), (%(a_1)s, %(b_1)s)"
+        flat_params = {"a_0": 1, "b_0": "x", "a_1": 2, "b_1": "y"}
+
+    Constraints:
+    - Query must contain exactly one VALUES (...) clause
+    - All input lists must have the same non-zero length
+    """
+    keys = list(inputs.keys())
+
+    lengths = {len(v) for v in inputs.values()}
+    if len(lengths) != 1:
+        raise ValueError("All batch input lists must have the same length.")
+
+    n_rows = lengths.pop()
+    if n_rows == 0:
+        raise ValueError("Batch inputs must contain at least one row.")
+
+    # Find all VALUES (...) clauses (case-insensitive)
+    # The inner group handles %(...)s placeholders which contain literal parentheses
+    matches = list(re.finditer(r"VALUES\s*\(((?:[^()]*|\([^)]*\))*)\)", query, re.IGNORECASE))
+    if not matches:
+        raise ValueError("Batch inputs require a query template with a VALUES (...) clause.")
+    if len(matches) > 1:
+        raise ValueError("Batch query expansion supports exactly one VALUES (...) clause.")
+
+    match = matches[0]
+    values_content = match.group(1)
+
+    value_rows = []
+    flat_params: Dict[str, Any] = {}
+
+    for i in range(n_rows):
+        row_content = values_content
+        for key in keys:
+            row_content = re.sub(
+                rf"%\({re.escape(key)}\)s",
+                f"%({key}_{i})s",
+                row_content,
+            )
+            flat_params[f"{key}_{i}"] = inputs[key][i]
+        value_rows.append(f"({row_content})")
+
+    before = query[: match.start()]
+    after = query[match.end() :]
+
+    expanded_query = before + "VALUES " + ", ".join(value_rows) + after
+    return expanded_query, flat_params
+
+
+class SnowflakeConnector(AsyncConnector):
+    name = "Snowflake Connector"
+    task_type_name = TASK_TYPE
+    metadata_type = SnowflakeJobMetadata
+
+    async def create(
+        self,
+        task_template: TaskTemplate,
+        inputs: Optional[Dict[str, Any]] = None,
+        snowflake_private_key: Optional[str] = None,
+        snowflake_private_key_passphrase: Optional[str] = None,
+        **kwargs,
+    ) -> SnowflakeJobMetadata:
+        """
+        Submit a query to Snowflake asynchronously.
+
+        Args:
+            task_template: The Flyte task template containing the SQL query and configuration.
+            inputs: Optional dictionary of input parameters for parameterized queries.
+            snowflake_private_key: The private key content set as a Flyte secret.
+            snowflake_private_key_passphrase: The passphrase for the private key set as a Flyte secret, if any.
+
+        Returns:
+            A SnowflakeJobMetadata object containing the query ID and link to the query dashboard.
+        """
+        custom = json_format.MessageToDict(task_template.custom) if task_template.custom else {}
+
+        account = custom.get("account")
+        if not account:
+            raise ValueError("Missing Snowflake account. Set it through task configuration.")
+
+        user = custom.get("user")
+        database = custom.get("database")
+        schema = custom.get("schema", "PUBLIC")
+        warehouse = custom.get("warehouse")
+
+        if not all([user, database, warehouse]):
+            raise ValueError("User, database and warehouse must be specified in the task configuration.")
+
+        # Get additional connection parameters from custom config
+        connection_kwargs = custom.get("connection_kwargs", {})
+
+        conn = await _get_snowflake_connection(
+            account=account,
+            user=user,
+            database=database,
+            schema=schema,
+            warehouse=warehouse,
+            private_key_content=snowflake_private_key,
+            private_key_passphrase=snowflake_private_key_passphrase,
+            **connection_kwargs,
+        )
+
+        query = task_template.sql.statement
+        batch = custom.get("batch", False)
+
+        def _execute_query():
+            cursor = conn.cursor()
+
+            if batch and inputs:
+                expanded_query, flat_params = _expand_batch_query(query, inputs)
+                cursor.execute_async(expanded_query, flat_params)
+            else:
+                cursor.execute_async(query, inputs)
+
+            query_id = cursor.sfqid
+            cursor.close()
+
+            return query_id
+
+        loop = asyncio.get_running_loop()
+        query_id = await loop.run_in_executor(None, _execute_query)
+
+        logger.info(f"Snowflake query submitted with ID: {query_id}")
+
+        return SnowflakeJobMetadata(
+            account=account,
+            user=user,
+            database=database,
+            schema=schema,
+            warehouse=warehouse,
+            query_id=query_id,
+            has_output=task_template.interface.outputs is not None
+            and len(task_template.interface.outputs.variables) > 0,
+            connection_kwargs=connection_kwargs,
+        )
+
+    async def get(
+        self,
+        resource_meta: SnowflakeJobMetadata,
+        snowflake_private_key: Optional[str] = None,
+        snowflake_private_key_passphrase: Optional[str] = None,
+        **kwargs,
+    ) -> Resource:
+        """
+        Poll the status of a Snowflake query.
+
+        Args:
+            resource_meta: The SnowflakeJobMetadata containing the query ID.
+            snowflake_private_key: The private key content set as a Flyte secret.
+            snowflake_private_key_passphrase: The passphrase for the private key set as a Flyte secret, if any.
+
+        Returns:
+            A Resource object containing the query results and a link to the query dashboard.
+        """
+        conn = await _get_snowflake_connection(
+            account=resource_meta.account,
+            user=resource_meta.user,
+            database=resource_meta.database,
+            schema=resource_meta.schema,
+            warehouse=resource_meta.warehouse,
+            private_key_content=snowflake_private_key,
+            private_key_passphrase=snowflake_private_key_passphrase,
+            **(resource_meta.connection_kwargs or {}),
+        )
+
+        log_link = TaskLog(
+            uri=_construct_query_link(resource_meta.account, resource_meta.query_id),
+            name="Snowflake Dashboard",
+            ready=True,
+            link_type=TaskLog.DASHBOARD,
+        )
+
+        def _get_query_status():
+            try:
+                status = conn.get_query_status_throw_if_error(resource_meta.query_id)
+                return status, None
+            except Exception as e:
+                return None, str(e)
+
+        loop = asyncio.get_running_loop()
+        status, error = await loop.run_in_executor(None, _get_query_status)
+
+        if error:
+            logger.error(f"Snowflake query failed: {error}")
+            return Resource(phase=TaskExecution.FAILED, message=error, log_links=[log_link])
+
+        # Map Snowflake status to Flyte phase
+        # Snowflake statuses: RUNNING, SUCCESS, FAILED_WITH_ERROR, ABORTING, etc.
+        cur_phase = convert_to_flyte_phase(status.name)
+        outputs = None
+
+        if cur_phase == TaskExecution.SUCCEEDED and resource_meta.has_output:
+            # Construct the output URI for the results
+            output_location = (
+                f"snowflake://{resource_meta.user}/{resource_meta.account}/{resource_meta.warehouse}/"
+                f"{resource_meta.database}/{resource_meta.schema}/{resource_meta.query_id}"
+            )
+            outputs = {"results": DataFrame(uri=output_location)}
+
+        return Resource(phase=cur_phase, message=status.name, log_links=[log_link], outputs=outputs)
+
+    async def delete(
+        self,
+        resource_meta: SnowflakeJobMetadata,
+        snowflake_private_key: Optional[str] = None,
+        snowflake_private_key_passphrase: Optional[str] = None,
+        **kwargs,
+    ):
+        """
+        Cancel a running Snowflake query.
+
+        Args:
+            resource_meta: The SnowflakeJobMetadata containing the query ID.
+            snowflake_private_key: The private key content set as a Flyte secret.
+            snowflake_private_key_passphrase: The passphrase for the private key set as a Flyte secret, if any.
+        """
+        conn = await _get_snowflake_connection(
+            account=resource_meta.account,
+            user=resource_meta.user,
+            database=resource_meta.database,
+            schema=resource_meta.schema,
+            warehouse=resource_meta.warehouse,
+            private_key_content=snowflake_private_key,
+            private_key_passphrase=snowflake_private_key_passphrase,
+            **(resource_meta.connection_kwargs or {}),
+        )
+
+        def _cancel_query():
+            cursor = conn.cursor()
+            try:
+                cursor.execute(f"SELECT SYSTEM$CANCEL_QUERY('{resource_meta.query_id}')")
+            finally:
+                cursor.close()
+                conn.close()
+
+        loop = asyncio.get_running_loop()
+        await loop.run_in_executor(None, _cancel_query)
+
+        logger.info(f"Snowflake query {resource_meta.query_id} cancelled")
+
+
+ConnectorRegistry.register(SnowflakeConnector())

flyteplugins/snowflake/dataframe.py

@@ -0,0 +1,124 @@
+import os
+import re
+import typing
+from typing import Optional
+
+from flyte._utils import lazy_module
+from flyte.io._dataframe.dataframe import DataFrame, DataFrameDecoder, DataFrameEncoder
+from flyteidl2.core import literals_pb2, types_pb2
+
+if typing.TYPE_CHECKING:
+    import pandas as pd
+
+    import snowflake.connector
+else:
+    pd = lazy_module("pandas")
+
+SNOWFLAKE = "snowflake"
+PROTOCOL_SEP = "\\/|://|:"
+
+
+def _get_private_key(private_key_content: str, private_key_passphrase: Optional[str] = None) -> bytes:
+    """Decode a PEM private key and return it in DER format."""
+    from cryptography.hazmat.backends import default_backend
+    from cryptography.hazmat.primitives import serialization
+
+    private_key_bytes = private_key_content.strip().encode()
+    password = private_key_passphrase.encode() if private_key_passphrase else None
+
+    private_key = serialization.load_pem_private_key(
+        private_key_bytes,
+        password=password,
+        backend=default_backend(),
+    )
+
+    return private_key.private_bytes(
+        encoding=serialization.Encoding.DER,
+        format=serialization.PrivateFormat.PKCS8,
+        encryption_algorithm=serialization.NoEncryption(),
+    )
+
+
+def _get_connection(
+    user: str,
+    account: str,
+    database: str,
+    schema: str,
+    warehouse: str,
+) -> "snowflake.connector.SnowflakeConnection":
+    """Create a Snowflake connection using environment-provided credentials."""
+    import snowflake.connector
+
+    conn_params: dict[str, typing.Any] = {
+        "user": user,
+        "account": account,
+        "database": database,
+        "schema": schema,
+        "warehouse": warehouse,
+    }
+
+    # The secrets will be injected as environment variables.
+    private_key_content = os.environ.get("SNOWFLAKE_PRIVATE_KEY")
+    if private_key_content:
+        private_key_passphrase = os.environ.get("SNOWFLAKE_PRIVATE_KEY_PASSPHRASE")
+        conn_params["private_key"] = _get_private_key(private_key_content, private_key_passphrase)
+
+    return snowflake.connector.connect(**conn_params)
+
+
+def _write_to_sf(dataframe: DataFrame):
+    if not dataframe.uri:
+        raise ValueError("dataframe.uri cannot be None.")
+
+    from snowflake.connector.pandas_tools import write_pandas
+
+    uri = typing.cast(str, dataframe.uri)
+    _, user, account, warehouse, database, schema, table = re.split(PROTOCOL_SEP, uri)
+    df = typing.cast("pd.DataFrame", dataframe.val)
+
+    conn = _get_connection(user, account, database, schema, warehouse)
+    write_pandas(conn, df, table)
+
+
+def _read_from_sf(
+    flyte_value: literals_pb2.StructuredDataset,
+    current_task_metadata: literals_pb2.StructuredDatasetMetadata,
+) -> "pd.DataFrame":
+    uri = flyte_value.uri
+    if not uri:
+        raise ValueError("flyte_value.uri cannot be empty.")
+
+    _, user, account, warehouse, database, schema, query_id = re.split(PROTOCOL_SEP, uri)
+
+    conn = _get_connection(user, account, database, schema, warehouse)
+    cs = conn.cursor()
+    cs.get_results_from_sfqid(query_id)
+    return cs.fetch_pandas_all()
+
+
+class PandasToSnowflakeEncodingHandlers(DataFrameEncoder):
+    def __init__(self):
+        super().__init__(pd.DataFrame, SNOWFLAKE, "")
+
+    async def encode(
+        self,
+        dataframe: DataFrame,
+        structured_dataset_type: types_pb2.StructuredDatasetType,
+    ) -> literals_pb2.StructuredDataset:
+        _write_to_sf(dataframe)
+        return literals_pb2.StructuredDataset(
+            uri=typing.cast(str, dataframe.uri),
+            metadata=literals_pb2.StructuredDatasetMetadata(structured_dataset_type=structured_dataset_type),
+        )
+
+
+class SnowflakeToPandasDecodingHandler(DataFrameDecoder):
+    def __init__(self):
+        super().__init__(pd.DataFrame, SNOWFLAKE, "")
+
+    async def decode(
+        self,
+        flyte_value: literals_pb2.StructuredDataset,
+        current_task_metadata: literals_pb2.StructuredDatasetMetadata,
+    ) -> "pd.DataFrame":
+        return _read_from_sf(flyte_value, current_task_metadata)

flyteplugins/snowflake/task.py

@@ -0,0 +1,132 @@
+import re
+from dataclasses import dataclass
+from typing import Any, Dict, Optional, Type
+
+from flyte.connectors import AsyncConnectorExecutorMixin
+from flyte.extend import TaskTemplate
+from flyte.models import NativeInterface, SerializationContext
+from flyteidl2.core import tasks_pb2
+
+
+@dataclass
+class SnowflakeConfig(object):
+    """
+    Configure a Snowflake Task using a `SnowflakeConfig` object.
+
+    Additional connection parameters (role, authenticator, session_parameters, etc.) can be passed
+    via connection_kwargs.
+    See: https://docs.snowflake.com/en/developer-guide/python-connector/python-connector-api
+
+    Args:
+        account: The Snowflake account identifier.
+        database: The Snowflake database name.
+        schema: The Snowflake schema name.
+        warehouse: The Snowflake warehouse name.
+        user: The Snowflake user name.
+        connection_kwargs: Optional dictionary of additional Snowflake connection parameters.
+    """
+
+    account: str
+    database: str
+    schema: str
+    warehouse: str
+    user: str
+    connection_kwargs: Optional[Dict[str, Any]] = None
+
+
+class Snowflake(AsyncConnectorExecutorMixin, TaskTemplate):
+    _TASK_TYPE = "snowflake"
+
+    def __init__(
+        self,
+        name: str,
+        query_template: str,
+        plugin_config: SnowflakeConfig,
+        inputs: Optional[Dict[str, Type]] = None,
+        output_dataframe_type: Optional[Type] = None,
+        secret_group: Optional[str] = None,
+        snowflake_private_key: Optional[str] = None,
+        snowflake_private_key_passphrase: Optional[str] = None,
+        batch: bool = False,
+        **kwargs,
+    ):
+        """
+        Task to run parameterized SQL queries against Snowflake.
+
+        Args:
+            name: The name of this task.
+            query_template: The actual query to run. This can be parameterized using Python's
+                printf-style string formatting with named parameters (e.g. %(param_name)s).
+            plugin_config: `SnowflakeConfig` object containing connection metadata.
+            inputs: Name and type of inputs specified as a dictionary.
+            output_dataframe_type: If some data is produced by this query, then you can specify the
+                output dataframe type.
+            secret_group: Optional group for secrets in the secret store. The environment variable
+                name is auto-generated from ``{secret_group}_{key}``, uppercased with hyphens
+                replaced by underscores. If omitted, the key alone is used.
+            snowflake_private_key: The secret key for the Snowflake private key (key-pair auth).
+            snowflake_private_key_passphrase: The secret key for the private key passphrase
+                (if encrypted).
+            batch: When True, list inputs are expanded into a multi-row VALUES clause. The
+                query_template should contain a single ``VALUES (%(col)s, ...)`` placeholder
+                and each input should be a list of equal length.
+
+        Note: For password authentication or other auth methods, pass them via `connection_kwargs`.
+        """
+        outputs = None
+        if output_dataframe_type is not None:
+            outputs = {"results": output_dataframe_type}
+
+        super().__init__(
+            name=name,
+            interface=NativeInterface(
+                {k: (v, None) for k, v in inputs.items()} if inputs else {},
+                outputs or {},
+            ),
+            task_type=self._TASK_TYPE,
+            image=None,
+            **kwargs,
+        )
+
+        self.output_dataframe_type = output_dataframe_type
+        self.plugin_config = plugin_config
+        self.query_template = re.sub(r"\s+", " ", query_template.replace("\n", " ").replace("\t", " ")).strip()
+        self.batch = batch
+        self.secret_group = secret_group
+        self.snowflake_private_key = snowflake_private_key
+        self.snowflake_private_key_passphrase = snowflake_private_key_passphrase
+
+    def _to_env_var(self, key: str) -> str:
+        """Generate an environment variable name from the secret group and key."""
+        env_var = f"{self.secret_group}_{key}" if self.secret_group else key
+        return env_var.replace("-", "_").upper()
+
+    def custom_config(self, sctx: SerializationContext) -> Optional[Dict[str, Any]]:
+        config = {
+            "account": self.plugin_config.account,
+            "database": self.plugin_config.database,
+            "schema": self.plugin_config.schema,
+            "warehouse": self.plugin_config.warehouse,
+            "user": self.plugin_config.user,
+        }
+
+        if self.batch:
+            config["batch"] = True
+
+        # Add additional connection parameters
+        if self.plugin_config.connection_kwargs:
+            config["connection_kwargs"] = self.plugin_config.connection_kwargs
+
+        secrets = {}
+        if self.snowflake_private_key is not None:
+            secrets["snowflake_private_key"] = self._to_env_var(self.snowflake_private_key)
+        if self.snowflake_private_key_passphrase is not None:
+            secrets["snowflake_private_key_passphrase"] = self._to_env_var(self.snowflake_private_key_passphrase)
+        if secrets:
+            config["secrets"] = secrets
+
+        return config
+
+    def sql(self, sctx: SerializationContext) -> Optional[str]:
+        sql = tasks_pb2.Sql(statement=self.query_template, dialect=tasks_pb2.Sql.Dialect.ANSI)
+        return sql

flyteplugins_snowflake-2.0.0.dist-info/METADATA

@@ -0,0 +1,178 @@
+Metadata-Version: 2.4
+Name: flyteplugins-snowflake
+Version: 2.0.0
+Summary: Snowflake plugin for flyte
+Author-email: Kevin Su <pingsutw@users.noreply.github.com>
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+Requires-Dist: flyte[connector]
+Requires-Dist: snowflake-connector-python[pandas]
+Requires-Dist: cryptography
+
+# Snowflake Plugin for Flyte
+
+Run Snowflake SQL queries as Flyte tasks with parameterized inputs, key-pair authentication, batch inserts, and DataFrame support.
+
+## Installation
+
+```bash
+pip install flyteplugins-snowflake
+```
+
+## Quick start
+
+```python
+from flyteplugins.snowflake import Snowflake, SnowflakeConfig
+
+import flyte
+
+config = SnowflakeConfig(
+    account="myorg-myaccount",
+    user="flyte_user",
+    database="ANALYTICS",
+    schema="PUBLIC",
+    warehouse="COMPUTE_WH",
+)
+
+query = Snowflake(
+    name="count_users",
+    query_template="SELECT COUNT(*) FROM users",
+    plugin_config=config,
+    snowflake_private_key="snowflake-pk",
+)
+```
+
+## Authentication
+
+The plugin supports Snowflake [key-pair authentication](https://docs.snowflake.com/en/user-guide/key-pair-auth). Pass secret keys via `snowflake_private_key` (and optionally `snowflake_private_key_passphrase`).
+
+```python
+task = Snowflake(
+    name="my_task",
+    query_template="SELECT 1",
+    plugin_config=config,
+    snowflake_private_key="private-key",
+    snowflake_private_key_passphrase="passphrase",
+    # Generates env vars: PRIVATE_KEY, PASSPHRASE
+)
+```
+
+For other auth methods (password, OAuth, etc.), pass them via `connection_kwargs`:
+
+```python
+config = SnowflakeConfig(
+    account="myorg-myaccount",
+    user="flyte_user",
+    database="ANALYTICS",
+    schema="PUBLIC",
+    warehouse="COMPUTE_WH",
+    connection_kwargs={"password": "...", "role": "ADMIN"},
+)
+```
+
+## Parameterized queries
+
+Use `%(name)s` placeholders and typed `inputs`:
+
+```python
+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",
+)
+```
+
+## Batch inserts
+
+Set `batch=True` to expand list inputs into multi-row `VALUES` clauses:
+
+```python
+insert_rows = Snowflake(
+    name="insert_users",
+    query_template="INSERT INTO users (id, name, age) VALUES (%(id)s, %(name)s, %(age)s)",
+    plugin_config=config,
+    inputs={"id": list[int], "name": list[str], "age": list[int]},
+    snowflake_private_key="snowflake-pk",
+    batch=True,
+)
+
+# Calling with id=[1,2], name=["Alice","Bob"], age=[30,25] expands to:
+# INSERT INTO users (id, name, age) VALUES (%(id_0)s, %(name_0)s, %(age_0)s), (%(id_1)s, %(name_1)s, %(age_1)s)
+```
+
+## Reading results as DataFrames
+
+Set `output_dataframe_type` to get query results as a pandas DataFrame:
+
+```python
+import pandas as pd
+
+select_task = Snowflake(
+    name="get_users",
+    query_template="SELECT * FROM users",
+    plugin_config=config,
+    output_dataframe_type=pd.DataFrame,
+    snowflake_private_key="snowflake-pk",
+)
+```
+
+## Full example
+
+```python
+import pandas as pd
+from flyteplugins.snowflake import Snowflake, SnowflakeConfig
+
+import flyte
+
+config = SnowflakeConfig(
+    user="KEVIN",
+    account="PWGJLTH-XKB21544",
+    database="FLYTE",
+    schema="PUBLIC",
+    warehouse="COMPUTE_WH",
+)
+
+insert_task = Snowflake(
+    name="insert_rows",
+    inputs={"id": list[int], "name": list[str], "age": list[int]},
+    plugin_config=config,
+    query_template="INSERT INTO FLYTE.PUBLIC.TEST (ID, NAME, AGE) VALUES (%(id)s, %(name)s, %(age)s)",
+    snowflake_private_key="snowflake",
+    batch=True,
+)
+
+select_task = Snowflake(
+    name="select_all",
+    output_dataframe_type=pd.DataFrame,
+    plugin_config=config,
+    query_template="SELECT * FROM FLYTE.PUBLIC.TEST",
+    snowflake_private_key="snowflake",
+)
+
+snowflake_env = flyte.TaskEnvironment.from_task("snowflake_env", insert_task, select_task)
+
+env = flyte.TaskEnvironment(
+    name="example_env",
+    image=flyte.Image.from_debian_base().with_pip_packages("flyteplugins-snowflake"),
+    secrets=[flyte.Secret(key="snowflake", as_env_var="SNOWFLAKE_PRIVATE_KEY")],
+    depends_on=[snowflake_env],
+)
+
+
+@env.task
+async def main(ids: list[int], names: list[str], ages: list[int]) -> float:
+    await insert_task(id=ids, name=names, age=ages)
+    df = await select_task()
+    return df["AGE"].mean().item()
+
+
+if __name__ == "__main__":
+    flyte.init_from_config()
+    run = flyte.with_runcontext(mode="remote").run(
+        main, ids=[123, 456], names=["Kevin", "Alice"], ages=[30, 25],
+    )
+    print(run.url)
+```

flyteplugins_snowflake-2.0.0.dist-info/RECORD

@@ -0,0 +1,9 @@
+flyteplugins/snowflake/__init__.py,sha256=pduVR0IVn0gejuQoP3KjkUfO4OGypHocm4lFh_vnZXQ,1673
+flyteplugins/snowflake/connector.py,sha256=ahFjYUVuYYvgW41pQ86spHQFP7HIAMa_PdkJvu4UkeI,14292
+flyteplugins/snowflake/dataframe.py,sha256=yUkRVCGWZ0IqbK9CIaKsp-MkwshJdbJPd49jWlzM7BM,3995
+flyteplugins/snowflake/task.py,sha256=jDdYMjD5F7JZGOu2Mp5ax1kqV3R7wz6BIVbw7-KAaGE,5416
+flyteplugins_snowflake-2.0.0.dist-info/METADATA,sha256=g77dxtEKw9fcSlWlkolw9Imacs9G41eWjBKoOsVFvZ0,4603
+flyteplugins_snowflake-2.0.0.dist-info/WHEEL,sha256=YCfwYGOYMi5Jhw2fU4yNgwErybb2IX5PEwBKV4ZbdBo,91
+flyteplugins_snowflake-2.0.0.dist-info/entry_points.txt,sha256=JwU7E683aUhA8NdHlUkF3RioiWf7aHE-c3mMgkUP4qM,83
+flyteplugins_snowflake-2.0.0.dist-info/top_level.txt,sha256=cgd779rPu9EsvdtuYgUxNHHgElaQvPn74KhB5XSeMBE,13
+flyteplugins_snowflake-2.0.0.dist-info/RECORD,,

flyteplugins_snowflake-2.0.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (82.0.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

flyteplugins_snowflake-2.0.0.dist-info/entry_points.txt

@@ -0,0 +1,2 @@
+[flyte.connectors]
+snowflake = flyteplugins.snowflake.connector:SnowflakeConnector

flyteplugins_snowflake-2.0.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+flyteplugins