Fix Snowflake Agent Bug

flytekit/core/context_manager.py

@@ -367,6 +367,12 @@
         Retrieves a secret using the resolution order -> Env followed by file. If not found raises a ValueError
         param encode_mode, defines the mode to open files, it can either be "r" to read file, or "rb" to read binary file
         """
+
+        from flytekit.configuration.plugin import get_plugin
+
+        if not get_plugin().secret_requires_group():
+            group, group_version = None, None
+
         env_var = self.get_secrets_env_var(group, key, group_version)
         fpath = self.get_secrets_file(group, key, group_version)
         v = os.environ.get(env_var)

flytekit/core/type_engine.py

@@ -983,6 +983,7 @@
             register_arrow_handlers,
             register_bigquery_handlers,
             register_pandas_handlers,
+            register_snowflake_handlers,
         )
         from flytekit.types.structured.structured_dataset import DuplicateHandlerError
 
@@ -1015,6 +1016,11 @@
             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 @@
             "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,106 @@
+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"
+PROTOCOL_SEP = "\\/|://|:"
+
+
+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("private_key", "snowflake", encode_mode="r")
+
+    # Cryptography needs the string 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(PROTOCOL_SEP, 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(PROTOCOL_SEP, uri)
+
+    conn = snowflake.connector.connect(
+        user=user,
+        account=account,
+        private_key=get_private_key(),
+        database=database,
+        schema=schema,
+        warehouse=warehouse,
+    )
+
+    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

package.json

@@ -0,0 +1,5 @@
+{
+  "dependencies": {
+    "msgpack5": "^6.0.2"
+  }
+}

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,