Skip to content

Commit

Permalink
SNOW-1539600: Refactor structured type tests to explicitly set params (
Browse files Browse the repository at this point in the history
  • Loading branch information
sfc-gh-jrose authored Jul 19, 2024
1 parent a5ae8ed commit e9393a6
Showing 1 changed file with 20 additions and 3 deletions.
23 changes: 20 additions & 3 deletions test/integ/test_arrow_result.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@
import os
import random
import re
from contextlib import contextmanager
from datetime import date, datetime, time, timedelta, timezone

import numpy
Expand Down Expand Up @@ -198,6 +199,22 @@
]


@contextmanager
def structured_type_wrapped_conn(conn_cnx):
parameters = {}
if STRUCTURED_TYPES_SUPPORTED:
parameters = {
"python_connector_query_result_format": "arrow",
"ENABLE_STRUCTURED_TYPES_IN_CLIENT_RESPONSE": True,
"ENABLE_STRUCTURED_TYPES_NATIVE_ARROW_FORMAT": True,
"FORCE_ENABLE_STRUCTURED_TYPES_NATIVE_ARROW_FORMAT": True,
"IGNORE_CLIENT_VESRION_IN_STRUCTURED_TYPES_RESPONSE": True,
}

with conn_cnx(session_parameters=parameters) as conn:
yield conn


def serialize(value):
if isinstance(value, bytearray):
return value.hex()
Expand All @@ -214,7 +231,7 @@ def verify_datatypes(
conn_cnx, query, examples, schema, iceberg=False, pandas=False, deserialize=False
):
table_name = f"arrow_datatype_test_verifaction_table_{random_string(5)}"
with conn_cnx() as conn:
with structured_type_wrapped_conn(conn_cnx) as conn:
try:
conn.cursor().execute("alter session set use_cached_result=false")
iceberg_table, iceberg_config = (
Expand Down Expand Up @@ -271,7 +288,7 @@ def pandas_verify(cur, data, deserialize):
@pytest.mark.parametrize("datatype", ICEBERG_UNSUPPORTED_TYPES)
def test_iceberg_negative(datatype, conn_cnx):
table_name = f"arrow_datatype_test_verifaction_table_{random_string(5)}"
with conn_cnx() as conn:
with structured_type_wrapped_conn(conn_cnx) as conn:
try:
with pytest.raises(ProgrammingError):
conn.cursor().execute(
Expand Down Expand Up @@ -349,7 +366,7 @@ def test_structured_type_binds(conn_cnx):
json_data = [json.dumps(d) for d in data]
schema = "(num number, arr_b array(boolean), map map(varchar, int), obj object(city varchar, population float), arr_f array(float))"
table_name = f"arrow_structured_type_binds_test_{random_string(5)}"
with conn_cnx() as conn:
with structured_type_wrapped_conn(conn_cnx) as conn:
try:
conn.cursor().execute("alter session set enable_bind_stage_v2=Enable")
conn.cursor().execute(f"create table if not exists {table_name} {schema}")
Expand Down

0 comments on commit e9393a6

Please sign in to comment.