wehub-resource-sync
121 lines
3.8 KiB
Python
121 lines
3.8 KiB
Python
import base64
|
|
import os
|
|
import random
|
|
import string
|
|
from typing import Any, Dict, List, Tuple
|
|
|
|
import pytest
|
|
from snowflake.connector import connect
|
|
|
|
import ray
|
|
from ray.tests.conftest import * # noqa
|
|
|
|
# Note: Snowflake secrets are only used in postmerge authenticated tests.
|
|
|
|
|
|
@pytest.fixture
|
|
def connection_parameters():
|
|
private_key_b64 = os.getenv("SNOWFLAKE_PRIVATE_KEY")
|
|
private_key_bytes = base64.b64decode(private_key_b64)
|
|
parameters = {
|
|
"user": os.getenv("SNOWFLAKE_USER"),
|
|
"account": os.getenv("SNOWFLAKE_ACCOUNT"),
|
|
"database": os.getenv("SNOWFLAKE_DATABASE"),
|
|
"schema": os.getenv("SNOWFLAKE_SCHEMA"),
|
|
"warehouse": os.getenv("SNOWFLAKE_WAREHOUSE"),
|
|
"private_key": private_key_bytes,
|
|
}
|
|
|
|
yield parameters
|
|
|
|
|
|
@pytest.fixture
|
|
def temp_table(connection_parameters):
|
|
table_name = "".join([random.choice(string.ascii_uppercase) for _ in range(8)])
|
|
|
|
yield table_name
|
|
|
|
with connect(**connection_parameters) as connection, connection.cursor() as cursor:
|
|
cursor.execute(f"DROP TABLE IF EXISTS {table_name}")
|
|
connection.commit()
|
|
|
|
|
|
@pytest.mark.needs_credentials
|
|
def test_read(ray_start_regular_shared, connection_parameters):
|
|
# This query fetches a small dataset with a variety of column types.
|
|
query = "SELECT * FROM SNOWFLAKE_SAMPLE_DATA.TPCDS_SF100TCL.CALL_CENTER"
|
|
|
|
# Read the data and check contents.
|
|
dataset = ray.data.read_snowflake(query, connection_parameters)
|
|
actual_column_names = dataset.schema().names
|
|
actual_rows = [tuple(row.values()) for row in dataset.take_all()]
|
|
expected_column_names, expected_rows = execute(query, connection_parameters)
|
|
|
|
assert actual_column_names == expected_column_names
|
|
assert sorted(actual_rows) == sorted(expected_rows)
|
|
|
|
|
|
@pytest.mark.needs_credentials
|
|
def test_write(ray_start_regular_shared, temp_table, connection_parameters):
|
|
expected_column_names = ["title", "year", "score"]
|
|
expected_rows = [
|
|
("Monty Python and the Holy Grail", 1975, 8.2),
|
|
("And Now for Something Completely Different", 1971, 7.5),
|
|
]
|
|
|
|
# Create the table first
|
|
create_table_sql = f"""
|
|
CREATE TABLE IF NOT EXISTS {temp_table} (
|
|
"title" VARCHAR(255),
|
|
"year" INTEGER,
|
|
"score" FLOAT
|
|
)
|
|
"""
|
|
execute(create_table_sql, connection_parameters)
|
|
|
|
items = [dict(zip(expected_column_names, row)) for row in expected_rows]
|
|
dataset = ray.data.from_items(items)
|
|
|
|
dataset.write_snowflake(temp_table, connection_parameters)
|
|
actual_column_names, actual_rows = execute(
|
|
f"SELECT * FROM {temp_table}", connection_parameters
|
|
)
|
|
|
|
assert actual_column_names == expected_column_names
|
|
assert sorted(actual_rows) == sorted(expected_rows)
|
|
|
|
|
|
@pytest.mark.needs_credentials
|
|
def execute(
|
|
query: str, connection_parameters: Dict[str, str]
|
|
) -> Tuple[List[str], List[Tuple[Any]]]:
|
|
"""Execute a query on Snowflake and return the resulting data.
|
|
|
|
Args:
|
|
query: The SQL query to execute.
|
|
connection_parameters: Connection params for snowflake.
|
|
|
|
Returns:
|
|
A two-tuple containing the column names and rows.
|
|
"""
|
|
with connect(**connection_parameters) as connection, connection.cursor() as cursor:
|
|
cursor.execute(query)
|
|
column_names = [column_metadata.name for column_metadata in cursor.description]
|
|
rows = cursor.fetchall()
|
|
|
|
# TODO(mowen): Figure out how to actually handle the Decimal objects, we don't
|
|
# want a divergenece in behavior here.
|
|
# The Snowflake Python Connector represents numbers as `Decimal` objects.
|
|
# rows = [
|
|
# tuple(float(value) if isinstance(value, Decimal) else value for value in row)
|
|
# for row in rows
|
|
# ]
|
|
|
|
return column_names, rows
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
sys.exit(pytest.main(["-v", __file__]))
|