Fix Snowflake Agent Bug

flytekit/core/type_engine.py

@@ -983,6 +983,7 @@ def lazy_import_transformers(cls):
             register_arrow_handlers,
             register_bigquery_handlers,
             register_pandas_handlers,
+            register_snowflake_handlers,
         )
         from flytekit.types.structured.structured_dataset import DuplicateHandlerError
 
@@ -1015,6 +1016,11 @@ def lazy_import_transformers(cls):
             from flytekit.types import numpy  # noqa: F401
         if is_imported("PIL"):
             from flytekit.types.file import image  # noqa: F401
+        if is_imported("snowflake.connector"):
+            try:
+                register_snowflake_handlers()
+            except DuplicateHandlerError:
+                logger.debug("Transformer for snowflake is already registered.")
 
     @classmethod
     def to_literal_type(cls, python_type: Type) -> LiteralType:

flytekit/types/structured/__init__.py

@@ -68,3 +68,17 @@ def register_bigquery_handlers():
             "We won't register bigquery handler for structured dataset because "
             "we can't find the packages google-cloud-bigquery-storage and google-cloud-bigquery"
         )
+
+
+def register_snowflake_handlers():
+    try:
+        from .snowflake import PandasToSnowflakeEncodingHandlers, SnowflakeToPandasDecodingHandler
+
+        StructuredDatasetTransformerEngine.register(SnowflakeToPandasDecodingHandler())
+        StructuredDatasetTransformerEngine.register(PandasToSnowflakeEncodingHandlers())
+
+    except ImportError:
+        logger.info(
+            "We won't register snowflake handler for structured dataset because "
+            "we can't find package snowflake-connector-python"
+        )

flytekit/types/structured/snowflake.py

@@ -0,0 +1,105 @@
+import re
+import typing
+
+import pandas as pd
+import snowflake.connector
+from snowflake.connector.pandas_tools import write_pandas
+
+import flytekit
+from flytekit import FlyteContext
+from flytekit.models import literals
+from flytekit.models.types import StructuredDatasetType
+from flytekit.types.structured.structured_dataset import (
+    StructuredDataset,
+    StructuredDatasetDecoder,
+    StructuredDatasetEncoder,
+    StructuredDatasetMetadata,
+)
+
+SNOWFLAKE = "snowflake"
+
+
+def get_private_key() -> bytes:
+    from cryptography.hazmat.backends import default_backend
+    from cryptography.hazmat.primitives import serialization
+
+    pk_string = flytekit.current_context().secrets.get(None, "snowflake", encode_mode="r")
+    # cryptography needs str to be stripped and converted to bytes
+    pk_string = pk_string.strip().encode()
+    p_key = serialization.load_pem_private_key(pk_string, password=None, backend=default_backend())
+
+    pkb = p_key.private_bytes(
+        encoding=serialization.Encoding.DER,
+        format=serialization.PrivateFormat.PKCS8,
+        encryption_algorithm=serialization.NoEncryption(),
+    )
+
+    return pkb
+
+
+def _write_to_sf(structured_dataset: StructuredDataset):
+    if structured_dataset.uri is None:
+        raise ValueError("structured_dataset.uri cannot be None.")
+
+    uri = structured_dataset.uri
+    _, user, account, warehouse, database, schema, table = re.split("\\/|://|:", uri)
+    df = structured_dataset.dataframe
+
+    conn = snowflake.connector.connect(
+        user=user, account=account, private_key=get_private_key(), database=database, schema=schema, warehouse=warehouse
+    )
+
+    write_pandas(conn, df, table)
+
+
+def _read_from_sf(
+    flyte_value: literals.StructuredDataset, current_task_metadata: StructuredDatasetMetadata
+) -> pd.DataFrame:
+    if flyte_value.uri is None:
+        raise ValueError("structured_dataset.uri cannot be None.")
+
+    uri = flyte_value.uri
+    _, user, account, warehouse, database, schema, query_id = re.split("\\/|://|:", uri)
+
+    conn = snowflake.connector.connect(
+        user=user,
+        account=account,
+        private_key=get_private_key(),
+        database=database,
+        schema=schema,
+        warehouse=warehouse,
+        table="FLYTEAGENT.PUBLIC.TEST",
+    )
+
+    cs = conn.cursor()
+    cs.get_results_from_sfqid(query_id)
+    return cs.fetch_pandas_all()
+
+
+class PandasToSnowflakeEncodingHandlers(StructuredDatasetEncoder):
+    def __init__(self):
+        super().__init__(python_type=pd.DataFrame, protocol=SNOWFLAKE, supported_format="")
+
+    def encode(
+        self,
+        ctx: FlyteContext,
+        structured_dataset: StructuredDataset,
+        structured_dataset_type: StructuredDatasetType,
+    ) -> literals.StructuredDataset:
+        _write_to_sf(structured_dataset)
+        return literals.StructuredDataset(
+            uri=typing.cast(str, structured_dataset.uri), metadata=StructuredDatasetMetadata(structured_dataset_type)
+        )
+
+
+class SnowflakeToPandasDecodingHandler(StructuredDatasetDecoder):
+    def __init__(self):
+        super().__init__(pd.DataFrame, protocol=SNOWFLAKE, supported_format="")
+
+    def decode(
+        self,
+        ctx: FlyteContext,
+        flyte_value: literals.StructuredDataset,
+        current_task_metadata: StructuredDatasetMetadata,
+    ) -> pd.DataFrame:
+        return _read_from_sf(flyte_value, current_task_metadata)

