344 lines
11 KiB
Python
344 lines
11 KiB
Python
# WARNING: this file is auto-generated by 'async_to_sync.py'
|
|
# from the original file 'cursor_async.py'
|
|
# DO NOT CHANGE! Change the original file instead.
|
|
"""
|
|
Psycopg Cursor object.
|
|
"""
|
|
|
|
# Copyright (C) 2020 The Psycopg Team
|
|
|
|
from __future__ import annotations
|
|
|
|
from types import TracebackType
|
|
from typing import TYPE_CHECKING, Any, overload
|
|
from contextlib import contextmanager
|
|
from collections.abc import Iterable, Iterator
|
|
|
|
from . import errors as e
|
|
from . import pq
|
|
from .abc import Params, Query, QueryNoTemplate
|
|
from .copy import Copy, Writer
|
|
from .rows import Row, RowFactory, RowMaker
|
|
from ._compat import Self, Template
|
|
from ._pipeline import Pipeline
|
|
from ._cursor_base import BaseCursor
|
|
|
|
if TYPE_CHECKING:
|
|
from .connection import Connection
|
|
|
|
ACTIVE = pq.TransactionStatus.ACTIVE
|
|
|
|
|
|
class Cursor(BaseCursor["Connection[Any]", Row]):
|
|
__module__ = "psycopg"
|
|
__slots__ = ()
|
|
|
|
@overload
|
|
def __init__(self, connection: Connection[Row]): ...
|
|
|
|
@overload
|
|
def __init__(
|
|
self, connection: Connection[Any], *, row_factory: RowFactory[Row]
|
|
): ...
|
|
|
|
def __init__(
|
|
self, connection: Connection[Any], *, row_factory: RowFactory[Row] | None = None
|
|
):
|
|
super().__init__(connection)
|
|
self._row_factory = row_factory or connection.row_factory
|
|
|
|
def __enter__(self) -> Self:
|
|
return self
|
|
|
|
def __exit__(
|
|
self,
|
|
exc_type: type[BaseException] | None,
|
|
exc_val: BaseException | None,
|
|
exc_tb: TracebackType | None,
|
|
) -> None:
|
|
self.close()
|
|
|
|
def close(self) -> None:
|
|
"""
|
|
Close the current cursor and free associated resources.
|
|
"""
|
|
self._close()
|
|
|
|
@property
|
|
def row_factory(self) -> RowFactory[Row]:
|
|
"""Writable attribute to control how result rows are formed."""
|
|
return self._row_factory
|
|
|
|
@row_factory.setter
|
|
def row_factory(self, row_factory: RowFactory[Row]) -> None:
|
|
self._row_factory = row_factory
|
|
if self.pgresult:
|
|
self._make_row = row_factory(self)
|
|
|
|
def _make_row_maker(self) -> RowMaker[Row]:
|
|
return self._row_factory(self)
|
|
|
|
@overload
|
|
def execute(
|
|
self,
|
|
query: QueryNoTemplate,
|
|
params: Params | None = None,
|
|
*,
|
|
prepare: bool | None = None,
|
|
binary: bool | None = None,
|
|
) -> Self: ...
|
|
|
|
@overload
|
|
def execute(
|
|
self,
|
|
query: Template,
|
|
*,
|
|
prepare: bool | None = None,
|
|
binary: bool | None = None,
|
|
) -> Self: ...
|
|
|
|
def execute(
|
|
self,
|
|
query: Query,
|
|
params: Params | None = None,
|
|
*,
|
|
prepare: bool | None = None,
|
|
binary: bool | None = None,
|
|
) -> Self:
|
|
"""
|
|
Execute a query or command to the database.
|
|
"""
|
|
try:
|
|
with self._conn.lock:
|
|
self._conn.wait(
|
|
self._execute_gen(query, params, prepare=prepare, binary=binary)
|
|
)
|
|
except e._NO_TRACEBACK as ex:
|
|
raise ex.with_traceback(None)
|
|
return self
|
|
|
|
def executemany(
|
|
self, query: Query, params_seq: Iterable[Params], *, returning: bool = False
|
|
) -> None:
|
|
"""
|
|
Execute the same command with a sequence of input data.
|
|
"""
|
|
try:
|
|
with self._conn.lock:
|
|
if Pipeline.is_supported():
|
|
# If there is already a pipeline, ride it, in order to avoid
|
|
# sending unnecessary Sync.
|
|
if self._conn._pipeline:
|
|
self._conn.wait(
|
|
self._executemany_gen_pipeline(query, params_seq, returning)
|
|
)
|
|
else:
|
|
# Otherwise, make a new one
|
|
with self._conn._pipeline_nolock():
|
|
self._conn.wait(
|
|
self._executemany_gen_pipeline(
|
|
query, params_seq, returning
|
|
)
|
|
)
|
|
else:
|
|
self._conn.wait(
|
|
self._executemany_gen_no_pipeline(query, params_seq, returning)
|
|
)
|
|
except e._NO_TRACEBACK as ex:
|
|
raise ex.with_traceback(None)
|
|
|
|
def stream(
|
|
self,
|
|
query: Query,
|
|
params: Params | None = None,
|
|
*,
|
|
binary: bool | None = None,
|
|
size: int = 1,
|
|
) -> Iterator[Row]:
|
|
"""
|
|
Iterate row-by-row on a result from the database.
|
|
|
|
:param size: if greater than 1, results will be retrieved by chunks of
|
|
this size from the server (but still yielded row-by-row); this is only
|
|
available from version 17 of the libpq.
|
|
"""
|
|
if self._pgconn.pipeline_status:
|
|
raise e.ProgrammingError("stream() cannot be used in pipeline mode")
|
|
|
|
with self._conn.lock:
|
|
try:
|
|
self._conn.wait(
|
|
self._stream_send_gen(query, params, binary=binary, size=size)
|
|
)
|
|
first = True
|
|
while res := self._conn.wait(self._stream_fetchone_gen(first)):
|
|
for pos in range(res.ntuples):
|
|
yield self._tx.load_row(pos, self._make_row)
|
|
first = False
|
|
except e._NO_TRACEBACK as ex:
|
|
raise ex.with_traceback(None)
|
|
finally:
|
|
if self._pgconn.transaction_status == ACTIVE:
|
|
# Try to cancel the query, then consume the results
|
|
# already received.
|
|
self._conn._try_cancel()
|
|
try:
|
|
while self._conn.wait(self._stream_fetchone_gen(first=False)):
|
|
pass
|
|
except Exception:
|
|
pass
|
|
|
|
# Try to get out of ACTIVE state. Just do a single attempt, which
|
|
# should work to recover from an error or query cancelled.
|
|
try:
|
|
self._conn.wait(self._stream_fetchone_gen(first=False))
|
|
except Exception:
|
|
pass
|
|
|
|
def results(self) -> Iterator[Self]:
|
|
"""
|
|
Iterate across multiple record sets received by the cursor.
|
|
|
|
Multiple record sets are received after using `executemany()` with
|
|
`!returning=True` or using `execute()` with more than one query in the
|
|
command.
|
|
"""
|
|
if self.pgresult:
|
|
while True:
|
|
yield self
|
|
|
|
if not self.nextset():
|
|
break
|
|
|
|
def set_result(self, index: int) -> Self:
|
|
"""
|
|
Move to a specific result set.
|
|
|
|
:arg index: index of the result to go to
|
|
:type index: `!int`
|
|
|
|
More than one result will be available after executing calling
|
|
`executemany()` or `execute()` with more than one query.
|
|
|
|
`!index` is 0-based and supports negative values, counting from the end,
|
|
the same way you can index items in a list.
|
|
|
|
The function returns self, so that the result may be followed by a
|
|
fetch operation. See `results()` for details.
|
|
"""
|
|
if not -len(self._results) <= index < len(self._results):
|
|
raise IndexError(
|
|
f"index {index} out of range: {len(self._results)} result(s) available"
|
|
)
|
|
if index < 0:
|
|
index = len(self._results) + index
|
|
|
|
self._select_current_result(index)
|
|
return self
|
|
|
|
def fetchone(self) -> Row | None:
|
|
"""
|
|
Return the next record from the current result set.
|
|
|
|
Return `!None` the result set is finished.
|
|
|
|
:rtype: Row | None, with Row defined by `row_factory`
|
|
"""
|
|
self._fetch_pipeline()
|
|
res = self._check_result_for_fetch()
|
|
if self._pos < res.ntuples:
|
|
record = self._tx.load_row(self._pos, self._make_row)
|
|
self._pos += 1
|
|
return record
|
|
return None
|
|
|
|
def fetchmany(self, size: int = 0) -> list[Row]:
|
|
"""
|
|
Return the next `!size` records from the current result set.
|
|
|
|
`!size` default to `!self.arraysize` if not specified.
|
|
|
|
:rtype: Sequence[Row], with Row defined by `row_factory`
|
|
"""
|
|
self._fetch_pipeline()
|
|
res = self._check_result_for_fetch()
|
|
|
|
if not size:
|
|
size = self.arraysize
|
|
records = self._tx.load_rows(
|
|
self._pos, min(self._pos + size, res.ntuples), self._make_row
|
|
)
|
|
self._pos += len(records)
|
|
return records
|
|
|
|
def fetchall(self) -> list[Row]:
|
|
"""
|
|
Return all the remaining records from the current result set.
|
|
|
|
:rtype: Sequence[Row], with Row defined by `row_factory`
|
|
"""
|
|
self._fetch_pipeline()
|
|
res = self._check_result_for_fetch()
|
|
records = self._tx.load_rows(self._pos, res.ntuples, self._make_row)
|
|
self._pos = res.ntuples
|
|
return records
|
|
|
|
def __iter__(self) -> Self:
|
|
return self
|
|
|
|
def __next__(self) -> Row:
|
|
self._fetch_pipeline()
|
|
res = self._check_result_for_fetch()
|
|
if self._pos < res.ntuples:
|
|
record = self._tx.load_row(self._pos, self._make_row)
|
|
self._pos += 1
|
|
return record
|
|
raise StopIteration("no more records to return")
|
|
|
|
def scroll(self, value: int, mode: str = "relative") -> None:
|
|
"""
|
|
Move the cursor in the result set to a new position according to mode.
|
|
|
|
If `!mode` is ``'relative'`` (default), `!value` is taken as offset to
|
|
the current position in the result set; if set to ``'absolute'``,
|
|
`!value` states an absolute target position.
|
|
|
|
Raise `!IndexError` in case a scroll operation would leave the result
|
|
set. In this case the position will not change.
|
|
"""
|
|
self._fetch_pipeline()
|
|
self._scroll(value, mode)
|
|
|
|
@contextmanager
|
|
def copy(
|
|
self,
|
|
statement: Query,
|
|
params: Params | None = None,
|
|
*,
|
|
writer: Writer | None = None,
|
|
) -> Iterator[Copy]:
|
|
"""
|
|
Initiate a :sql:`COPY` operation and return an object to manage it.
|
|
"""
|
|
try:
|
|
with self._conn.lock:
|
|
self._conn.wait(self._start_copy_gen(statement, params))
|
|
|
|
with Copy(self, writer=writer) as copy:
|
|
yield copy
|
|
except e._NO_TRACEBACK as ex:
|
|
raise ex.with_traceback(None)
|
|
|
|
# If a fresher result has been set on the cursor by the Copy object,
|
|
# read its properties (especially rowcount).
|
|
self._select_current_result(0)
|
|
|
|
def _fetch_pipeline(self) -> None:
|
|
if (
|
|
self._execmany_returning is not False
|
|
and (not self.pgresult)
|
|
and self._conn._pipeline
|
|
):
|
|
with self._conn.lock:
|
|
self._conn.wait(self._conn._pipeline._fetch_gen(flush=True))
|