from . import abc
from .misc import ConninfoOption, PGnotify, PGresAttDesc
from .misc import error_message
-from ._enums import ConnStatus, DiagnosticField, ExecStatus, Format
+from ._enums import ConnStatus, DiagnosticField, ExecStatus, Format, Trace
from ._enums import Ping, PipelineStatus, PollingStatus, TransactionStatus
logger = logging.getLogger(__name__)
"Ping",
"DiagnosticField",
"Format",
+ "Trace",
"PGconn",
"PGnotify",
"Conninfo",
# Copyright (C) 2020-2021 The Psycopg Team
-from enum import IntEnum, auto
+from enum import IntEnum, IntFlag, auto
class ConnStatus(IntEnum):
"""Text parameter."""
BINARY = 1
"""Binary parameter."""
+
+
+class Trace(IntFlag):
+ """
+ Enum to control tracing of the client/server communication.
+ """
+
+ __module__ = "psycopg.pq"
+
+ SUPPRESS_TIMESTAMPS = 1
+ """Do not include timestamps in messages."""
+
+ REGRESS_MODE = 2
+ """Redact some fields, e.g. OIDs, from messages."""
pq = ctypes.cdll.LoadLibrary(libname)
+
+class FILE(Structure):
+ pass
+
+
+FILE_ptr = POINTER(FILE)
+
+if sys.platform == "linux":
+ libcname = ctypes.util.find_library("c")
+ assert libcname
+ libc = ctypes.cdll.LoadLibrary(libcname)
+
+ fdopen = libc.fdopen
+ fdopen.argtypes = (c_int, c_char_p)
+ fdopen.restype = FILE_ptr
+
+
# Get the libpq version to define what functions are available.
PQlibVersion = pq.PQlibVersion
PQgetCopyData.restype = c_int
+# 33.10. Control Functions
+
+PQtrace = pq.PQtrace
+PQtrace.argtypes = [PGconn_ptr, FILE_ptr]
+PQtrace.restype = None
+
+_PQsetTraceFlags = None
+
+if libpq_version >= 140000:
+ _PQsetTraceFlags = pq.PQsetTraceFlags
+ _PQsetTraceFlags.argtypes = [PGconn_ptr, c_int]
+ _PQsetTraceFlags.restype = None
+
+
+def PQsetTraceFlags(pgconn: PGconn_struct, flags: int) -> None:
+ if not _PQsetTraceFlags:
+ raise NotSupportedError(
+ f"PQsetTraceFlags requires libpq from PostgreSQL 14,"
+ f" {libpq_version} available instead"
+ )
+
+ _PQsetTraceFlags(pgconn, flags)
+
+
+PQuntrace = pq.PQuntrace
+PQuntrace.argtypes = [PGconn_ptr]
+PQuntrace.restype = None
+
# 33.11. Miscellaneous Functions
PQfreemem = pq.PQfreemem
"LP_c_int",
"LP_c_uint",
"LP_c_ulong",
+ "LP_FILE",
):
return f"pointer[{t.__name__[3:]}]"
from ctypes import Array, pointer
from ctypes import c_char, c_char_p, c_int, c_ubyte, c_uint, c_ulong
+class FILE: ...
+
+def fdopen(fd: int, mode: bytes) -> pointer[FILE]: ... # type: ignore[type-var]
+
Oid = c_uint
class PGconn_struct: ...
arg2: int,
arg3: Array[PGresAttDesc_struct], # type: ignore
) -> int: ...
+def PQtrace(
+ arg1: Optional[PGconn_struct],
+ arg2: pointer[FILE], # type: ignore[type-var]
+) -> None: ...
def PQencryptPasswordConn(
arg1: Optional[PGconn_struct],
arg2: bytes,
def PQgetCancel(arg1: Optional[PGconn_struct]) -> PGcancel_struct: ...
def PQfreeCancel(arg1: Optional[PGcancel_struct]) -> None: ...
def PQputCopyData(arg1: Optional[PGconn_struct], arg2: bytes, arg3: int) -> int: ...
+def PQsetTraceFlags(arg1: Optional[PGconn_struct], arg2: int) -> None: ...
+def PQuntrace(arg1: Optional[PGconn_struct]) -> None: ...
def PQfreemem(arg1: Any) -> None: ...
def _PQencryptPasswordConn(arg1: Optional[PGconn_struct], arg2: bytes, arg3: bytes, arg4: bytes) -> Optional[bytes]: ...
def PQmakeEmptyPGresult(arg1: Optional[PGconn_struct], arg2: int) -> PGresult_struct: ...
from typing import Any, Callable, List, Optional, Sequence, Tuple
from typing import Union, TYPE_CHECKING
-from ._enums import Format
+from ._enums import Format, Trace
from .._compat import Protocol
if TYPE_CHECKING:
def get_copy_data(self, async_: int) -> Tuple[int, memoryview]:
...
+ def trace(self, fileno: int) -> None:
+ ...
+
+ def set_trace_flags(self, flags: Trace) -> None:
+ ...
+
+ def untrace(self) -> None:
+ ...
+
def encrypt_password(
self, passwd: bytes, user: bytes, algorithm: Optional[bytes] = None
) -> bytes:
import os
import logging
+import sys
from weakref import ref
from functools import partial
from . import _pq_ctypes as impl
from .misc import PGnotify, ConninfoOption, PGresAttDesc
from .misc import error_message, connection_summary
-from ._enums import Format, ExecStatus
+from ._enums import Format, ExecStatus, Trace
if TYPE_CHECKING:
from . import abc
else:
return nbytes, memoryview(b"")
+ def trace(self, fileno: int) -> None:
+ if sys.platform != "linux":
+ raise e.NotSupportedError("only supported on Linux")
+ stream = impl.fdopen(fileno, b"w")
+ impl.PQtrace(self._pgconn_ptr, stream)
+
+ def set_trace_flags(self, flags: Trace) -> None:
+ impl.PQsetTraceFlags(self._pgconn_ptr, flags)
+
+ def untrace(self) -> None:
+ impl.PQuntrace(self._pgconn_ptr)
+
def encrypt_password(
self, passwd: bytes, user: bytes, algorithm: Optional[bytes] = None
) -> bytes:
# Copyright (C) 2020-2021 The Psycopg Team
+cdef extern from "stdio.h":
+
+ ctypedef struct FILE:
+ pass
+
cdef extern from "pg_config.h":
int PG_VERSION_NUM
int PQputCopyEnd(PGconn *conn, const char *errormsg)
int PQgetCopyData(PGconn *conn, char **buffer, int async)
+ # 33.10. Control Functions
+ void PQtrace(PGconn *conn, FILE *stream);
+ void PQsetTraceFlags(PGconn *conn, int flags);
+ void PQuntrace(PGconn *conn);
+
# 33.11. Miscellaneous Functions
void PQfreemem(void *ptr) nogil
void PQconninfoFree(PQconninfoOption *connOptions)
"""
pid_t getpid()
+from libc.stdio cimport fdopen
from cpython.mem cimport PyMem_Malloc, PyMem_Free
from cpython.bytes cimport PyBytes_AsString
from cpython.memoryview cimport PyMemoryView_FromObject
import logging
+import sys
-from psycopg.pq import Format as PqFormat
+from psycopg.pq import Format as PqFormat, Trace
from psycopg.pq.misc import PGnotify, connection_summary
from psycopg_c.pq cimport PQBuffer
else:
return nbytes, b"" # won't parse it, doesn't really be memoryview
+ def trace(self, fileno: int) -> None:
+ if sys.platform != "linux":
+ raise e.NotSupportedError("only supported on Linux")
+ stream = fdopen(fileno, b"w")
+ libpq.PQtrace(self._pgconn_ptr, stream)
+
+ def set_trace_flags(self, flags: Trace) -> None:
+ if libpq.PG_VERSION_NUM < 140000:
+ raise e.NotSupportedError(
+ f"PQsetTraceFlags requires libpq from PostgreSQL 14,"
+ f" {libpq.PG_VERSION_NUM} available instead"
+ )
+ libpq.PQsetTraceFlags(self._pgconn_ptr, flags)
+
+ def untrace(self) -> None:
+ libpq.PQuntrace(self._pgconn_ptr)
+
def encrypt_password(
self, const char *passwd, const char *user, algorithm = None
) -> bytes:
assert "hello error" in rec.message
+@pytest.mark.libpq("< 14")
+@pytest.mark.skipif("sys.platform != 'linux'")
+def test_trace_pre14(pgconn, tmp_path):
+ tracef = tmp_path / "trace"
+ with tracef.open("w") as f:
+ pgconn.trace(f.fileno())
+ with pytest.raises(psycopg.NotSupportedError):
+ pgconn.set_trace_flags(0)
+ pgconn.exec_(b"select 1")
+ pgconn.untrace()
+ pgconn.exec_(b"select 2")
+ traces = tracef.read_text()
+ assert "select 1" in traces
+ assert "select 2" not in traces
+
+
+@pytest.mark.libpq(">= 14")
+@pytest.mark.skipif("sys.platform != 'linux'")
+def test_trace(pgconn, tmp_path):
+ tracef = tmp_path / "trace"
+ with tracef.open("w") as f:
+ pgconn.trace(f.fileno())
+ pgconn.set_trace_flags(
+ pq.Trace.SUPPRESS_TIMESTAMPS | pq.Trace.REGRESS_MODE
+ )
+ pgconn.exec_(b"select 1")
+ pgconn.untrace()
+ pgconn.exec_(b"select 2")
+ traces = [line.split("\t") for line in tracef.read_text().splitlines()]
+ assert traces == [
+ ["F", "13", "Query", ' "select 1"'],
+ ["B", "33", "RowDescription", ' 1 "?column?" NNNN 0 NNNN 4 -1 0'],
+ ["B", "11", "DataRow", " 1 1 '1'"],
+ ["B", "13", "CommandComplete", ' "SELECT 1"'],
+ ["B", "5", "ReadyForQuery", " I"],
+ ]
+
+
+@pytest.mark.skipif("sys.platform == 'linux'")
+def test_trace_nonlinux(pgconn):
+ with pytest.raises(psycopg.NotSupportedError):
+ pgconn.trace(1)
+
+
@pytest.mark.libpq(">= 10")
def test_encrypt_password(pgconn):
enc = pgconn.encrypt_password(b"psycopg2", b"ashesh", b"md5")
projects / thirdparty / psycopg.git / commitdiff
