Skip to content

Commit 50cc1e2

Browse files
varun-edachali-dbxvarun-edachali-dbx
committed
run small queries with SEA during integration tests
Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com>
1 parent 0ce144d commit 50cc1e2

File tree

2 files changed

+162
-52
lines changed

2 files changed

+162
-52
lines changed

‎tests/e2e/test_driver.py‎

Lines changed: 108 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -196,10 +196,14 @@ def test_execute_async__long_running(self):
196196

197197
assert result[0].asDict() == {"count(1)": 0}
198198

199-
def test_execute_async__small_result(self):
199+
@pytest.mark.parametrize("extra_params", [
200+
{},
201+
{"use_sea": True, "use_cloud_fetch": False}
202+
])
203+
def test_execute_async__small_result(self, extra_params):
200204
small_result_query = "SELECT 1"
201205

202-
with self.cursor() as cursor:
206+
with self.cursor(extra_params) as cursor:
203207
cursor.execute_async(small_result_query)
204208

205209
## Fake sleep for 5 secs
@@ -328,8 +332,12 @@ def test_incorrect_query_throws_exception(self):
328332
cursor.execute("CREATE TABLE IF NOT EXISTS TABLE table_234234234")
329333
assert "table_234234234" in str(cm.value)
330334

331-
def test_create_table_will_return_empty_result_set(self):
332-
with self.cursor({}) as cursor:
335+
@pytest.mark.parametrize("extra_params", [
336+
{},
337+
{"use_sea": True, "use_cloud_fetch": False}
338+
])
339+
def test_create_table_will_return_empty_result_set(self, extra_params):
340+
with self.cursor(extra_params) as cursor:
333341
table_name = "table_{uuid}".format(uuid=str(uuid4()).replace("-", "_"))
334342
try:
335343
cursor.execute(
@@ -341,8 +349,12 @@ def test_create_table_will_return_empty_result_set(self):
341349
finally:
342350
cursor.execute("DROP TABLE IF EXISTS {}".format(table_name))
343351

344-
def test_get_tables(self):
345-
with self.cursor({}) as cursor:
352+
@pytest.mark.parametrize("extra_params", [
353+
{},
354+
{"use_sea": True, "use_cloud_fetch": False}
355+
])
356+
def test_get_tables(self, extra_params):
357+
with self.cursor(extra_params) as cursor:
346358
table_name = "table_{uuid}".format(uuid=str(uuid4()).replace("-", "_"))
347359
table_names = [table_name + "_1", table_name + "_2"]
348360

@@ -387,8 +399,12 @@ def test_get_tables(self):
387399
for table in table_names:
388400
cursor.execute("DROP TABLE IF EXISTS {}".format(table))
389401

390-
def test_get_columns(self):
391-
with self.cursor({}) as cursor:
402+
@pytest.mark.parametrize("extra_params", [
403+
{},
404+
{"use_sea": True, "use_cloud_fetch": False}
405+
])
406+
def test_get_columns(self, extra_params):
407+
with self.cursor(extra_params) as cursor:
392408
table_name = "table_{uuid}".format(uuid=str(uuid4()).replace("-", "_"))
393409
table_names = [table_name + "_1", table_name + "_2"]
394410

@@ -474,8 +490,12 @@ def test_get_columns(self):
474490
for table in table_names:
475491
cursor.execute("DROP TABLE IF EXISTS {}".format(table))
476492

477-
def test_escape_single_quotes(self):
478-
with self.cursor({}) as cursor:
493+
@pytest.mark.parametrize("extra_params", [
494+
{},
495+
{"use_sea": True, "use_cloud_fetch": False}
496+
])
497+
def test_escape_single_quotes(self, extra_params):
498+
with self.cursor(extra_params) as cursor:
479499
table_name = "table_{uuid}".format(uuid=str(uuid4()).replace("-", "_"))
480500
# Test escape syntax directly
481501
cursor.execute(
@@ -499,8 +519,12 @@ def test_escape_single_quotes(self):
499519
rows = cursor.fetchall()
500520
assert rows[0]["col_1"] == "you're"
501521

502-
def test_get_schemas(self):
503-
with self.cursor({}) as cursor:
522+
@pytest.mark.parametrize("extra_params", [
523+
{},
524+
{"use_sea": True, "use_cloud_fetch": False}
525+
])
526+
def test_get_schemas(self, extra_params):
527+
with self.cursor(extra_params) as cursor:
504528
database_name = "db_{uuid}".format(uuid=str(uuid4()).replace("-", "_"))
505529
try:
506530
cursor.execute("CREATE DATABASE IF NOT EXISTS {}".format(database_name))
@@ -517,8 +541,12 @@ def test_get_schemas(self):
517541
finally:
518542
cursor.execute("DROP DATABASE IF EXISTS {}".format(database_name))
519543

520-
def test_get_catalogs(self):
521-
with self.cursor({}) as cursor:
544+
@pytest.mark.parametrize("extra_params", [
545+
{},
546+
{"use_sea": True, "use_cloud_fetch": False}
547+
])
548+
def test_get_catalogs(self, extra_params):
549+
with self.cursor(extra_params) as cursor:
522550
cursor.catalogs()
523551
cursor.fetchall()
524552
catalogs_desc = cursor.description
@@ -527,27 +555,39 @@ def test_get_catalogs(self):
527555
]
528556

529557
@skipUnless(pysql_supports_arrow(), "arrow test need arrow support")
530-
def test_get_arrow(self):
558+
@pytest.mark.parametrize("extra_params", [
559+
{},
560+
{"use_sea": True, "use_cloud_fetch": False}
561+
])
562+
def test_get_arrow(self, extra_params):
531563
# These tests are quite light weight as the arrow fetch methods are used internally
532564
# by everything else
533-
with self.cursor({}) as cursor:
565+
with self.cursor(extra_params) as cursor:
534566
cursor.execute("SELECT * FROM range(10)")
535567
table_1 = cursor.fetchmany_arrow(1).to_pydict()
536568
assert table_1 == OrderedDict([("id", [0])])
537569

538570
table_2 = cursor.fetchall_arrow().to_pydict()
539571
assert table_2 == OrderedDict([("id", [1, 2, 3, 4, 5, 6, 7, 8, 9])])
540572

541-
def test_unicode(self):
573+
@pytest.mark.parametrize("extra_params", [
574+
{},
575+
{"use_sea": True, "use_cloud_fetch": False}
576+
])
577+
def test_unicode(self, extra_params):
542578
unicode_str = "数据砖"
543-
with self.cursor({}) as cursor:
579+
with self.cursor(extra_params) as cursor:
544580
cursor.execute("SELECT '{}'".format(unicode_str))
545581
results = cursor.fetchall()
546582
assert len(results) == 1 and len(results[0]) == 1
547583
assert results[0][0] == unicode_str
548584

549-
def test_cancel_during_execute(self):
550-
with self.cursor({}) as cursor:
585+
@pytest.mark.parametrize("extra_params", [
586+
{},
587+
{"use_sea": True, "use_cloud_fetch": False}
588+
])
589+
def test_cancel_during_execute(self, extra_params):
590+
with self.cursor(extra_params) as cursor:
551591

552592
def execute_really_long_query():
553593
cursor.execute(
@@ -578,8 +618,12 @@ def execute_really_long_query():
578618
assert len(cursor.fetchall()) == 3
579619

580620
@skipIf(pysql_has_version("<", "2"), "requires pysql v2")
581-
def test_can_execute_command_after_failure(self):
582-
with self.cursor({}) as cursor:
621+
@pytest.mark.parametrize("extra_params", [
622+
{},
623+
{"use_sea": True, "use_cloud_fetch": False}
624+
])
625+
def test_can_execute_command_after_failure(self, extra_params):
626+
with self.cursor(extra_params) as cursor:
583627
with pytest.raises(DatabaseError):
584628
cursor.execute("this is a sytnax error")
585629

@@ -589,8 +633,12 @@ def test_can_execute_command_after_failure(self):
589633
self.assertEqualRowValues(res, [[1]])
590634

591635
@skipIf(pysql_has_version("<", "2"), "requires pysql v2")
592-
def test_can_execute_command_after_success(self):
593-
with self.cursor({}) as cursor:
636+
@pytest.mark.parametrize("extra_params", [
637+
{},
638+
{"use_sea": True, "use_cloud_fetch": False}
639+
])
640+
def test_can_execute_command_after_success(self, extra_params):
641+
with self.cursor(extra_params) as cursor:
594642
cursor.execute("SELECT 1;")
595643
cursor.execute("SELECT 2;")
596644

@@ -602,8 +650,12 @@ def generate_multi_row_query(self):
602650
return query
603651

604652
@skipIf(pysql_has_version("<", "2"), "requires pysql v2")
605-
def test_fetchone(self):
606-
with self.cursor({}) as cursor:
653+
@pytest.mark.parametrize("extra_params", [
654+
{},
655+
{"use_sea": True, "use_cloud_fetch": False}
656+
])
657+
def test_fetchone(self, extra_params):
658+
with self.cursor(extra_params) as cursor:
607659
query = self.generate_multi_row_query()
608660
cursor.execute(query)
609661

@@ -614,8 +666,12 @@ def test_fetchone(self):
614666
assert cursor.fetchone() == None
615667

616668
@skipIf(pysql_has_version("<", "2"), "requires pysql v2")
617-
def test_fetchall(self):
618-
with self.cursor({}) as cursor:
669+
@pytest.mark.parametrize("extra_params", [
670+
{},
671+
{"use_sea": True, "use_cloud_fetch": False}
672+
])
673+
def test_fetchall(self, extra_params):
674+
with self.cursor(extra_params) as cursor:
619675
query = self.generate_multi_row_query()
620676
cursor.execute(query)
621677

@@ -624,26 +680,38 @@ def test_fetchall(self):
624680
assert cursor.fetchone() == None
625681

626682
@skipIf(pysql_has_version("<", "2"), "requires pysql v2")
627-
def test_fetchmany_when_stride_fits(self):
628-
with self.cursor({}) as cursor:
683+
@pytest.mark.parametrize("extra_params", [
684+
{},
685+
{"use_sea": True, "use_cloud_fetch": False}
686+
])
687+
def test_fetchmany_when_stride_fits(self, extra_params):
688+
with self.cursor(extra_params) as cursor:
629689
query = "SELECT * FROM range(4)"
630690
cursor.execute(query)
631691

632692
self.assertEqualRowValues(cursor.fetchmany(2), [[0], [1]])
633693
self.assertEqualRowValues(cursor.fetchmany(2), [[2], [3]])
634694

635695
@skipIf(pysql_has_version("<", "2"), "requires pysql v2")
636-
def test_fetchmany_in_excess(self):
637-
with self.cursor({}) as cursor:
696+
@pytest.mark.parametrize("extra_params", [
697+
{},
698+
{"use_sea": True, "use_cloud_fetch": False}
699+
])
700+
def test_fetchmany_in_excess(self, extra_params):
701+
with self.cursor(extra_params) as cursor:
638702
query = "SELECT * FROM range(4)"
639703
cursor.execute(query)
640704

641705
self.assertEqualRowValues(cursor.fetchmany(3), [[0], [1], [2]])
642706
self.assertEqualRowValues(cursor.fetchmany(3), [[3]])
643707

644708
@skipIf(pysql_has_version("<", "2"), "requires pysql v2")
645-
def test_iterator_api(self):
646-
with self.cursor({}) as cursor:
709+
@pytest.mark.parametrize("extra_params", [
710+
{},
711+
{"use_sea": True, "use_cloud_fetch": False}
712+
])
713+
def test_iterator_api(self, extra_params):
714+
with self.cursor(extra_params) as cursor:
647715
query = "SELECT * FROM range(4)"
648716
cursor.execute(query)
649717

@@ -803,8 +871,12 @@ def test_decimal_not_returned_as_strings_arrow(self):
803871
assert pyarrow.types.is_decimal(decimal_type)
804872

805873
@skipUnless(pysql_supports_arrow(), "arrow test needs arrow support")
806-
def test_catalogs_returns_arrow_table(self):
807-
with self.cursor() as cursor:
874+
@pytest.mark.parametrize("extra_params", [
875+
{},
876+
{"use_sea": True, "use_cloud_fetch": False}
877+
])
878+
def test_catalogs_returns_arrow_table(self, extra_params):
879+
with self.cursor(extra_params) as cursor:
808880
cursor.catalogs()
809881
results = cursor.fetchall_arrow()
810882
assert isinstance(results, pyarrow.Table)

0 commit comments

Comments
 (0)