Skip to content

Commit 5bc65ce

Browse files
committed
fix a test
1 parent 61d2435 commit 5bc65ce

File tree

1 file changed

+10
-1
lines changed

1 file changed

+10
-1
lines changed

bigframes/display/anywidget.py

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
from typing import Any, cast, Dict, Iterator, List, Optional, Type
2222
import uuid
2323

24+
import google.api_core.exceptions
2425
import pandas as pd
2526

2627
import bigframes
@@ -165,8 +166,16 @@ def _get_next_batch(self) -> bool:
165166
batch = next(iterator)
166167
self._cached_batches.append(batch)
167168
return True
168-
except StopIteration:
169+
except (
170+
StopIteration,
171+
google.api_core.exceptions.GoogleAPICallError,
172+
TypeError,
173+
ValueError,
174+
) as e:
169175
self._all_data_loaded = True
176+
if not isinstance(e, StopIteration):
177+
# If we fail to get a batch, assume no more data is available.
178+
self.row_count = 0
170179
return False
171180

172181
@property

0 commit comments

Comments
 (0)