Fix Snowflake Agent Bug

[Future-Outlier]

Tracking issue

flyteorg/flyte#3936

Why are the changes needed?

What changes were proposed in this pull request?

We want to support workflow like this.

from flytekit import kwtypes, workflow, task, ImageSpec, StructuredDataset
from flytekitplugins.snowflake import SnowflakeConfig, SnowflakeTask
from flytekit import Secret
import pandas as pd

flytekit_hash = "0c6fe09884849752cd77a9988a4b2bdc63980b7e"
flytekit = f"git+https://github.com/flyteorg/flytekit.git@{flytekit_hash}"
snowflake_plugins = f"git+https://github.com/flyteorg/flytekit.git@{flytekit_hash}#subdirectory=plugins/flytekit-snowflake"

image = ImageSpec(
    packages=[flytekit, snowflake_plugins, "pandas"],
    apt_packages=["git"],
    registry="localhost:30000",
)

"""
Define a Snowflake task to insert data into the FLYTEAGENT.PUBLIC.TEST table.
The `inputs` parameter specifies the types of the inputs using `kwtypes`.
The `query_template` uses Python string interpolation to insert these inputs into the SQL query.
The placeholders `%(id)s`, `%(name)s`, and `%(age)s` will be replaced by the actual values
provided when the task is executed.
"""

snowflake_task_insert_query = SnowflakeTask(
    name="insert-dynamic-query",
    inputs=kwtypes(id=int, name=str, age=int),
    task_config=SnowflakeConfig(
        account="WNITKUT-UJ60052",
        user="HANRU",
        database="FLYTEAGENT",
        schema="PUBLIC",
        warehouse="COMPUTE_WH",
        table="FLYTEAGENT.PUBLIC.TEST",
    ),
    query_template="""
            INSERT INTO FLYTEAGENT.PUBLIC.TEST (ID, NAME, AGE)
            VALUES (%(id)s, %(name)s, %(age)s);
            """,
)

snowflake_task_templatized_query = SnowflakeTask(
    name="test-simple-query",
    output_schema_type=StructuredDataset,
    # Define inputs as well as their types that can be used to customize the query.
    # inputs=kwtypes(nation_key=int),
    task_config=SnowflakeConfig(
        account="WNITKUT-UJ60052",
        user="HANRU",
        database="FLYTEAGENT",
        schema="PUBLIC",
        warehouse="COMPUTE_WH",
        table="FLYTEAGENT.PUBLIC.TEST",
    ),
    query_template="SELECT * FROM FLYTEAGENT.PUBLIC.TEST LIMIT 1;",
)

@task(container_image=image, secret_requests=[Secret(
      group="snowflake",
      key="private_key",)])
def print_head(input_sd: StructuredDataset) -> pd.DataFrame:
    df = input_sd.open(pd.DataFrame).all()
    print(df)
    return df


@task(container_image=image, secret_requests=[Secret(
      group="snowflake",
      key="private_key",)])
def write_table() -> StructuredDataset:
    df = pd.DataFrame({
        "ID": [1, 2, 3],
        "NAME": ["union", "union", "union"],
        "AGE": [30, 30, 30]
    })
    print(df)
    return StructuredDataset(dataframe=df, uri="snowflake://HANRU/WNITKUT-UJ60052/COMPUTE_WH/FLYTEAGENT/PUBLIC/TEST")

@workflow
def wf() -> StructuredDataset:
    sd = snowflake_task_templatized_query()
    t1 = print_head(input_sd=sd)
    insert_query = snowflake_task_insert_query(id=1, name="Flyte", age=30)
    sd2 = snowflake_task_templatized_query()

    sd >> t1 >> insert_query >> sd2
    write_table()
    return print_head(input_sd=sd2)


if __name__ == "__main__":
    print(f"Running {__file__} main...")
    wf()
    print("Done!")

How was this patch tested?

Setup process

Screenshots

local execution:

remote execution:

Check all the applicable boxes

• I updated the documentation accordingly.
• All new and existing tests passed.
• All commits are signed-off.

Related PRs

#1811

Docs link

[Future-Outlier]

example:

from flytekit import kwtypes, workflow, StructuredDataset
from flytekitplugins.snowflake import SnowflakeConfig, SnowflakeTask

