flyteplugins-snowflake 2.0.0b54__py3-none-any.whl

flyteplugins/snowflake/__init__.py

@@ -0,0 +1,4 @@
+from flyteplugins.snowflake.connector import SnowflakeConnector
+from flyteplugins.snowflake.task import Snowflake, SnowflakeConfig
+
+__all__ = ["Snowflake", "SnowflakeConfig", "SnowflakeConnector"]

flyteplugins/snowflake/connector.py

@@ -0,0 +1,271 @@
+import asyncio
+from dataclasses import dataclass
+from typing import Any, Dict, Optional
+
+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 import connector
+
+TASK_TYPE = "snowflake"
+
+
+@dataclass
+class SnowflakeJobMetadata(ResourceMeta):
+    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.
+    """
+    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,
+) -> connector.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
+    """
+
+    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 connector.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."""
+    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"
+
+
+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.
+        """
+        custom = json_format.MessageToDict(task_template.custom) if task_template.custom else {}
+
+        account = custom.get("account")
+        if not account:
+            raise ValueError("Missing Snowflake account. Please 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
+
+        def _execute_query():
+            cursor = conn.cursor()
+
+            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.
+        """
+        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.account}/{resource_meta.database}/"
+                f"{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.
+        """
+        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/task.py

@@ -0,0 +1,103 @@
+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.io import DataFrame
+from flyte.models import NativeInterface, SerializationContext
+from flyteidl2.core import tasks_pb2
+
+
+@dataclass
+class SnowflakeConfig(object):
+    """
+    SnowflakeConfig should be used to configure a Snowflake Task.
+
+    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
+    """
+
+    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[DataFrame]] = None,
+        snowflake_private_key: Optional[str] = None,
+        snowflake_private_key_passphrase: Optional[str] = None,
+        **kwargs,
+    ):
+        """
+        To be used to query Snowflake databases.
+
+        :param name: The name of this task, should be unique in the project
+        :param query_template: The actual query to run. We use Flyte's Golang templating format for query templating.
+        :param plugin_config: SnowflakeConfig object (includes connection_kwargs for additional parameters)
+        :param inputs: Name and type of inputs specified as an ordered dictionary
+        :param output_dataframe_type: If some data is produced by this query, then you can specify the
+         output dataframe type.
+        :param snowflake_private_key: The name of the secret containing the Snowflake private key for key-pair auth.
+        :param snowflake_private_key_passphrase: The name of the secret containing the private key passphrase
+            (if encrypted).
+
+        Note: For password authentication or other auth methods, pass them via plugin_config.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.snowflake_private_key = snowflake_private_key
+        self.snowflake_private_key_passphrase = snowflake_private_key_passphrase
+
+    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,
+        }
+
+        # 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.snowflake_private_key
+        if self.snowflake_private_key_passphrase is not None:
+            secrets["snowflake_private_key_passphrase"] = 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.0b54.dist-info/METADATA

@@ -0,0 +1,42 @@
+Metadata-Version: 2.4
+Name: flyteplugins-snowflake
+Version: 2.0.0b54
+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
+Requires-Dist: cryptography
+
+# Snowflake Plugin for Flyte
+
+This plugin provides Snowflake integration for Flyte, enabling you to run Snowflake queries as Flyte tasks.
+
+## Installation
+
+```bash
+pip install flyteplugins-snowflake
+```
+
+## Usage
+
+```python
+from flyteplugins.snowflake import Snowflake, SnowflakeConfig
+
+config = SnowflakeConfig(
+    account="myaccount.us-east-1",
+    user="myuser",
+    database="mydb",
+    schema="PUBLIC",
+    warehouse="mywarehouse",
+)
+
+task = Snowflake(
+    name="my_query",
+    query_template="INSERT INTO FLYTE.PUBLIC.TEST (ID, NAME, AGE) VALUES (%(id)s, %(name)s, %(age)s);",
+    plugin_config=config,
+    inputs={"id": int, "name": str, "age": int},
+    snowflake_private_key="snowflake-private-key-secret",
+)
+```

flyteplugins_snowflake-2.0.0b54.dist-info/RECORD

@@ -0,0 +1,9 @@
+flyteplugins/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+flyteplugins/snowflake/__init__.py,sha256=o44Xb0MqEYmjLDx-lWR5JHyAX16ReCHTW1andZiGJ5U,197
+flyteplugins/snowflake/connector.py,sha256=-PTgNflQcTLn4qVP4b61aoyspTjIah76lRFpGGUHamg,9308
+flyteplugins/snowflake/task.py,sha256=KbOA9wdQ60h1vrlhb30SBy9b8QJjwtQ1cDCJ_QQ0Aho,4129
+flyteplugins_snowflake-2.0.0b54.dist-info/METADATA,sha256=_juIGsQVQKe7me65vRc2oZluL8OEBKX4lSoJTaQmZhM,1054
+flyteplugins_snowflake-2.0.0b54.dist-info/WHEEL,sha256=wUyA8OaulRlbfwMtmQsvNngGrxQHAvkKcvRmdizlJi0,92
+flyteplugins_snowflake-2.0.0b54.dist-info/entry_points.txt,sha256=JwU7E683aUhA8NdHlUkF3RioiWf7aHE-c3mMgkUP4qM,83
+flyteplugins_snowflake-2.0.0b54.dist-info/top_level.txt,sha256=cgd779rPu9EsvdtuYgUxNHHgElaQvPn74KhB5XSeMBE,13
+flyteplugins_snowflake-2.0.0b54.dist-info/RECORD,,

flyteplugins_snowflake-2.0.0b54.dist-info/WHEEL

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

flyteplugins_snowflake-2.0.0b54.dist-info/entry_points.txt

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

flyteplugins_snowflake-2.0.0b54.dist-info/top_level.txt

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