]> git.ipfire.org Git - thirdparty/psycopg.git/commitdiff
Added C implementation of int dumpers (text, binary)
authorDaniele Varrazzo <daniele.varrazzo@gmail.com>
Sat, 21 Nov 2020 20:17:26 +0000 (20:17 +0000)
committerDaniele Varrazzo <daniele.varrazzo@gmail.com>
Sat, 21 Nov 2020 20:34:29 +0000 (20:34 +0000)
psycopg3/psycopg3/types/numeric.py
psycopg3_c/psycopg3_c/adapt.pyx
psycopg3_c/psycopg3_c/pq_cython.pxd
psycopg3_c/psycopg3_c/pq_cython.pyx
psycopg3_c/psycopg3_c/types/numeric.pyx

index 8281199e7a0dd98a1ad617f00e9cb53f2cdd28b5..2a70a8e6f2f1914eaa0b52d0e3952f9bd76806d0 100644 (file)
@@ -57,11 +57,17 @@ class Oid(int):
 
 
 class NumberDumper(Dumper):
-    _special: Dict[bytes, bytes] = {}
-
     def dump(self, obj: Any) -> bytes:
         return str(obj).encode("utf8")
 
+    def quote(self, obj: Any) -> bytes:
+        value = self.dump(obj)
+        return value if obj >= 0 else b" " + value
+
+
+class SpecialValuesDumper(NumberDumper):
+    _special: Dict[bytes, bytes] = {}
+
     def quote(self, obj: Any) -> bytes:
         value = self.dump(obj)
 
@@ -83,7 +89,7 @@ class IntBinaryDumper(IntDumper):
 
 
 @Dumper.text(float)