snowflake_task_templatized_query = SnowflakeTask(
    name="test-simple-query",
    # inputs={},
    output_schema_type=StructuredDataset,

    # Define inputs as well as their types that can be used to customize the query.
    # inputs=kwtypes(nation_key=int),
    task_config=SnowflakeConfig(
        account="WNITKUT-UJ60052",
        user="HANRU",
        database="FLYTEAGENT",
        schema="PUBLIC",
        warehouse="COMPUTE_WH",
    ),
    # query_template="SELECT 1"
    query_template="SELECT * FROM FLYTEAGENT.PUBLIC.TEST;",
)

if __name__ == "__main__":
    sd = snowflake_task_templatized_query()
    print(sd)
    # print(snowflake_task_no_io())
    # print(snowflake_wf(nation_key=10))

[fiedlerNr9]

Thats amazing work!

When reading the StructuredDataset in another task, I stumbled upon snowflake protocol not being supported:

from flytekit import kwtypes, workflow, task, ImageSpec, StructuredDataset
from flytekitplugins.snowflake import SnowflakeConfig, SnowflakeTask
import pandas as pd

snowflake_task_no_io = SnowflakeTask(
    name="something",
    inputs=kwtypes(x=int),
    query_template="SELECT * FROM JANTEST.PUBLIC.PEOPLE LIMIT 2;",
    output_schema_type=StructuredDataset,
    task_config=SnowflakeConfig(
        user="FIEDLERNR9",
        account="FJTBLYP-JK84112",
        database="JANTEST",
        schema="PUBLIC",
        warehouse="COMPUTE_WH",
    ),
)

@task
def print_head(input_sd: StructuredDataset):
    df = input_sd.open(pd.DataFrame).all()
    print(df)


@workflow
def wf():
    sd = snowflake_task_no_io(x=1)
    print_head(input_sd=sd)

full stacktrace:

