chore: import upstream snapshot with attribution

python/ray/data/tests/datasource/test_snowflake.py

@@ -0,0 +1,120 @@
+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__]))