Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: DB API cursors are now iterable #618

Merged
merged 2 commits into from Apr 16, 2021
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
2 changes: 1 addition & 1 deletion google/cloud/bigquery/dbapi/_helpers.py
Expand Up @@ -276,7 +276,7 @@ def decorate_public_methods(klass):
"""Apply ``_raise_on_closed()`` decorator to public instance methods.
"""
for name in dir(klass):
if name.startswith("_"):
if name.startswith("_") and name != "__iter__":
continue

member = getattr(klass, name)
Expand Down
4 changes: 4 additions & 0 deletions google/cloud/bigquery/dbapi/cursor.py
Expand Up @@ -365,6 +365,10 @@ def setinputsizes(self, sizes):
def setoutputsize(self, size, column=None):
"""No-op, but for consistency raise an error if cursor is closed."""

def __iter__(self):
self._try_fetch()
return iter(self._query_data)


def _format_operation_list(operation, parameters):
"""Formats parameters in operation in the way BigQuery expects.
Expand Down
24 changes: 24 additions & 0 deletions tests/unit/test_dbapi_cursor.py
Expand Up @@ -178,6 +178,7 @@ def test_raises_error_if_closed(self):
"fetchone",
"setinputsizes",
"setoutputsize",
"__iter__",
)

for method in method_names:
Expand Down Expand Up @@ -611,6 +612,29 @@ def test_executemany_w_dml(self):
self.assertIsNone(cursor.description)
self.assertEqual(cursor.rowcount, 12)

def test_is_iterable(self):
from google.cloud.bigquery import dbapi

connection = dbapi.connect(
self._mock_client(rows=[("hello", "there", 7), ("good", "bye", -3)])
)
cursor = connection.cursor()
cursor.execute("SELECT foo, bar, baz FROM hello_world WHERE baz < 42;")

rows_iter = iter(cursor)

row = next(rows_iter)
self.assertEqual(row, ("hello", "there", 7))
row = next(rows_iter)
self.assertEqual(row, ("good", "bye", -3))
self.assertRaises(StopIteration, next, rows_iter)

self.assertEqual(
list(cursor),
[],
"Iterating again over the same results should produce no rows.",
)

def test__format_operation_w_dict(self):
from google.cloud.bigquery.dbapi import cursor

Expand Down