Skip to content

Commit 9cfaa16

Browse files
committed
Fix columns property to handle lazy SegmentIterator
Commit a2a1001 made fetch() return a lazy SegmentIterator for spooled results but did not update the columns property which still used `+= self.fetch()`. This raises TypeError when _result.rows is an itertools.chain object and silently materializes the iterator when _result.rows is a list. Apply the same isinstance branching already present in execute().
1 parent a2a1001 commit 9cfaa16

2 files changed

Lines changed: 42 additions & 1 deletion

File tree

tests/integration/test_dbapi_integration.py

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1973,6 +1973,33 @@ def test_spooled_segments_iterator_protocol(trino_connection):
19731973
assert count == 60175, f"Expected 60175 rows, got {count}"
19741974

19751975

1976+
@pytest.mark.skipif(
1977+
trino_version() <= 466,
1978+
reason="spooling protocol was introduced in version 466"
1979+
)
1980+
def test_spooled_segments_lazy_description(trino_connection):
1981+
"""Verify that accessing cursor.description does not materialize the lazy spooled iterator."""
1982+
if trino_connection._client_session.encoding is None:
1983+
pytest.skip("spooling requires an encoding")
1984+
1985+
cur = trino_connection.cursor()
1986+
cur.execute("SELECT * FROM tpch.tiny.lineitem")
1987+
1988+
assert not isinstance(cur._query._result._rows, list), (
1989+
f"Expected lazy iterator for spooled results, got {type(cur._query._result._rows)}"
1990+
)
1991+
1992+
desc = cur.description
1993+
assert desc is not None
1994+
assert len(desc) > 0
1995+
1996+
assert not isinstance(cur._query._result._rows, list), (
1997+
f"Expected lazy iterator after description access, got {type(cur._query._result._rows)}"
1998+
)
1999+
2000+
assert len(cur.fetchall()) == 60175
2001+
2002+
19762003
def get_cursor(legacy_prepared_statements, run_trino):
19772004
host, port = run_trino
19782005

trino/client.py

Lines changed: 15 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -880,7 +880,21 @@ def columns(self):
880880
while not self._columns and not self.finished and not self.cancelled:
881881
# Columns are not returned immediately after query is submitted.
882882
# Continue fetching data until columns information is available and push fetched rows into buffer.
883-
self._result.rows += self.fetch()
883+
#
884+
# Two protocols produce rows differently:
885+
# - Direct: fetch() returns a list - accumulate into the existing list.
886+
# - Spooling: fetch() returns a lazy iterator - replace rows and stop,
887+
# because we cannot cheaply check iterator length.
888+
new_rows = self.fetch()
889+
if isinstance(new_rows, list):
890+
self._result.rows += new_rows
891+
else:
892+
try:
893+
first_row = next(new_rows)
894+
self._result.rows = itertools.chain([first_row], new_rows)
895+
break
896+
except StopIteration:
897+
self._result.rows = []
884898
return self._columns
885899

886900
@property

0 commit comments

Comments
 (0)