334 lines
9.7 KiB
Python
334 lines
9.7 KiB
Python
|
"""
|
||
|
|
Generators implementing communication protocols with the libpq
|
||
|
|
|
||
|
|
Certain operations (connection, querying) are an interleave of libpq calls and
|
||
|
|
waiting for the socket to be ready. This module contains the code to execute
|
||
|
|
the operations, yielding a polling state whenever there is to wait. The
|
||
|
|
functions in the `waiting` module are the ones who wait more or less
|
||
|
|
cooperatively for the socket to be ready and make these generators continue.
|
||
|
|
|
||
|
|
All these generators yield pairs (fileno, `Wait`) whenever an operation would
|
||
|
|
block. The generator can be restarted sending the appropriate `Ready` state
|
||
|
|
when the file descriptor is ready.
|
||
|
|
|
||
|
|
"""
|
||
|
|
|
||
|
|
# Copyright (C) 2020 The Psycopg Team
|
||
|
|
|
||
|
|
import logging
|
||
|
|
from typing import List, Optional, Union
|
||
|
|
|
||
|
|
from . import pq
|
||
|
|
from . import errors as e
|
||
|
|
from .abc import Buffer, PipelineCommand, PQGen, PQGenConn
|
||
|
|
from .pq.abc import PGconn, PGresult
|
||
|
|
from .waiting import Wait, Ready
|
||
|
|
from ._compat import Deque
|
||
|
|
from ._cmodule import _psycopg
|
||
|
|
from ._encodings import pgconn_encoding, conninfo_encoding
|
||
|
|
|
||
|
|
OK = pq.ConnStatus.OK
|
||
|
|
BAD = pq.ConnStatus.BAD
|
||
|
|
|
||
|
|
POLL_OK = pq.PollingStatus.OK
|
||
|
|
POLL_READING = pq.PollingStatus.READING
|
||
|
|
POLL_WRITING = pq.PollingStatus.WRITING
|
||
|
|
POLL_FAILED = pq.PollingStatus.FAILED
|
||
|
|
|
||
|
|
COMMAND_OK = pq.ExecStatus.COMMAND_OK
|
||
|
|
COPY_OUT = pq.ExecStatus.COPY_OUT
|
||
|
|
COPY_IN = pq.ExecStatus.COPY_IN
|
||
|
|
COPY_BOTH = pq.ExecStatus.COPY_BOTH
|
||
|
|
PIPELINE_SYNC = pq.ExecStatus.PIPELINE_SYNC
|
||
|
|
|
||
|
|
WAIT_R = Wait.R
|
||
|
|
WAIT_W = Wait.W
|
||
|
|
WAIT_RW = Wait.RW
|
||
|
|
READY_R = Ready.R
|
||
|
|
READY_W = Ready.W
|
||
|
|
READY_RW = Ready.RW
|
||
|
|
|
||
|
|
logger = logging.getLogger(__name__)
|
||
|
|
|
||
|
|
|
||
|
|
def _connect(conninfo: str) -> PQGenConn[PGconn]:
|
||
|
|
"""
|
||
|
|
Generator to create a database connection without blocking.
|
||
|
|
|
||
|
|
"""
|
||
|
|
conn = pq.PGconn.connect_start(conninfo.encode())
|
||
|
|
while True:
|
||
|
|
if conn.status == BAD:
|
||
|
|
encoding = conninfo_encoding(conninfo)
|
||
|
|
raise e.OperationalError(
|
||
|
|
f"connection is bad: {pq.error_message(conn, encoding=encoding)}",
|
||
|
|
pgconn=conn,
|
||
|
|
)
|
||
|
|
|
||
|
|
status = conn.connect_poll()
|
||
|
|
if status == POLL_OK:
|
||
|
|
break
|
||
|
|
elif status == POLL_READING:
|
||
|
|
yield conn.socket, WAIT_R
|
||
|
|
elif status == POLL_WRITING:
|
||
|
|
yield conn.socket, WAIT_W
|
||
|
|
elif status == POLL_FAILED:
|
||
|
|
encoding = conninfo_encoding(conninfo)
|
||
|
|
raise e.OperationalError(
|
||
|
|
f"connection failed: {pq.error_message(conn, encoding=encoding)}",
|
||
|
|
pgconn=e.finish_pgconn(conn),
|
||
|
|
)
|
||
|
|
else:
|
||
|
|
raise e.InternalError(
|
||
|
|
f"unexpected poll status: {status}", pgconn=e.finish_pgconn(conn)
|
||
|
|
)
|
||
|
|
|
||
|
|
conn.nonblocking = 1
|
||
|
|
return conn
|
||
|
|
|
||
|
|
|
||
|
|
def _execute(pgconn: PGconn) -> PQGen[List[PGresult]]:
|
||
|
|
"""
|
||
|
|
Generator sending a query and returning results without blocking.
|
||
|
|
|
||
|
|
The query must have already been sent using `pgconn.send_query()` or
|
||
|
|
similar. Flush the query and then return the result using nonblocking
|
||
|
|
functions.
|
||
|
|
|
||
|
|
Return the list of results returned by the database (whether success
|
||
|
|
or error).
|
||
|
|
"""
|
||
|
|
yield from _send(pgconn)
|
||
|
|
rv = yield from _fetch_many(pgconn)
|
||
|
|
return rv
|
||
|
|
|
||
|
|
|
||
|
|
def _send(pgconn: PGconn) -> PQGen[None]:
|
||
|
|
"""
|
||
|
|
Generator to send a query to the server without blocking.
|
||
|
|
|
||
|
|
The query must have already been sent using `pgconn.send_query()` or
|
||
|
|
similar. Flush the query and then return the result using nonblocking
|
||
|
|
functions.
|
||
|
|
|
||
|
|
After this generator has finished you may want to cycle using `fetch()`
|
||
|
|
to retrieve the results available.
|
||
|
|
"""
|
||
|
|
while True:
|
||
|
|
f = pgconn.flush()
|
||
|
|
if f == 0:
|
||
|
|
break
|
||
|
|
|
||
|
|
ready = yield WAIT_RW
|
||
|
|
if ready & READY_R:
|
||
|
|
# This call may read notifies: they will be saved in the
|
||
|
|
# PGconn buffer and passed to Python later, in `fetch()`.
|
||
|
|
pgconn.consume_input()
|
||
|
|
|
||
|
|
|
||
|
|
def _fetch_many(pgconn: PGconn) -> PQGen[List[PGresult]]:
|
||
|
|
"""
|
||
|
|
Generator retrieving results from the database without blocking.
|
||
|
|
|
||
|
|
The query must have already been sent to the server, so pgconn.flush() has
|
||
|
|
already returned 0.
|
||
|
|
|
||
|
|
Return the list of results returned by the database (whether success
|
||
|
|
or error).
|
||
|
|
"""
|
||
|
|
results: List[PGresult] = []
|
||
|
|
while True:
|
||
|
|
res = yield from _fetch(pgconn)
|
||
|
|
if not res:
|
||
|
|
break
|
||
|
|
|
||
|
|
results.append(res)
|
||
|
|
status = res.status
|
||
|
|
if status == COPY_IN or status == COPY_OUT or status == COPY_BOTH:
|
||
|
|
# After entering copy mode the libpq will create a phony result
|
||
|
|
# for every request so let's break the endless loop.
|
||
|
|
break
|
||
|
|
|
||
|
|
if status == PIPELINE_SYNC:
|
||
|
|
# PIPELINE_SYNC is not followed by a NULL, but we return it alone
|
||
|
|
# similarly to other result sets.
|
||
|
|
assert len(results) == 1, results
|
||
|
|
break
|
||
|
|
|
||
|
|
return results
|
||
|
|
|
||
|
|
|
||
|
|
def _fetch(pgconn: PGconn) -> PQGen[Optional[PGresult]]:
|
||
|
|
"""
|
||
|
|
Generator retrieving a single result from the database without blocking.
|
||
|
|
|
||
|
|
The query must have already been sent to the server, so pgconn.flush() has
|
||
|
|
already returned 0.
|
||
|
|
|
||
|
|
Return a result from the database (whether success or error).
|
||
|
|
"""
|
||
|
|
if pgconn.is_busy():
|
||
|
|
yield WAIT_R
|
||
|
|
while True:
|
||
|
|
pgconn.consume_input()
|
||
|
|
if not pgconn.is_busy():
|
||
|
|
break
|
||
|
|
yield WAIT_R
|
||
|
|
|
||
|
|
_consume_notifies(pgconn)
|
||
|
|
|
||
|
|
return pgconn.get_result()
|
||
|
|
|
||
|
|
|
||
|
|
def _pipeline_communicate(
|
||
|
|
pgconn: PGconn, commands: Deque[PipelineCommand]
|
||
|
|
) -> PQGen[List[List[PGresult]]]:
|
||
|
|
"""Generator to send queries from a connection in pipeline mode while also
|
||
|
|
receiving results.
|
||
|
|
|
||
|
|
Return a list results, including single PIPELINE_SYNC elements.
|
||
|
|
"""
|
||
|
|
results = []
|
||
|
|
|
||
|
|
while True:
|
||
|
|
ready = yield WAIT_RW
|
||
|
|
|
||
|
|
if ready & READY_R:
|
||
|
|
pgconn.consume_input()
|
||
|
|
_consume_notifies(pgconn)
|
||
|
|
|
||
|
|
res: List[PGresult] = []
|
||
|
|
while not pgconn.is_busy():
|
||
|
|
r = pgconn.get_result()
|
||
|
|
if r is None:
|
||
|
|
if not res:
|
||
|
|
break
|
||
|
|
results.append(res)
|
||
|
|
res = []
|
||
|
|
else:
|
||
|
|
status = r.status
|
||
|
|
if status == PIPELINE_SYNC:
|
||
|
|
assert not res
|
||
|
|
results.append([r])
|
||
|
|
elif status == COPY_IN or status == COPY_OUT or status == COPY_BOTH:
|
||
|
|
# This shouldn't happen, but insisting hard enough, it will.
|
||
|
|
# For instance, in test_executemany_badquery(), with the COPY
|
||
|
|
# statement and the AsyncClientCursor, which disables
|
||
|
|
# prepared statements).
|
||
|
|
# Bail out from the resulting infinite loop.
|
||
|
|
raise e.NotSupportedError(
|
||
|
|
"COPY cannot be used in pipeline mode"
|
||
|
|
)
|
||
|
|
else:
|
||
|
|
res.append(r)
|
||
|
|
|
||
|
|
if ready & READY_W:
|
||
|
|
pgconn.flush()
|
||
|
|
if not commands:
|
||
|
|
break
|
||
|
|
commands.popleft()()
|
||
|
|
|
||
|
|
return results
|
||
|
|
|
||
|
|
|
||
|
|
def _consume_notifies(pgconn: PGconn) -> None:
|
||
|
|
# Consume notifies
|
||
|
|
while True:
|
||
|
|
n = pgconn.notifies()
|
||
|
|
if not n:
|
||
|
|
break
|
||
|
|
if pgconn.notify_handler:
|
||
|
|
pgconn.notify_handler(n)
|
||
|
|
|
||
|
|
|
||
|
|
def notifies(pgconn: PGconn) -> PQGen[List[pq.PGnotify]]:
|
||
|
|
yield WAIT_R
|
||
|
|
pgconn.consume_input()
|
||
|
|
|
||
|
|
ns = []
|
||
|
|
while True:
|
||
|
|
n = pgconn.notifies()
|
||
|
|
if n:
|
||
|
|
ns.append(n)
|
||
|
|
else:
|
||
|
|
break
|
||
|
|
|
||
|
|
return ns
|
||
|
|
|
||
|
|
|
||
|
|
def copy_from(pgconn: PGconn) -> PQGen[Union[memoryview, PGresult]]:
|
||
|
|
while True:
|
||
|
|
nbytes, data = pgconn.get_copy_data(1)
|
||
|
|
if nbytes != 0:
|
||
|
|
break
|
||
|
|
|
||
|
|
# would block
|
||
|
|
yield WAIT_R
|
||
|
|
pgconn.consume_input()
|
||
|
|
|
||
|
|
if nbytes > 0:
|
||
|
|
# some data
|
||
|
|
return data
|
||
|
|
|
||
|
|
# Retrieve the final result of copy
|
||
|
|
results = yield from _fetch_many(pgconn)
|
||
|
|
if len(results) > 1:
|
||
|
|
# TODO: too brutal? Copy worked.
|
||
|
|
raise e.ProgrammingError("you cannot mix COPY with other operations")
|
||
|
|
result = results[0]
|
||
|
|
if result.status != COMMAND_OK:
|
||
|
|
encoding = pgconn_encoding(pgconn)
|
||
|
|
raise e.error_from_result(result, encoding=encoding)
|
||
|
|
|
||
|
|
return result
|
||
|
|
|
||
|
|
|
||
|
|
def copy_to(pgconn: PGconn, buffer: Buffer) -> PQGen[None]:
|
||
|
|
# Retry enqueuing data until successful.
|
||
|
|
#
|
||
|
|
# WARNING! This can cause an infinite loop if the buffer is too large. (see
|
||
|
|
# ticket #255). We avoid it in the Copy object by splitting a large buffer
|
||
|
|
# into smaller ones. We prefer to do it there instead of here in order to
|
||
|
|
# do it upstream the queue decoupling the writer task from the producer one.
|
||
|
|
while pgconn.put_copy_data(buffer) == 0:
|
||
|
|
yield WAIT_W
|
||
|
|
|
||
|
|
|
||
|
|
def copy_end(pgconn: PGconn, error: Optional[bytes]) -> PQGen[PGresult]:
|
||
|
|
# Retry enqueuing end copy message until successful
|
||
|
|
while pgconn.put_copy_end(error) == 0:
|
||
|
|
yield WAIT_W
|
||
|
|
|
||
|
|
# Repeat until it the message is flushed to the server
|
||
|
|
while True:
|
||
|
|
yield WAIT_W
|
||
|
|
f = pgconn.flush()
|
||
|
|
if f == 0:
|
||
|
|
break
|
||
|
|
|
||
|
|
# Retrieve the final result of copy
|
||
|
|
(result,) = yield from _fetch_many(pgconn)
|
||
|
|
if result.status != COMMAND_OK:
|
||
|
|
encoding = pgconn_encoding(pgconn)
|
||
|
|
raise e.error_from_result(result, encoding=encoding)
|
||
|
|
|
||
|
|
return result
|
||
|
|
|
||
|
|
|
||
|
|
# Override functions with fast versions if available
|
||
|
|
if _psycopg:
|
||
|
|
connect = _psycopg.connect
|
||
|
|
execute = _psycopg.execute
|
||
|
|
send = _psycopg.send
|
||
|
|
fetch_many = _psycopg.fetch_many
|
||
|
|
fetch = _psycopg.fetch
|
||
|
|
pipeline_communicate = _psycopg.pipeline_communicate
|
||
|
|
|
||
|
|
else:
|
||
|
|
connect = _connect
|
||
|
|
execute = _execute
|
||
|
|
send = _send
|
||
|
|
fetch_many = _fetch_many
|
||
|
|
fetch = _fetch
|
||
|
|
pipeline_communicate = _pipeline_communicate
|