-
Notifications
You must be signed in to change notification settings - Fork 287
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
* fix snowflake agent bug Signed-off-by: Future-Outlier <[email protected]> * a work version Signed-off-by: Future-Outlier <[email protected]> * Snowflake work version Signed-off-by: Future-Outlier <[email protected]> * fix secret encode Signed-off-by: Future-Outlier <[email protected]> * all works, I am so happy Signed-off-by: Future-Outlier <[email protected]> * improve additional protocol Signed-off-by: Future-Outlier <[email protected]> * fix tests Signed-off-by: Future-Outlier <[email protected]> * Fix Tests Signed-off-by: Future-Outlier <[email protected]> * update agent Signed-off-by: Kevin Su <[email protected]> * Add snowflake test Signed-off-by: Kevin Su <[email protected]> * nit Signed-off-by: Kevin Su <[email protected]> * sd Signed-off-by: Kevin Su <[email protected]> * snowflake loglinks Signed-off-by: Future-Outlier <[email protected]> * add metadata Signed-off-by: Future-Outlier <[email protected]> * secret Signed-off-by: Kevin Su <[email protected]> * nit Signed-off-by: Kevin Su <[email protected]> * remove table Signed-off-by: Future-Outlier <[email protected]> * add comment for get private key Signed-off-by: Future-Outlier <[email protected]> * update comments: Signed-off-by: Future-Outlier <[email protected]> * Fix Tests Signed-off-by: Future-Outlier <[email protected]> * update comments Signed-off-by: Future-Outlier <[email protected]> * update comments Signed-off-by: Future-Outlier <[email protected]> * Better Secrets Signed-off-by: Future-Outlier <[email protected]> * use union secret Signed-off-by: Future-Outlier <[email protected]> * Update Changes Signed-off-by: Future-Outlier <[email protected]> * use if not get_plugin().secret_requires_group() Signed-off-by: Future-Outlier <[email protected]> * Use Union SDK Signed-off-by: Future-Outlier <[email protected]> * Update Signed-off-by: Future-Outlier <[email protected]> * Fix Secrets Signed-off-by: Future-Outlier <[email protected]> * Fix Secrets Signed-off-by: Future-Outlier <[email protected]> * remove pacakge.json Signed-off-by: Future-Outlier <[email protected]> * lint Signed-off-by: Future-Outlier <[email protected]> * add snowflake-connector-python Signed-off-by: Future-Outlier <[email protected]> * fix test_snowflake Signed-off-by: Future-Outlier <[email protected]> * Try to fix tests Signed-off-by: Future-Outlier <[email protected]> * fix tests Signed-off-by: Future-Outlier <[email protected]> * Try Fix snowflake Import Signed-off-by: Future-Outlier <[email protected]> * snowflake test passed Signed-off-by: Future-Outlier <[email protected]> --------- Signed-off-by: Future-Outlier <[email protected]> Signed-off-by: Kevin Su <[email protected]> Co-authored-by: Kevin Su <[email protected]>
- Loading branch information
1 parent
4f96b33
commit 1b67f16
Showing
13 changed files
with
298 additions
and
55 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.