<#1799 Snowflake agent> Pull request opened by <hh...
# flyte-github
a
#1799 Snowflake agent Pull request opened by hhcs9527 TL;DR Support Snowflake with an agent, let the user is able to run it either locally or remotely. Type ☐ Bug Fix ☑︎ Feature ☐ Plugin Are all requirements met? ☑︎ Code completed ☐ Smoke tested ☑︎ Unit tests added ☐ Code documentation added ☐ Any pending items have an associated Issue Complete description Support Snowflake with an agent, and let the user be able to run it either locally or remotely. Here is an example.
Copy code
try:
    from typing import Annotated
except ImportError:
    from typing_extensions import Annotated

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

# we need to type our account, and user, and prepare the private key to build a connection.
# Note that, the table field is for building the correct uri for future reference.

snowflake_task_templatized_query2 = SnowflakeTask(
    name="sf",
    # Define inputs as well as their types that can be used to customize the query.1
    inputs=kwtypes(number=int),
    output_schema_type=StructuredDataset ,
    task_config=SnowflakeConfig( user='hhcs9527',
    account='ed61997.europe-west4.gcp',
    database='AGENT',
    schema='PEOPLE',
    table='INTRO',
    warehouse='COMPUTE_WH'),
    query_template="SELECT AGE, NAME from INTRO limit %(number)s",
)

@task
def convert_bq_table_to_pandas_dataframe(sd: DogeCoinDataset) -> pd.DataFrame:
    return sd.open(pd.DataFrame).all()

@workflow
def sf_query_wf(number: int) -> pd.DataFrame:
    sd = snowflake_task_templatized_query2(number=number)
    return convert_bq_table_to_pandas_dataframe(sd=sd)

if __name__ == "__main__":
    print(sf_query_wf(number=2)
Tracking Issue Flyte Agent Ecosystem Follow-up issue Add encoder/decoder in structureDataset for snowflake. flyteorg/flytekit ✅ All checks have passed 30/30 successful checks
#1799 Snowflake agent Pull request ready for review by hhcs9527 flyteorg/flytekit