299 lines
10 KiB
Python
299 lines
10 KiB
Python
|
"""
|
||
|
|
commands pipeline management
|
||
|
|
"""
|
||
|
|
|
||
|
|
# Copyright (C) 2021 The Psycopg Team
|
||
|
|
|
||
|
|
import logging
|
||
|
|
from types import TracebackType
|
||
|
|
from typing import Any, List, Optional, Union, Tuple, Type, TypeVar, TYPE_CHECKING
|
||
|
|
from typing_extensions import TypeAlias
|
||
|
|
|
||
|
|
from . import pq
|
||
|
|
from . import errors as e
|
||
|
|
from .abc import PipelineCommand, PQGen
|
||
|
|
from ._compat import Deque
|
||
|
|
from .pq.misc import connection_summary
|
||
|
|
from ._encodings import pgconn_encoding
|
||
|
|
from ._preparing import Key, Prepare
|
||
|
|
from .generators import pipeline_communicate, fetch_many, send
|
||
|
|
|
||
|
|
if TYPE_CHECKING:
|
||
|
|
from .pq.abc import PGresult
|
||
|
|
from .cursor import BaseCursor
|
||
|
|
from .connection import BaseConnection, Connection
|
||
|
|
from .connection_async import AsyncConnection
|
||
|
|
|
||
|
|
|
||
|
|
PendingResult: TypeAlias = Union[
|
||
|
|
None, Tuple["BaseCursor[Any, Any]", Optional[Tuple[Key, Prepare, bytes]]]
|
||
|
|
]
|
||
|
|
|
||
|
|
FATAL_ERROR = pq.ExecStatus.FATAL_ERROR
|
||
|
|
PIPELINE_ABORTED = pq.ExecStatus.PIPELINE_ABORTED
|
||
|
|
BAD = pq.ConnStatus.BAD
|
||
|
|
|
||
|
|
ACTIVE = pq.TransactionStatus.ACTIVE
|
||
|
|
|
||
|
|
logger = logging.getLogger("psycopg")
|
||
|
|
|
||
|
|
|
||
|
|
class BasePipeline:
|
||
|
|
command_queue: Deque[PipelineCommand]
|
||
|
|
result_queue: Deque[PendingResult]
|
||
|
|
_is_supported: Optional[bool] = None
|
||
|
|
|
||
|
|
def __init__(self, conn: "BaseConnection[Any]") -> None:
|
||
|
|
self._conn = conn
|
||
|
|
self.pgconn = conn.pgconn
|
||
|
|
self.command_queue = Deque[PipelineCommand]()
|
||
|
|
self.result_queue = Deque[PendingResult]()
|
||
|
|
self.level = 0
|
||
|
|
|
||
|
|
def __repr__(self) -> str:
|
||
|
|
cls = f"{self.__class__.__module__}.{self.__class__.__qualname__}"
|
||
|
|
info = connection_summary(self._conn.pgconn)
|
||
|
|
return f"<{cls} {info} at 0x{id(self):x}>"
|
||
|
|
|
||
|
|
@property
|
||
|
|
def status(self) -> pq.PipelineStatus:
|
||
|
|
return pq.PipelineStatus(self.pgconn.pipeline_status)
|
||
|
|
|
||
|
|
@classmethod
|
||
|
|
def is_supported(cls) -> bool:
|
||
|
|
"""Return `!True` if the psycopg libpq wrapper supports pipeline mode."""
|
||
|
|
if BasePipeline._is_supported is None:
|
||
|
|
BasePipeline._is_supported = not cls._not_supported_reason()
|
||
|
|
return BasePipeline._is_supported
|
||
|
|
|
||
|
|
@classmethod
|
||
|
|
def _not_supported_reason(cls) -> str:
|
||
|
|
"""Return the reason why the pipeline mode is not supported.
|
||
|
|
|
||
|
|
Return an empty string if pipeline mode is supported.
|
||
|
|
"""
|
||
|
|
# Support only depends on the libpq functions available in the pq
|
||
|
|
# wrapper, not on the database version.
|
||
|
|
if pq.version() < 140000:
|
||
|
|
return (
|
||
|
|
f"libpq too old {pq.version()};"
|
||
|
|
" v14 or greater required for pipeline mode"
|
||
|
|
)
|
||
|
|
|
||
|
|
if pq.__build_version__ < 140000:
|
||
|
|
return (
|
||
|
|
f"libpq too old: module built for {pq.__build_version__};"
|
||
|
|
" v14 or greater required for pipeline mode"
|
||
|
|
)
|
||
|
|
|
||
|
|
return ""
|
||
|
|
|
||
|
|
def _enter_gen(self) -> PQGen[None]:
|
||
|
|
if not self.is_supported():
|
||
|
|
raise e.NotSupportedError(
|
||
|
|
f"pipeline mode not supported: {self._not_supported_reason()}"
|
||
|
|
)
|
||
|
|
if self.level == 0:
|
||
|
|
self.pgconn.enter_pipeline_mode()
|
||
|
|
elif self.command_queue or self.pgconn.transaction_status == ACTIVE:
|
||
|
|
# Nested pipeline case.
|
||
|
|
# Transaction might be ACTIVE when the pipeline uses an "implicit
|
||
|
|
# transaction", typically in autocommit mode. But when entering a
|
||
|
|
# Psycopg transaction(), we expect the IDLE state. By sync()-ing,
|
||
|
|
# we make sure all previous commands are completed and the
|
||
|
|
# transaction gets back to IDLE.
|
||
|
|
yield from self._sync_gen()
|
||
|
|
self.level += 1
|
||
|
|
|
||
|
|
def _exit(self, exc: Optional[BaseException]) -> None:
|
||
|
|
self.level -= 1
|
||
|
|
if self.level == 0 and self.pgconn.status != BAD:
|
||
|
|
try:
|
||
|
|
self.pgconn.exit_pipeline_mode()
|
||
|
|
except e.OperationalError as exc2:
|
||
|
|
# Notice that this error might be pretty irrecoverable. It
|
||
|
|
# happens on COPY, for instance: even if sync succeeds, exiting
|
||
|
|
# fails with "cannot exit pipeline mode with uncollected results"
|
||
|
|
if exc:
|
||
|
|
logger.warning("error ignored exiting %r: %s", self, exc2)
|
||
|
|
else:
|
||
|
|
raise exc2.with_traceback(None)
|
||
|
|
|
||
|
|
def _sync_gen(self) -> PQGen[None]:
|
||
|
|
self._enqueue_sync()
|
||
|
|
yield from self._communicate_gen()
|
||
|
|
yield from self._fetch_gen(flush=False)
|
||
|
|
|
||
|
|
def _exit_gen(self) -> PQGen[None]:
|
||
|
|
"""
|
||
|
|
Exit current pipeline by sending a Sync and fetch back all remaining results.
|
||
|
|
"""
|
||
|
|
try:
|
||
|
|
self._enqueue_sync()
|
||
|
|
yield from self._communicate_gen()
|
||
|
|
finally:
|
||
|
|
# No need to force flush since we emitted a sync just before.
|
||
|
|
yield from self._fetch_gen(flush=False)
|
||
|
|
|
||
|
|
def _communicate_gen(self) -> PQGen[None]:
|
||
|
|
"""Communicate with pipeline to send commands and possibly fetch
|
||
|
|
results, which are then processed.
|
||
|
|
"""
|
||
|
|
fetched = yield from pipeline_communicate(self.pgconn, self.command_queue)
|
||
|
|
exception = None
|
||
|
|
for results in fetched:
|
||
|
|
queued = self.result_queue.popleft()
|
||
|
|
try:
|
||
|
|
self._process_results(queued, results)
|
||
|
|
except e.Error as exc:
|
||
|
|
if exception is None:
|
||
|
|
exception = exc
|
||
|
|
if exception is not None:
|
||
|
|
raise exception
|
||
|
|
|
||
|
|
def _fetch_gen(self, *, flush: bool) -> PQGen[None]:
|
||
|
|
"""Fetch available results from the connection and process them with
|
||
|
|
pipeline queued items.
|
||
|
|
|
||
|
|
If 'flush' is True, a PQsendFlushRequest() is issued in order to make
|
||
|
|
sure results can be fetched. Otherwise, the caller may emit a
|
||
|
|
PQpipelineSync() call to ensure the output buffer gets flushed before
|
||
|
|
fetching.
|
||
|
|
"""
|
||
|
|
if not self.result_queue:
|
||
|
|
return
|
||
|
|
|
||
|
|
if flush:
|
||
|
|
self.pgconn.send_flush_request()
|
||
|
|
yield from send(self.pgconn)
|
||
|
|
|
||
|
|
exception = None
|
||
|
|
while self.result_queue:
|
||
|
|
results = yield from fetch_many(self.pgconn)
|
||
|
|
if not results:
|
||
|
|
# No more results to fetch, but there may still be pending
|
||
|
|
# commands.
|
||
|
|
break
|
||
|
|
queued = self.result_queue.popleft()
|
||
|
|
try:
|
||
|
|
self._process_results(queued, results)
|
||
|
|
except e.Error as exc:
|
||
|
|
if exception is None:
|
||
|
|
exception = exc
|
||
|
|
if exception is not None:
|
||
|
|
raise exception
|
||
|
|
|
||
|
|
def _process_results(
|
||
|
|
self, queued: PendingResult, results: List["PGresult"]
|
||
|
|
) -> None:
|
||
|
|
"""Process a results set fetched from the current pipeline.
|
||
|
|
|
||
|
|
This matches 'results' with its respective element in the pipeline
|
||
|
|
queue. For commands (None value in the pipeline queue), results are
|
||
|
|
checked directly. For prepare statement creation requests, update the
|
||
|
|
cache. Otherwise, results are attached to their respective cursor.
|
||
|
|
"""
|
||
|
|
if queued is None:
|
||
|
|
(result,) = results
|
||
|
|
if result.status == FATAL_ERROR:
|
||
|
|
raise e.error_from_result(result, encoding=pgconn_encoding(self.pgconn))
|
||
|
|
elif result.status == PIPELINE_ABORTED:
|
||
|
|
raise e.PipelineAborted("pipeline aborted")
|
||
|
|
else:
|
||
|
|
cursor, prepinfo = queued
|
||
|
|
if prepinfo:
|
||
|
|
key, prep, name = prepinfo
|
||
|
|
# Update the prepare state of the query.
|
||
|
|
cursor._conn._prepared.validate(key, prep, name, results)
|
||
|
|
cursor._set_results_from_pipeline(results)
|
||
|
|
|
||
|
|
def _enqueue_sync(self) -> None:
|
||
|
|
"""Enqueue a PQpipelineSync() command."""
|
||
|
|
self.command_queue.append(self.pgconn.pipeline_sync)
|
||
|
|
self.result_queue.append(None)
|
||
|
|
|
||
|
|
|
||
|
|
class Pipeline(BasePipeline):
|
||
|
|
"""Handler for connection in pipeline mode."""
|
||
|
|
|
||
|
|
__module__ = "psycopg"
|
||
|
|
_conn: "Connection[Any]"
|
||
|
|
_Self = TypeVar("_Self", bound="Pipeline")
|
||
|
|
|
||
|
|
def __init__(self, conn: "Connection[Any]") -> None:
|
||
|
|
super().__init__(conn)
|
||
|
|
|
||
|
|
def sync(self) -> None:
|
||
|
|
"""Sync the pipeline, send any pending command and receive and process
|
||
|
|
all available results.
|
||
|
|
"""
|
||
|
|
try:
|
||
|
|
with self._conn.lock:
|
||
|
|
self._conn.wait(self._sync_gen())
|
||
|
|
except e._NO_TRACEBACK as ex:
|
||
|
|
raise ex.with_traceback(None)
|
||
|
|
|
||
|
|
def __enter__(self: _Self) -> _Self:
|
||
|
|
with self._conn.lock:
|
||
|
|
self._conn.wait(self._enter_gen())
|
||
|
|
return self
|
||
|
|
|
||
|
|
def __exit__(
|
||
|
|
self,
|
||
|
|
exc_type: Optional[Type[BaseException]],
|
||
|
|
exc_val: Optional[BaseException],
|
||
|
|
exc_tb: Optional[TracebackType],
|
||
|
|
) -> None:
|
||
|
|
try:
|
||
|
|
with self._conn.lock:
|
||
|
|
self._conn.wait(self._exit_gen())
|
||
|
|
except Exception as exc2:
|
||
|
|
# Don't clobber an exception raised in the block with this one
|
||
|
|
if exc_val:
|
||
|
|
logger.warning("error ignored terminating %r: %s", self, exc2)
|
||
|
|
else:
|
||
|
|
raise exc2.with_traceback(None)
|
||
|
|
finally:
|
||
|
|
self._exit(exc_val)
|
||
|
|
|
||
|
|
|
||
|
|
class AsyncPipeline(BasePipeline):
|
||
|
|
"""Handler for async connection in pipeline mode."""
|
||
|
|
|
||
|
|
__module__ = "psycopg"
|
||
|
|
_conn: "AsyncConnection[Any]"
|
||
|
|
_Self = TypeVar("_Self", bound="AsyncPipeline")
|
||
|
|
|
||
|
|
def __init__(self, conn: "AsyncConnection[Any]") -> None:
|
||
|
|
super().__init__(conn)
|
||
|
|
|
||
|
|
async def sync(self) -> None:
|
||
|
|
try:
|
||
|
|
async with self._conn.lock:
|
||
|
|
await self._conn.wait(self._sync_gen())
|
||
|
|
except e._NO_TRACEBACK as ex:
|
||
|
|
raise ex.with_traceback(None)
|
||
|
|
|
||
|
|
async def __aenter__(self: _Self) -> _Self:
|
||
|
|
async with self._conn.lock:
|
||
|
|
await self._conn.wait(self._enter_gen())
|
||
|
|
return self
|
||
|
|
|
||
|
|
async def __aexit__(
|
||
|
|
self,
|
||
|
|
exc_type: Optional[Type[BaseException]],
|
||
|
|
exc_val: Optional[BaseException],
|
||
|
|
exc_tb: Optional[TracebackType],
|
||
|
|
) -> None:
|
||
|
|
try:
|
||
|
|
async with self._conn.lock:
|
||
|
|
await self._conn.wait(self._exit_gen())
|
||
|
|
except Exception as exc2:
|
||
|
|
# Don't clobber an exception raised in the block with this one
|
||
|
|
if exc_val:
|
||
|
|
logger.warning("error ignored terminating %r: %s", self, exc2)
|
||
|
|
else:
|
||
|
|
raise exc2.with_traceback(None)
|
||
|
|
finally:
|
||
|
|
self._exit(exc_val)
|