-class FloatDumper(NumberDumper):
+class FloatDumper(SpecialValuesDumper):
     oid = builtins["float8"].oid
 
     _special = {
@@ -100,7 +106,7 @@ class FloatBinaryDumper(NumberDumper):
 
 
 @Dumper.text(Decimal)
-class DecimalDumper(NumberDumper):
+class DecimalDumper(SpecialValuesDumper):
     oid = builtins["numeric"].oid
 
     _special = {
index c099bd5eee9c27cf7ba61a6e63fb9cf9b268c22b..fd1a336e397d56583bc51ccbd47a8e73a957a981 100644 (file)
@@ -17,8 +17,9 @@ from typing import Any
 
 from cpython.bytes cimport PyBytes_AsStringAndSize
 
-from psycopg3_c.adapt cimport cloader_func, get_context_func
 from psycopg3_c cimport libpq as impl
+from psycopg3_c.adapt cimport cloader_func, get_context_func
+from psycopg3_c.pq_cython cimport Escaping
 
 from psycopg3.pq.enums import Format
 
@@ -26,14 +27,67 @@ import logging
 logger = logging.getLogger("psycopg3.adapt")
 
 
+cdef class CDumper:
+    cdef object src
+    cdef object context
+    cdef object connection
+
+    def __init__(self, src: type, context: AdaptContext = None):
+        self.src = src
+        self.context = context
+        self.connection = _connection_from_context(context)
+
+    def dump(self, obj: Any) -> bytes:
+        raise NotImplementedError()
+
+    def quote(self, obj: Any) -> bytes:
+        # TODO: can be optimized
+        cdef bytes value = self.dump(obj)
+        cdef bytes tmp
+        cdef Escaping esc
+
+        if self.connection:
+            esc = Escaping(self.connection.pgconn)
+            return esc.escape_literal(value)
+
+        else:
+            esc = Escaping()
+            tmp = esc.escape_string(value)
+            return b"'%s'" % tmp
+
+    @property
+    def oid(self) -> int:
+        return 0
+
+    @classmethod
+    def register(
+        cls,
+        src: Union[type, str],
+        context: AdaptContext = None,
+        format: Format = Format.TEXT,
+    ) -> None:
+        if not isinstance(src, (str, type)):
+            raise TypeError(
+                f"dumpers should be registered on classes, got {src} instead"
+            )
+        from psycopg3.adapt import Dumper
+
+        where = context.dumpers if context else Dumper.globals
+        where[src, format] = cls
+
+    @classmethod
+    def register_binary(
+        cls, src: Union[type, str], context: AdaptContext = None
+    ) -> None:
+        cls.register(src, context, format=Format.BINARY)
+
+
 cdef class CLoader:
     cdef impl.Oid oid
     cdef object context
     cdef object connection
 
     def __init__(self, oid: int, context: "AdaptContext" = None):
-        from psycopg3.adapt import _connection_from_context
-
         self.oid = oid
         self.context = context
         self.connection = _connection_from_context(context)
@@ -71,6 +125,11 @@ cdef class CLoader:
         cls.register(oid, context, format=Format.BINARY)
 
 
+cdef _connection_from_context(object context):
+    from psycopg3.adapt import _connection_from_context
+    return _connection_from_context(context)
+
+
 def register_builtin_c_adapters():
     """
     Register all the builtin optimized adpaters.
index edfe82c8c6caa8fee30cd4e10af383882d187df5..8b94ea0d8930ccd04db1283b13d36dfd282654e0 100644 (file)
@@ -32,3 +32,7 @@ cdef class PGcancel:
 
     @staticmethod
     cdef PGcancel _from_ptr(impl.PGcancel *ptr)
+
+
+cdef class Escaping:
+    cdef PGconn conn
index c310f5a0435ea7b7e8f50fb4adfeaa3182e1ec7a..1b1c8530dfa29444683757fc6e820d3268bcb30b 100644 (file)
@@ -790,8 +790,6 @@ class Conninfo:
 
 
 cdef class Escaping:
-    cdef PGconn conn
-
     def __init__(self, conn: Optional[PGconn] = None):
         self.conn = conn
 
index f9161fc0427e7ce083177a1cd5ab0c739e385a73..17e362beeb026fb0e77181b4fa3c914e961728d1 100644 (file)
@@ -5,17 +5,49 @@ Cython adapters for numeric types.
 # Copyright (C) 2020 The Psycopg Team
 
 from libc.stdint cimport *
-from psycopg3_c.endian cimport be16toh, be32toh, be64toh
+from psycopg3_c.endian cimport be16toh, be32toh, be64toh, htobe64
 
-from cpython.long cimport PyLong_FromString, PyLong_FromLong
+from cpython.long cimport PyLong_FromString, PyLong_FromLong, PyLong_AsLongLong
 from cpython.long cimport PyLong_FromLongLong, PyLong_FromUnsignedLong
 from cpython.float cimport PyFloat_FromDouble
 
-# work around https://github.com/cython/cython/issues/3909
 cdef extern from "Python.h":
+    # work around https://github.com/cython/cython/issues/3909
     double PyOS_string_to_double(
         const char *s, char **endptr, object overflow_exception) except? -1.0
 
+    int PyOS_snprintf(char *str, size_t size, const char *format, ...)
+
+
+cdef class IntDumper(CDumper):
+    oid = 20  # TODO: int8 oid
+
+    def dump(self, obj: Any) -> bytes:
+        cdef char buf[22]
+        cdef long long val = PyLong_AsLongLong(obj)
+        cdef int written = PyOS_snprintf(buf, sizeof(buf), "%lld", val)
+        return buf[:written]
+
+    def quote(self, obj: Any) -> bytes:
+        cdef char buf[23]
+        cdef long long val = PyLong_AsLongLong(obj)
+        cdef int written
+        if val >= 0:
+            written = PyOS_snprintf(buf, sizeof(buf), "%lld", val)
+        else:
+            written = PyOS_snprintf(buf, sizeof(buf), " %lld", val)
+
+        return buf[:written]
+
+
+cdef class IntBinaryDumper(IntDumper):
+    def dump(self, obj: Any) -> bytes:
+        cdef long long val = PyLong_AsLongLong(obj)
+        cdef uint64_t *ptvar = <uint64_t *>(&val)
+        cdef int64_t beval = htobe64(ptvar[0])
+        cdef char *buf = <char *>&beval
+        return buf[:sizeof(beval)]
+
 
 cdef class IntLoader(CLoader):
     cdef object cload(self, const char *data, size_t length):
@@ -79,7 +111,9 @@ cdef void register_numeric_c_adapters():
     logger.debug("registering optimised numeric c adapters")
 
     from psycopg3.oids import builtins
-    from psycopg3.adapt import Loader
+
+    IntDumper.register(int)
+    IntBinaryDumper.register_binary(int)
 
     IntLoader.register(builtins["int2"].oid)
     IntLoader.register(builtins["int4"].oid)