flytekit/types/structured/structured_dataset.py

@@ -6,7 +6,7 @@
 import typing
 from abc import ABC, abstractmethod
 from dataclasses import dataclass, field, is_dataclass
-from typing import Dict, Generator, Optional, Type, Union
+from typing import Dict, Generator, List, Optional, Type, Union
 
 from dataclasses_json import config
 from fsspec.utils import get_protocol
@@ -222,7 +222,12 @@ def extract_cols_and_format(
 
 
 class StructuredDatasetEncoder(ABC):
-    def __init__(self, python_type: Type[T], protocol: Optional[str] = None, supported_format: Optional[str] = None):
+    def __init__(
+        self,
+        python_type: Type[T],
+        protocol: Optional[str] = None,
+        supported_format: Optional[str] = None,
+    ):
         """
         Extend this abstract class, implement the encode function, and register your concrete class with the
         StructuredDatasetTransformerEngine class in order for the core flytekit type engine to handle
@@ -284,7 +289,13 @@ def encode(
 
 
 class StructuredDatasetDecoder(ABC):
-    def __init__(self, python_type: Type[DF], protocol: Optional[str] = None, supported_format: Optional[str] = None):
+    def __init__(
+        self,
+        python_type: Type[DF],
+        protocol: Optional[str] = None,
+        supported_format: Optional[str] = None,
+        additional_protocols: Optional[List[str]] = None,
+    ):
         """
         Extend this abstract class, implement the decode function, and register your concrete class with the
         StructuredDatasetTransformerEngine class in order for the core flytekit type engine to handle

plugins/flytekit-bigquery/flytekitplugins/bigquery/task.py

@@ -38,7 +38,7 @@ def __init__(
         self,
         name: str,
         query_template: str,
-        task_config: Optional[BigQueryConfig],
+        task_config: BigQueryConfig,
         inputs: Optional[Dict[str, Type]] = None,
         output_structured_dataset_type: Optional[Type[StructuredDataset]] = None,
         **kwargs,

plugins/flytekit-snowflake/flytekitplugins/snowflake/agent.py

@@ -1,18 +1,18 @@
 from dataclasses import dataclass
 from typing import Optional
 
-from flyteidl.core.execution_pb2 import TaskExecution
+from flyteidl.core.execution_pb2 import TaskExecution, TaskLog
 
-from flytekit import FlyteContextManager, StructuredDataset, lazy_module, logger
+import flytekit
+from flytekit import FlyteContextManager, StructuredDataset, logger
 from flytekit.core.type_engine import TypeEngine
 from flytekit.extend.backend.base_agent import AgentRegistry, AsyncAgentBase, Resource, ResourceMeta
 from flytekit.extend.backend.utils import convert_to_flyte_phase
 from flytekit.models import literals
 from flytekit.models.literals import LiteralMap
 from flytekit.models.task import TaskTemplate
 from flytekit.models.types import LiteralType, StructuredDatasetType
-
-snowflake_connector = lazy_module("snowflake.connector")
+from snowflake import connector as sc
 
 TASK_TYPE = "snowflake"
 SNOWFLAKE_PRIVATE_KEY = "snowflake_private_key"
@@ -27,15 +27,16 @@ class SnowflakeJobMetadata(ResourceMeta):
     warehouse: str
     table: str
     query_id: str
+    has_output: bool
 
 
 def get_private_key():
     from cryptography.hazmat.backends import default_backend
     from cryptography.hazmat.primitives import serialization
 
-    import flytekit
-
-    pk_string = flytekit.current_context().secrets.get(SNOWFLAKE_PRIVATE_KEY, encode_mode="rb")
+    pk_string = flytekit.current_context().secrets.get(SNOWFLAKE_PRIVATE_KEY, encode_mode="r")
+    # cryptography needs str to be stripped and converted to bytes
+    pk_string = pk_string.strip().encode()
     p_key = serialization.load_pem_private_key(pk_string, password=None, backend=default_backend())
 
     pkb = p_key.private_bytes(
@@ -47,8 +48,8 @@ def get_private_key():
     return pkb
 
 
-def get_connection(metadata: SnowflakeJobMetadata) -> snowflake_connector:
-    return snowflake_connector.connect(
+def get_connection(metadata: SnowflakeJobMetadata) -> sc:
+    return sc.connect(
         user=metadata.user,
         account=metadata.account,
         private_key=get_private_key(),
@@ -69,10 +70,11 @@ async def create(
     ) -> SnowflakeJobMetadata:
         ctx = FlyteContextManager.current_context()
         literal_types = task_template.interface.inputs
-        params = TypeEngine.literal_map_to_kwargs(ctx, inputs, literal_types=literal_types) if inputs else None
+
+        params = TypeEngine.literal_map_to_kwargs(ctx, inputs, literal_types=literal_types) if inputs.literals else None
 
         config = task_template.config
-        conn = snowflake_connector.connect(
+        conn = sc.connect(
             user=config["user"],
             account=config["account"],
             private_key=get_private_key(),
@@ -82,7 +84,7 @@ async def create(
         )
 
         cs = conn.cursor()
-        cs.execute_async(task_template.sql.statement, params=params)
+        cs.execute_async(task_template.sql.statement, params)
 
         return SnowflakeJobMetadata(
             user=config["user"],
@@ -91,34 +93,42 @@ async def create(
             schema=config["schema"],
             warehouse=config["warehouse"],
             table=config["table"],
-            query_id=str(cs.sfqid),
+            query_id=cs.sfqid,
+            has_output=task_template.interface.outputs is not None and len(task_template.interface.outputs) > 0,
         )
 
     async def get(self, resource_meta: SnowflakeJobMetadata, **kwargs) -> Resource:
         conn = get_connection(resource_meta)
         try:
             query_status = conn.get_query_status_throw_if_error(resource_meta.query_id)
-        except snowflake_connector.ProgrammingError as err:
+        except sc.ProgrammingError as err:
             logger.error("Failed to get snowflake job status with error:", err.msg)
             return Resource(phase=TaskExecution.FAILED)
+
+        log_link = TaskLog(
+            uri=construct_query_link(resource_meta=resource_meta),
+            name="Snowflake Query Details",
+        )
+        # The snowflake job's state is determined by query status.
+        # https://github.com/snowflakedb/snowflake-connector-python/blob/main/src/snowflake/connector/constants.py#L373
         cur_phase = convert_to_flyte_phase(str(query_status.name))
         res = None
 
-        if cur_phase == TaskExecution.SUCCEEDED:
+        if cur_phase == TaskExecution.SUCCEEDED and resource_meta.has_output:
             ctx = FlyteContextManager.current_context()
-            output_metadata = f"snowflake://{resource_meta.user}:{resource_meta.account}/{resource_meta.warehouse}/{resource_meta.database}/{resource_meta.schema}/{resource_meta.table}"
+            uri = f"snowflake://{resource_meta.user}:{resource_meta.account}/{resource_meta.warehouse}/{resource_meta.database}/{resource_meta.schema}/{resource_meta.query_id}"
             res = literals.LiteralMap(
                 {
                     "results": TypeEngine.to_literal(
                         ctx,
-                        StructuredDataset(uri=output_metadata),
+                        StructuredDataset(uri=uri),
                         StructuredDataset,
                         LiteralType(structured_dataset_type=StructuredDatasetType(format="")),
                     )
                 }
-            ).to_flyte_idl()
+            )
 
-        return Resource(phase=cur_phase, outputs=res)
+        return Resource(phase=cur_phase, outputs=res, log_links=[log_link])
 
     async def delete(self, resource_meta: SnowflakeJobMetadata, **kwargs):
         conn = get_connection(resource_meta)
@@ -131,4 +141,17 @@ async def delete(self, resource_meta: SnowflakeJobMetadata, **kwargs):
             conn.close()
 
 
+def construct_query_link(resource_meta: SnowflakeJobMetadata) -> str:
+    base_url = "https://app.snowflake.com"
+
+    # Extract the account and region (assuming the format is account-region, you might need to adjust this based on your actual account format)
+    account_parts = resource_meta.account.split("-")
+    account = account_parts[0]
+    region = account_parts[1] if len(account_parts) > 1 else ""
+
+    url = f"{base_url}/{region}/{account}/#/compute/history/queries/{resource_meta.query_id}/detail"
+
+    return url
+
+
 AgentRegistry.register(SnowflakeAgent())

plugins/flytekit-snowflake/flytekitplugins/snowflake/task.py

@@ -19,20 +19,22 @@
 class SnowflakeConfig(object):
     """
     SnowflakeConfig should be used to configure a Snowflake Task.
+    You can use the query below to retrieve all metadata for this config.
+
+    SELECT
+        CURRENT_USER() AS "User",
+        CONCAT(CURRENT_ORGANIZATION_NAME(), '-', CURRENT_ACCOUNT_NAME()) AS "Account",
+        CURRENT_DATABASE() AS "Database",
+        CURRENT_SCHEMA() AS "Schema",
+        CURRENT_WAREHOUSE() AS "Warehouse";
     """
 
-    # The user to query against
-    user: Optional[str] = None
-    # The account to query against
-    account: Optional[str] = None
-    # The database to query against
-    database: Optional[str] = None
-    # The optional schema to separate query execution.
-    schema: Optional[str] = None
-    # The optional warehouse to set for the given Snowflake query
-    warehouse: Optional[str] = None
-    # The optional table to set for the given Snowflake query
-    table: Optional[str] = None
+    user: str
+    account: str
+    database: str
+    schema: str
+    warehouse: str
+    table: str
 
 
 class SnowflakeTask(AsyncAgentExecutorMixin, SQLTask[SnowflakeConfig]):
@@ -47,7 +49,7 @@ def __init__(
         self,
         name: str,
         query_template: str,
-        task_config: Optional[SnowflakeConfig] = None,
+        task_config: SnowflakeConfig,
         inputs: Optional[Dict[str, Type]] = None,
         output_schema_type: Optional[Type[StructuredDataset]] = None,
         **kwargs,
@@ -63,13 +65,13 @@ def __init__(
         :param output_schema_type: If some data is produced by this query, then you can specify the output schema type
         :param kwargs: All other args required by Parent type - SQLTask
         """
+
         outputs = None
         if output_schema_type is not None:
             outputs = {
                 "results": output_schema_type,
             }
-        if task_config is None:
-            task_config = SnowflakeConfig()
+
         super().__init__(
             name=name,
             task_config=task_config,