|
17 | 17 |
|
18 | 18 |
|
19 | 19 | from collections import deque
|
| 20 | +from warnings import warn |
20 | 21 |
|
21 | 22 | from ..._async_compat.util import AsyncUtil
|
22 | 23 | from ...data import DataDehydrator
|
@@ -248,11 +249,11 @@ async def _buffer(self, n=None):
|
248 | 249 | record_buffer.append(record)
|
249 | 250 | if n is not None and len(record_buffer) >= n:
|
250 | 251 | break
|
251 |
| - self._exhausted = False |
252 | 252 | if n is None:
|
253 | 253 | self._record_buffer = record_buffer
|
254 | 254 | else:
|
255 | 255 | self._record_buffer.extend(record_buffer)
|
| 256 | + self._exhausted = not self._record_buffer |
256 | 257 |
|
257 | 258 | async def _buffer_all(self):
|
258 | 259 | """Sets the Result object in an detached state by fetching all records
|
@@ -286,12 +287,20 @@ def keys(self):
|
286 | 287 | """
|
287 | 288 | return self._keys
|
288 | 289 |
|
| 290 | + async def _exhaust(self): |
| 291 | + # Exhaust the result, ditching all remaining records. |
| 292 | + if not self._exhausted: |
| 293 | + self._discarding = True |
| 294 | + self._record_buffer.clear() |
| 295 | + async for _ in self: |
| 296 | + pass |
| 297 | + |
289 | 298 | async def _tx_end(self):
|
290 | 299 | # Handle closure of the associated transaction.
|
291 | 300 | #
|
292 | 301 | # This will consume the result and mark it at out of scope.
|
293 | 302 | # Subsequent calls to `next` will raise a ResultConsumedError.
|
294 |
| - await self.consume() |
| 303 | + await self._exhaust() |
295 | 304 | self._out_of_scope = True
|
296 | 305 |
|
297 | 306 | async def consume(self):
|
@@ -330,42 +339,85 @@ async def get_two_tx(tx):
|
330 | 339 |
|
331 | 340 | :returns: The :class:`neo4j.ResultSummary` for this result
|
332 | 341 | """
|
333 |
| - if self._exhausted is False: |
334 |
| - self._discarding = True |
335 |
| - async for _ in self: |
336 |
| - pass |
| 342 | + if self._exhausted: |
| 343 | + if self._out_of_scope: |
| 344 | + raise ResultConsumedError(self, _RESULT_OUT_OF_SCOPE_ERROR) |
| 345 | + if self._consumed: |
| 346 | + raise ResultConsumedError(self, _RESULT_CONSUMED_ERROR) |
| 347 | + else: |
| 348 | + await self._exhaust() |
337 | 349 |
|
338 | 350 | summary = self._obtain_summary()
|
339 | 351 | self._consumed = True
|
340 | 352 | return summary
|
341 | 353 |
|
342 |
| - async def single(self): |
343 |
| - """Obtain the next and only remaining record from this result if available else return None. |
| 354 | + async def single(self, strict=False): |
| 355 | + """Obtain the next and only remaining record or None. |
| 356 | +
|
344 | 357 | Calling this method always exhausts the result.
|
345 | 358 |
|
346 | 359 | A warning is generated if more than one record is available but
|
347 | 360 | the first of these is still returned.
|
348 | 361 |
|
349 |
| - :returns: the next :class:`neo4j.AsyncRecord`. |
| 362 | + :param strict: |
| 363 | + If :const:`True`, raise a :class:`neo4j.ResultNotSingleError` |
| 364 | + instead of returning None if there is more than one record or |
| 365 | + warning if there are more than 1 record. |
| 366 | + :const:`False` by default. |
| 367 | + :type strict: bool |
350 | 368 |
|
351 |
| - :raises ResultNotSingleError: if not exactly one record is available. |
352 |
| - :raises ResultConsumedError: if the transaction from which this result was |
353 |
| - obtained has been closed. |
| 369 | + :returns: the next :class:`neo4j.Record` or :const:`None` if none remain |
| 370 | + :warns: if more than one record is available |
| 371 | +
|
| 372 | + :raises ResultNotSingleError: |
| 373 | + If ``strict=True`` and not exactly one record is available. |
| 374 | + :raises ResultConsumedError: if the transaction from which this result |
| 375 | + was obtained has been closed. |
| 376 | +
|
| 377 | + .. versionchanged:: 5.0 |
| 378 | + Added ``strict`` parameter. |
354 | 379 | """
|
355 | 380 | await self._buffer(2)
|
356 |
| - if not self._record_buffer: |
| 381 | + buffer = self._record_buffer |
| 382 | + self._record_buffer = deque() |
| 383 | + await self._exhaust() |
| 384 | + if not buffer: |
| 385 | + if not strict: |
| 386 | + return None |
357 | 387 | raise ResultNotSingleError(
|
358 | 388 | self,
|
359 | 389 | "No records found. "
|
360 | 390 | "Make sure your query returns exactly one record."
|
361 | 391 | )
|
362 |
| - elif len(self._record_buffer) > 1: |
363 |
| - raise ResultNotSingleError( |
364 |
| - self, |
365 |
| - "More than one record found. " |
366 |
| - "Make sure your query returns exactly one record." |
367 |
| - ) |
368 |
| - return self._record_buffer.popleft() |
| 392 | + elif len(buffer) > 1: |
| 393 | + res = buffer.popleft() |
| 394 | + if not strict: |
| 395 | + warn("Expected a result with a single record, " |
| 396 | + "but found multiple.") |
| 397 | + return res |
| 398 | + else: |
| 399 | + raise ResultNotSingleError( |
| 400 | + self, |
| 401 | + "More than one record found. " |
| 402 | + "Make sure your query returns exactly one record." |
| 403 | + ) |
| 404 | + return buffer.popleft() |
| 405 | + |
| 406 | + async def fetch(self, n): |
| 407 | + """Obtain up to n records from this result. |
| 408 | +
|
| 409 | + :param n: the maximum number of records to fetch. |
| 410 | + :type n: int |
| 411 | +
|
| 412 | + :returns: list of :class:`neo4j.AsyncRecord` |
| 413 | +
|
| 414 | + .. versionadded:: 5.0 |
| 415 | + """ |
| 416 | + await self._buffer(n) |
| 417 | + return [ |
| 418 | + self._record_buffer.popleft() |
| 419 | + for _ in range(min(n, len(self._record_buffer))) |
| 420 | + ] |
369 | 421 |
|
370 | 422 | async def peek(self):
|
371 | 423 | """Obtain the next record from this result without consuming it.
|
|
0 commit comments