pyflyte run test_snowflake.py wf
Running Execution on local.
╭─────────────────────────────── Traceback (most recent call last) ────────────────────────────────╮
│ /opt/homebrew/anaconda3/envs/fbin/lib/python3.9/site-packages/flytekit/exceptions/scopes.py:242  │
│ in user_entry_point                                                                              │
│                                                                                                  │
│   239 │   │   │   │   │   raise type(exc)(f"Error encountered while executing '{fn_name}':\n     │
│   240 │   │   else:                                                                              │
│   241 │   │   │   try:                                                                           │
│ ❱ 242 │   │   │   │   return wrapped(*args, **kwargs)                                            │
│   243 │   │   │   except FlyteScopedException as scoped:                                         │
│   244 │   │   │   │   raise scoped                                                               │
│   245 │   │   │   except _user_exceptions.FlyteUserException:                                    │
│                                                                                                  │
│ /Users/janfiedler/Documents/Customers/Fortune-Brands/plugin-test/test_snowflake.py:38 in         │
│ print_head                                                                                       │
│                                                                                                  │
│   35                                                                                             │
│   36 @task                                                                                       │
│   37 def print_head(input_sd: StructuredDataset):                                                │
│ ❱ 38 │   df = input_sd.open(pd.DataFrame).all()                                                  │
│   39 │   print(df)                                                                               │
│   40                                                                                             │
│   41                                                                                             │
│                                                                                                  │
│ /opt/homebrew/anaconda3/envs/fbin/lib/python3.9/site-packages/flytekit/types/structured/structur │
│ ed_dataset.py:141 in all                                                                         │
│                                                                                                  │
│   138 │   │   if self._dataframe_type is None:                                                   │
│   139 │   │   │   raise ValueError("No dataframe type set. Use open() to set the local datafra   │
│   140 │   │   ctx = FlyteContextManager.current_context()                                        │
│ ❱ 141 │   │   return flyte_dataset_transformer.open_as(ctx, self.literal, self._dataframe_type   │
│   142 │                                                                                          │
│   143 │   def iter(self) -> Generator[DF, None, None]:                                           │
│   144 │   │   if self._dataframe_type is None:                                                   │
│                                                                                                  │
│ /opt/homebrew/anaconda3/envs/fbin/lib/python3.9/site-packages/flytekit/types/structured/structur │
│ ed_dataset.py:854 in open_as                                                                     │
│                                                                                                  │
│   851 │   │   """                                                                                │
│   852 │   │   protocol = get_protocol(sd.uri)                                                    │
│   853 │   │   decoder = self.get_decoder(df_type, protocol, sd.metadata.structured_dataset_typ   │
│ ❱ 854 │   │   result = decoder.decode(ctx, sd, updated_metadata)                                 │
│   855 │   │   if isinstance(result, types.GeneratorType):                                        │
│   856 │   │   │   raise ValueError(f"Decoder {decoder} returned iterator {result} but whole va   │
│   857 │   │   return result                                                                      │
│                                                                                                  │
│ /opt/homebrew/anaconda3/envs/fbin/lib/python3.9/site-packages/flytekit/types/structured/basic_df │
│ s.py:137 in decode                                                                               │
│                                                                                                  │
│   134 │   │   if current_task_metadata.structured_dataset_type and current_task_metadata.struc   │
│   135 │   │   │   columns = [c.name for c in current_task_metadata.structured_dataset_type.col   │
│   136 │   │   try:                                                                               │
│ ❱ 137 │   │   │   return pd.read_parquet(uri, columns=columns, storage_options=kwargs)           │
│   138 │   │   except NoCredentialsError:                                                         │
│   139 │   │   │   logger.debug("S3 source detected, attempting anonymous S3 access")             │
│   140 │   │   │   kwargs = get_pandas_storage_options(uri=uri, data_config=ctx.file_access.dat   │
│                                                                                                  │
│ /opt/homebrew/anaconda3/envs/fbin/lib/python3.9/site-packages/pandas/io/parquet.py:667 in        │
│ read_parquet                                                                                     │
│                                                                                                  │
│   664 │   │   use_nullable_dtypes = False                                                        │
│   665 │   check_dtype_backend(dtype_backend)                                                     │
│   666 │                                                                                          │
│ ❱ 667 │   return impl.read(                                                                      │
│   668 │   │   path,                                                                              │
│   669 │   │   columns=columns,                                                                   │
│   670 │   │   filters=filters,                                                                   │
│                                                                                                  │
│ /opt/homebrew/anaconda3/envs/fbin/lib/python3.9/site-packages/pandas/io/parquet.py:267 in read   │
│                                                                                                  │
│   264 │   │   if manager == "array":                                                             │
│   265 │   │   │   to_pandas_kwargs["split_blocks"] = True  # type: ignore[assignment]            │
│   266 │   │                                                                                      │
│ ❱ 267 │   │   path_or_handle, handles, filesystem = _get_path_or_handle(                         │
│   268 │   │   │   path,                                                                          │
│   269 │   │   │   filesystem,                                                                    │
│   270 │   │   │   storage_options=storage_options,                                               │
│                                                                                                  │
│ /opt/homebrew/anaconda3/envs/fbin/lib/python3.9/site-packages/pandas/io/parquet.py:122 in        │
│ _get_path_or_handle                                                                              │
│                                                                                                  │
│   119 │   │   │   │   pass                                                                       │
│   120 │   │   if fs is None:                                                                     │
│   121 │   │   │   fsspec = import_optional_dependency("fsspec")                                  │
│ ❱ 122 │   │   │   fs, path_or_handle = fsspec.core.url_to_fs(                                    │
│   123 │   │   │   │   path_or_handle, **(storage_options or {})                                  │
│   124 │   │   │   )                                                                              │
│   125 │   elif storage_options and (not is_url(path_or_handle) or mode != "rb"):                 │
│                                                                                                  │
│ /opt/homebrew/anaconda3/envs/fbin/lib/python3.9/site-packages/fsspec/core.py:396 in url_to_fs    │
│                                                                                                  │
│   393 │   │   "num",                                                                             │
│   394 │   }                                                                                      │
│   395 │   kwargs = {k: v for k, v in kwargs.items() if k not in known_kwargs}                    │
│ ❱ 396 │   chain = _un_chain(url, kwargs)                                                         │
│   397 │   inkwargs = {}                                                                          │
│   398 │   # Reverse iterate the chain, creating a nested target_* structure                      │
│   399 │   for i, ch in enumerate(reversed(chain)):                                               │
│                                                                                                  │
│ /opt/homebrew/anaconda3/envs/fbin/lib/python3.9/site-packages/fsspec/core.py:344 in _un_chain    │
│                                                                                                  │
│   341 │   kwargs = kwargs.copy()                                                                 │
│   342 │   for bit in reversed(bits):                                                             │
│   343 │   │   protocol = kwargs.pop("protocol", None) or split_protocol(bit)[0] or "file"        │
│ ❱ 344 │   │   cls = get_filesystem_class(protocol)                                               │
│   345 │   │   extra_kwargs = cls._get_kwargs_from_urls(bit)                                      │
│   346 │   │   kws = kwargs.pop(protocol, {})                                                     │
│   347 │   │   if bit is bits[0]:                                                                 │
│                                                                                                  │
│ /opt/homebrew/anaconda3/envs/fbin/lib/python3.9/site-packages/fsspec/registry.py:239 in          │
│ get_filesystem_class                                                                             │
│                                                                                                  │
│   236 │                                                                                          │
│   237 │   if protocol not in registry:                                                           │
│   238 │   │   if protocol not in known_implementations:                                          │
│ ❱ 239 │   │   │   raise ValueError(f"Protocol not known: {protocol}")                            │
│   240 │   │   bit = known_implementations[protocol]                                              │
│   241 │   │   try:                                                                               │
│   242 │   │   │   register_implementation(protocol, _import_class(bit["class"]))                 │
╰──────────────────────────────────────────────────────────────────────────────────────────────────╯
ValueError: Protocol not known: snowflake

During handling of the above exception, another exception occurred:

╭─────────────────────────────── Traceback (most recent call last) ────────────────────────────────╮
│ /opt/homebrew/anaconda3/envs/fbin/lib/python3.9/site-packages/flytekit/exceptions/scopes.py:217  │
│ in user_entry_point                                                                              │
│                                                                                                  │
│   214 │   │   │   # See comment at this location for system_entry_point                          │
│   215 │   │   │   fn_name = wrapped.__name__                                                     │
│   216 │   │   │   try:                                                                           │
│ ❱ 217 │   │   │   │   return wrapped(*args, **kwargs)                                            │
│   218 │   │   │   except FlyteScopedException as exc:                                            │
│   219 │   │   │   │   raise exc.type(f"Error encountered while executing '{fn_name}':\n  {exc.   │
│   220 │   │   │   except Exception as exc:                                                       │
│                                                                                                  │
│ /Users/janfiedler/Documents/Customers/Fortune-Brands/plugin-test/test_snowflake.py:46 in wf      │
│                                                                                                  │
│   43 def wf():                                                                                   │
│   44 │   sd = snowflake_task_no_io(x=1)                                                          │
│   45 │   # print(sd)                                                                             │
│ ❱ 46 │   print_head(input_sd=sd)                                                                 │
│   47                                                                                             │
│                                                                                                  │
│ /opt/homebrew/anaconda3/envs/fbin/lib/python3.9/site-packages/flytekit/core/base_task.py:358 in  │
│ __call__                                                                                         │
│                                                                                                  │
│   355 │   │   return create_task_output(vals, self.python_interface)                             │
│   356 │                                                                                          │
│   357 │   def __call__(self, *args: object, **kwargs: object) -> Union[Tuple[Promise], Promise   │
│ ❱ 358 │   │   return flyte_entity_call_handler(self, *args, **kwargs)  # type: ignore            │
│   359 │                                                                                          │
│   360 │   def compile(self, ctx: FlyteContext, *args, **kwargs):                                 │
│   361 │   │   raise Exception("not implemented")                                                 │
│                                                                                                  │
│ /opt/homebrew/anaconda3/envs/fbin/lib/python3.9/site-packages/flytekit/core/promise.py:1248 in   │
│ flyte_entity_call_handler                                                                        │
│                                                                                                  │
│   1245 │   │   │   │   │   return create_task_output(vals, entity.python_interface)              │
│   1246 │   │   │   │   else:                                                                     │
│   1247 │   │   │   │   │   return None                                                           │
│ ❱ 1248 │   │   │   return cast(LocallyExecutable, entity).local_execute(ctx, **kwargs)           │
│   1249 │   else:                                                                                 │
│   1250 │   │   mode = cast(LocallyExecutable, entity).local_execution_mode()                     │
│   1251 │   │   omt = OutputMetadataTracker()                                                     │
│                                                                                                  │
│ /opt/homebrew/anaconda3/envs/fbin/lib/python3.9/site-packages/flytekit/core/base_task.py:335 in  │
│ local_execute                                                                                    │
│                                                                                                  │
│   332 │   │   │   # This code should mirror the call to `sandbox_execute` in the above cache c   │
│   333 │   │   │   # Code is simpler with duplication and less metaprogramming, but introduces    │
│   334 │   │   │   # if one is changed and not the other.                                         │
│ ❱ 335 │   │   │   outputs_literal_map = self.sandbox_execute(ctx, input_literal_map)             │
│   336 │   │                                                                                      │
│   337 │   │   if inspect.iscoroutine(outputs_literal_map):                                       │
│   338 │   │   │   return outputs_literal_map                                                     │
│                                                                                                  │
│ /opt/homebrew/anaconda3/envs/fbin/lib/python3.9/site-packages/flytekit/core/base_task.py:415 in  │
│ sandbox_execute                                                                                  │
│                                                                                                  │
│   412 │   │   es = cast(ExecutionState, ctx.execution_state)                                     │
│   413 │   │   b = cast(ExecutionParameters, es.user_space_params).with_task_sandbox()            │
│   414 │   │   ctx = ctx.current_context().with_execution_state(es.with_params(user_space_param   │
│ ❱ 415 │   │   return self.dispatch_execute(ctx, input_literal_map)                               │
│   416 │                                                                                          │
│   417 │   @abstractmethod                                                                        │
│   418 │   def dispatch_execute(                                                                  │
│                                                                                                  │
│ /opt/homebrew/anaconda3/envs/fbin/lib/python3.9/site-packages/flytekit/core/base_task.py:733 in  │
│ dispatch_execute                                                                                 │
│                                                                                                  │
│   730 │   │   │   #   a workflow or a subworkflow etc                                            │
│   731 │   │   │   logger.info(f"Invoking {self.name} with inputs: {native_inputs}")              │
│   732 │   │   │   with timeit("Execute user level code"):                                        │
│ ❱ 733 │   │   │   │   native_outputs = self.execute(**native_inputs)                             │
│   734 │   │   │                                                                                  │
│   735 │   │   │   if inspect.iscoroutine(native_outputs):                                        │
│   736 │   │   │   │   # If native outputs is a coroutine, then this is an eager workflow.        │
│                                                                                                  │
│ /opt/homebrew/anaconda3/envs/fbin/lib/python3.9/site-packages/flytekit/core/python_function_task │
│ .py:199 in execute                                                                               │
│                                                                                                  │
│   196 │   │   handle dynamic tasks or you will no longer be able to use the task as a dynamic    │
│   197 │   │   """                                                                                │
│   198 │   │   if self.execution_mode == self.ExecutionBehavior.DEFAULT:                          │
│ ❱ 199 │   │   │   return exception_scopes.user_entry_point(self._task_function)(**kwargs)        │
│   200 │   │   elif self.execution_mode == self.ExecutionBehavior.EAGER:                          │
│   201 │   │   │   # if the task is a coroutine function, inject the context object so that the   │
│   202 │   │   │   # has access to the FlyteContext.                                              │
│                                                                                                  │
│ /opt/homebrew/anaconda3/envs/fbin/lib/python3.9/site-packages/flytekit/exceptions/scopes.py:148  │
│ in f                                                                                             │
│                                                                                                  │
│   145 │   def inner_decorator(inner_f):                                                          │
│   146 │   │   @_wraps(inner_f)                                                                   │
│   147 │   │   def f(*args, **kwargs):                                                            │
│ ❱ 148 │   │   │   return outer_f(inner_f, args, kwargs)                                          │
│   149 │   │                                                                                      │
│   150 │   │   return f                                                                           │
│   151                                                                                            │
│                                                                                                  │
│ /opt/homebrew/anaconda3/envs/fbin/lib/python3.9/site-packages/flytekit/exceptions/scopes.py:253  │
│ in user_entry_point                                                                              │
│                                                                                                  │
│   250 │   │   │   │   # This is why this function exists - arbitrary exceptions that we don't    │
│   251 │   │   │   │   # interpreted as user exceptions.                                          │
│   252 │   │   │   │   # This will also catch FlyteUserException re-raised by the system_entry_   │
│ ❱ 253 │   │   │   │   raise FlyteScopedUserException(*_exc_info())                               │
│   254 │   finally:                                                                               │
│   255 │   │   _CONTEXT_STACK.pop()                                                               │
│   256                                                                                            │
╰──────────────────────────────────────────────────────────────────────────────────────────────────╯
FlyteScopedUserException: Protocol not known: snowflake

The above exception was the direct cause of the following exception:

ValueError: Error encountered while executing 'wf':
  Protocol not known: snowflake

[Future-Outlier]

TODO:

0. CHECK IF table is required in snowflake
1. FIX TESTS
2. list things to do to improve the doc.
   (1) how to get all config you need in 1 query? It's doable in snowflake console.
   (2) add examples and also mentioned that we have to build our own image for snowflake dependency and mount the secret to read it correctly.

[Future-Outlier]

todo: add flyte-comformance

[codecov]

Codecov Report

Attention: Patch coverage is 7.57576% with 61 lines in your changes missing coverage. Please review.

Project coverage is 78.35%. Comparing base (eae31c0) to head (c4b641e).
Report is 13 commits behind head on master.

Files	Patch %	Lines
flytekit/types/structured/snowflake.py	0.00%	48 Missing ⚠️
flytekit/types/structured/__init__.py	14.28%	6 Missing ⚠️
flytekit/core/type_engine.py	0.00%	4 Missing and 1 partial ⚠️
flytekit/core/context_manager.py	33.33%	1 Missing and 1 partial ⚠️

Additional details and impacted files

@@             Coverage Diff             @@
##           master    #2605       +/-   ##
===========================================
+ Coverage   47.17%   78.35%   +31.17%     
===========================================
  Files         230      189       -41     
  Lines       21322    19110     -2212     
  Branches     3711     3719        +8     
===========================================
+ Hits        10059    14973     +4914     
+ Misses      11154     3464     -7690     
- Partials      109      673      +564     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[Future-Outlier]

We should use this

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";

[Future-Outlier]

If we want to use the secret in the union cloud, we have to add union in our Python packages.

image = ImageSpec(
    packages=[flytekit, snowflake_plugins, "union"],
    apt_packages=["git"],
    registry="futureoutlier",
)

secret_requests=[Secret(
      group=None,
      key="snowflake",)]

If we want to use the secret in the flyte sandbox, we have to remove union in our Python packages.

image = ImageSpec(
    packages=[flytekit, snowflake_plugins,],
    apt_packages=["git"],
    registry="futureoutlier",
)
secret_requests=[Secret(
      group="private_key",
      key="snowflake",)]

Does this looks good to all of you?
cc @pingsutw @kumare3 @thomasjpfan @wild-endeavor @eapolinario @samhita-alla

[thomasjpfan]

Is this secret file based?

I'll prefer if the user can specific the secret and then it's passed here. I only see it possible by having the user call register_snowflake_handlers:

secret = Secret(key="abc", group="xyz")

# `get_private_key` uses `secret.{key, group}` to get the secret
register_snowflake_handlers(secret=secret)

@task(secret_request=[secret])

[Future-Outlier]

Is this secret file based?
The secret is a file.

I'll prefer if the user can specific the secret and then it's passed here.
This is a problem, we haven't figured out how to make users specify the secret in a better way.
But maybe flyte connection can solve this problem in the future.

Do you think it's ok to merge this implementation now and change it's behavior by flyte connection in the future?

[thomasjpfan]

Even with flyte connection, the key and group in this PR is still hardcoded, so I do not think it resolves my concern. Is the design of flyte connection to hard code the secret key?

[kumare3]

this is not right, i think we should ignore group in the SecretManager itself. Not in every plugin. this cannot scale and work universally

[kumare3]

cc @samhita-alla can you please comment and help here

[kumare3]

Lgtm - cc @thomasjpfan and @katrogan

[Future-Outlier]

It works now