Added support for datetime extensions

CHANGELOG.md

@@ -4,7 +4,7 @@
 * Drop support for `loop` argument in the `Connection` (fixes #18)
 
 **New features:**
-* Added support for `Decimal` and `UUID` types natively
+* Added support for `Decimal`, `UUID` and `datetime` types natively using MessagePack extensions
 * Added support for SQL prepared statements with `Connection.prepare()` method and
   `PreparedStatement` class
 * Added support for interactive transactions and streams (fixes #21)

asynctnt/iproto/bit.pxd

@@ -0,0 +1,35 @@
+from libc.stdint cimport uint64_t, uint32_t, uint16_t
+from libc.string cimport memcpy
+
+cdef inline uint64_t load_u64(const void * p):
+    cdef:
+        uint64_t res
+
+    res = 0
+    memcpy(&res, p, sizeof(res))
+    return res
+
+cdef inline uint64_t load_u32(const void * p):
+    cdef:
+        uint32_t res
+
+    res = 0
+    memcpy(&res, p, sizeof(res))
+    return res
+
+cdef inline uint64_t load_u16(const void * p):
+    cdef:
+        uint16_t res
+
+    res = 0
+    memcpy(&res, p, sizeof(res))
+    return res
+
+cdef inline void store_u64(void * p, uint64_t v):
+    memcpy(p, &v, sizeof(v))
+
+cdef inline void store_u32(void * p, uint32_t v):
+    memcpy(p, &v, sizeof(v))
+
+cdef inline void store_u16(void * p, uint16_t v):
+    memcpy(p, &v, sizeof(v))

asynctnt/iproto/buffer.pxd

@@ -42,6 +42,7 @@ cdef class WriteBuffer:
                              const char *data, uint32_t len) except NULL
     cdef char *mp_encode_decimal(self, char *p, object value) except NULL
     cdef char *mp_encode_uuid(self, char *p, object value) except NULL
+    cdef char *mp_encode_datetime(self, char *p, object value) except NULL
     cdef char *mp_encode_array(self, char *p, uint32_t len) except NULL
     cdef char *mp_encode_map(self, char *p, uint32_t len) except NULL
     cdef char *mp_encode_list(self, char *p, list arr) except NULL

asynctnt/iproto/buffer.pyx

@@ -8,6 +8,7 @@ cimport cpython.unicode
 
 from cpython.mem cimport PyMem_Malloc, PyMem_Realloc, PyMem_Free
 from cpython.ref cimport PyObject
+from cpython.datetime cimport datetime
 
 from libc.string cimport memcpy
 from libc.stdint cimport uint32_t, uint64_t, int64_t, uint8_t
@@ -261,6 +262,24 @@ cdef class WriteBuffer:
         self._length += (p - begin)
         return p
 
+    cdef char *mp_encode_datetime(self, char *p, object value) except NULL:
+        cdef:
+            char *begin
+            uint32_t length
+            datetime pydt
+            IProtoDateTime dt
+
+        pydt = <datetime> value
+        datetime_zero(&dt)
+        datetime_from_py(pydt, &dt)
+
+        length = datetime_len(&dt)
+        p = begin = self._ensure_allocated(p, mp_sizeof_ext(length))
+        p = mp_encode_extl(p, tarantool.MP_DATETIME, length)
+        p = datetime_encode(p, &dt)
+        self._length += (p - begin)
+        return p
+
     cdef char *mp_encode_array(self, char *p, uint32_t len) except NULL:
         cdef char *begin
         p = begin = self._ensure_allocated(p, mp_sizeof_array(len))
@@ -385,6 +404,9 @@ cdef class WriteBuffer:
         elif isinstance(o, dict):
             return self.mp_encode_dict(p, <dict> o)
 
+        elif isinstance(o, datetime):
+            return self.mp_encode_datetime(p, o)
+
         elif isinstance(o, Decimal):
             return self.mp_encode_decimal(p, o)
 

asynctnt/iproto/const.pxi

@@ -27,3 +27,5 @@ DEF SCRAMBLE_SIZE = 20
 
 DEF SPACE_VSPACE = 281
 DEF SPACE_VINDEX = 289
+
+DEF DATETIME_TAIL_SZ = 4 + 2 + 2

asynctnt/iproto/ext.pxd

@@ -1,11 +1,31 @@
-from libc.stdint cimport uint32_t, uint8_t
+from libc.stdint cimport uint32_t, uint8_t, int64_t, int32_t, int16_t
 from libc cimport math
+from cpython.datetime cimport datetime
 
 cdef inline uint32_t bcd_len(uint32_t digits_len):
     return <uint32_t> math.floor(digits_len / 2) + 1
 
 cdef uint32_t decimal_len(int exponent, uint32_t digits_count)
-cdef char *decimal_encode(char *p, uint32_t digits_count, uint8_t sign, tuple digits, int exponent)
-cdef object decimal_decode(const char **p, uint32_t length)
+cdef char *decimal_encode(char *p,
+                          uint32_t digits_count,
+                          uint8_t sign,
+                          tuple digits,
+                          int exponent) except NULL
+cdef object decimal_decode(const char ** p, uint32_t length)
 
-cdef object uuid_decode(const char **p, uint32_t length)
+cdef object uuid_decode(const char ** p, uint32_t length)
+
+cdef struct IProtoDateTime:
+    int64_t seconds
+    int32_t nsec
+    int16_t tzoffset
+    int16_t tzindex
+
+cdef void datetime_zero(IProtoDateTime *dt)
+cdef uint32_t datetime_len(IProtoDateTime *dt)
+cdef char *datetime_encode(char *p, IProtoDateTime *dt) except NULL
+cdef int datetime_decode(const char ** p,
+                         uint32_t length,
+                         IProtoDateTime *dt) except -1
+cdef void datetime_from_py(datetime ob, IProtoDateTime *dt)
+cdef object datetime_to_py(IProtoDateTime *dt)

asynctnt/iproto/ext.pyx

@@ -1,3 +1,8 @@
+from libc.stdint cimport uint32_t
+from cpython.datetime cimport PyDateTimeAPI, timedelta_new, datetime, datetime_tzinfo
+cimport cpython.datetime
+from libc.string cimport memcpy
+
 from uuid import UUID
 from decimal import Decimal
 
@@ -13,7 +18,11 @@ cdef uint32_t decimal_len(int exponent, uint32_t digits_count):
 
     return length
 
-cdef char *decimal_encode(char *p, uint32_t digits_count, uint8_t sign, tuple digits, int exponent):
+cdef char *decimal_encode(char *p,
+                          uint32_t digits_count,
+                          uint8_t sign,
+                          tuple digits,
+                          int exponent) except NULL:
     cdef:
         int i
         uint8_t byte
@@ -118,8 +127,85 @@ cdef object decimal_decode(const char ** p, uint32_t length):
 
     return Decimal((<object> <int> sign, digits, <object> exponent))
 
-
-cdef object uuid_decode(const char **p, uint32_t length):
+cdef object uuid_decode(const char ** p, uint32_t length):
     data = cpython.bytes.PyBytes_FromStringAndSize(p[0], length)
     p[0] += length
     return UUID(bytes=data)
+
+cdef inline void datetime_zero(IProtoDateTime *dt):
+    dt.seconds = 0
+    dt.nsec = 0
+    dt.tzoffset = 0
+    dt.tzindex = 0
+
+cdef inline uint32_t datetime_len(IProtoDateTime *dt):
+    cdef uint32_t sz
+    sz = sizeof(int64_t)
+    if dt.nsec != 0 or dt.tzoffset != 0 or dt.tzindex != 0:
+        return sz + DATETIME_TAIL_SZ
+    return sz
+
+cdef char *datetime_encode(char *p, IProtoDateTime *dt) except NULL:
+    store_u64(p, dt.seconds)
+    p += sizeof(dt.seconds)
+    if dt.nsec != 0 or dt.tzoffset != 0 or dt.tzindex != 0:
+        memcpy(p, &dt.nsec, DATETIME_TAIL_SZ)
+        p += DATETIME_TAIL_SZ
+    return p
+
+cdef int datetime_decode(
+        const char ** p,
+        uint32_t length,
+        IProtoDateTime *dt
+) except -1:
+    delta = None
+    tz = None
+
+    dt.seconds = load_u64(p[0])
+    p[0] += sizeof(dt.seconds)
+    length -= sizeof(dt.seconds)
+
+    if length == 0:
+        return 0
+
+    if length != DATETIME_TAIL_SZ:
+        raise ValueError("invalid datetime size. got {} extra bytes".format(
+            length
+        ))
+
+    dt.nsec = load_u32(p[0])
+    p[0] += 4
+    dt.tzoffset = load_u16(p[0])
+    p[0] += 2
+    dt.tzindex = load_u16(p[0])
+    p[0] += 2
+
+cdef void datetime_from_py(datetime ob, IProtoDateTime *dt):
+    cdef:
+        double ts
+        int offset
+    ts = <double> ob.timestamp()
+    dt.seconds = <int64_t> ts
+    dt.nsec = <int32_t> ((ts - <double> dt.seconds) * 1000000) * 1000
+
+    if datetime_tzinfo(ob) is not None:
+        offset = ob.utcoffset().total_seconds()
+        dt.tzoffset = <int16_t> (offset / 60)
+
+cdef object datetime_to_py(IProtoDateTime *dt):
+    cdef:
+        double timestamp
+        object tz
+
+    tz = None
+
+    if dt.tzoffset != 0:
+        delta = timedelta_new(0, <int> dt.tzoffset * 60, 0)
+        tz = timezone_new(delta)
+
+    timestamp = dt.seconds + (<double> dt.nsec) / 1e9
+    return PyDateTimeAPI.DateTime_FromTimestamp(
+        <object> PyDateTimeAPI.DateTimeType,
+        (timestamp,) if tz is None else (timestamp, tz),
+        <object> NULL,
+    )

asynctnt/iproto/protocol.pxd

@@ -6,6 +6,7 @@ include "const.pxi"
 include "cmsgpuck.pxd"
 include "xd.pxd"
 include "python.pxd"
+include "bit.pxd"
 
 include "unicodeutil.pxd"
 include "schema.pxd"

asynctnt/iproto/protocol.pyx

@@ -1,6 +1,8 @@
 # cython: language_level=3
 
 cimport cpython.dict
+from cpython.datetime cimport import_datetime
+import_datetime()
 
 import asyncio
 import enum

asynctnt/iproto/python.pxd

@@ -1,3 +1,76 @@
+from cpython.version cimport PY_VERSION_HEX
+
 cdef extern from "Python.h":
     char *PyByteArray_AS_STRING(object obj)
     int Py_REFCNT(object obj)
+
+
+cdef extern from "datetime.h":
+    """
+    /* Backport for Python 2.x */
+    #if PY_MAJOR_VERSION < 3
+        #ifndef PyDateTime_DELTA_GET_DAYS
+            #define PyDateTime_DELTA_GET_DAYS(o) (((PyDateTime_Delta*)o)->days)
+        #endif
+        #ifndef PyDateTime_DELTA_GET_SECONDS
+            #define PyDateTime_DELTA_GET_SECONDS(o) (((PyDateTime_Delta*)o)->seconds)
+        #endif
+        #ifndef PyDateTime_DELTA_GET_MICROSECONDS
+            #define PyDateTime_DELTA_GET_MICROSECONDS(o) (((PyDateTime_Delta*)o)->microseconds)
+        #endif
+    #endif
+    /* Backport for Python < 3.6 */
+    #if PY_VERSION_HEX < 0x030600a4
+        #ifndef PyDateTime_TIME_GET_FOLD
+            #define PyDateTime_TIME_GET_FOLD(o) ((void)(o), 0)
+        #endif
+        #ifndef PyDateTime_DATE_GET_FOLD
+            #define PyDateTime_DATE_GET_FOLD(o) ((void)(o), 0)
+        #endif
+    #endif
+    /* Backport for Python < 3.6 */
+    #if PY_VERSION_HEX < 0x030600a4
+        #define __Pyx_DateTime_DateTimeWithFold(year, month, day, hour, minute, second, microsecond, tz, fold) \
+            ((void)(fold), PyDateTimeAPI->DateTime_FromDateAndTime(year, month, day, hour, minute, second, \
+                microsecond, tz, PyDateTimeAPI->DateTimeType))
+        #define __Pyx_DateTime_TimeWithFold(hour, minute, second, microsecond, tz, fold) \
+            ((void)(fold), PyDateTimeAPI->Time_FromTime(hour, minute, second, microsecond, tz, PyDateTimeAPI->TimeType))
+    #else /* For Python 3.6+ so that we can pass tz */
+        #define __Pyx_DateTime_DateTimeWithFold(year, month, day, hour, minute, second, microsecond, tz, fold) \
+            PyDateTimeAPI->DateTime_FromDateAndTimeAndFold(year, month, day, hour, minute, second, \
+                microsecond, tz, fold, PyDateTimeAPI->DateTimeType)
+        #define __Pyx_DateTime_TimeWithFold(hour, minute, second, microsecond, tz, fold) \
+            PyDateTimeAPI->Time_FromTimeAndFold(hour, minute, second, microsecond, tz, fold, PyDateTimeAPI->TimeType)
+    #endif
+    /* Backport for Python < 3.7 */
+    #if PY_VERSION_HEX < 0x030700b1
+        #define __Pyx_TimeZone_UTC NULL
+        #define __Pyx_TimeZone_FromOffset(offset) ((void)(offset), (PyObject*)NULL)
+        #define __Pyx_TimeZone_FromOffsetAndName(offset, name) ((void)(offset), (void)(name), (PyObject*)NULL)
+    #else
+        #define __Pyx_TimeZone_UTC PyDateTime_TimeZone_UTC
+        #define __Pyx_TimeZone_FromOffset(offset) PyTimeZone_FromOffset(offset)
+        #define __Pyx_TimeZone_FromOffsetAndName(offset, name) PyTimeZone_FromOffsetAndName(offset, name)
+    #endif
+    /* Backport for Python < 3.10 */
+    #if PY_VERSION_HEX < 0x030a00a1
+        #ifndef PyDateTime_TIME_GET_TZINFO
+            #define PyDateTime_TIME_GET_TZINFO(o) \
+                ((((PyDateTime_Time*)o)->hastzinfo) ? ((PyDateTime_Time*)o)->tzinfo : Py_None)
+        #endif
+        #ifndef PyDateTime_DATE_GET_TZINFO
+            #define PyDateTime_DATE_GET_TZINFO(o) \
+                ((((PyDateTime_DateTime*)o)->hastzinfo) ? ((PyDateTime_DateTime*)o)->tzinfo : Py_None)
+        #endif
+    #endif
+    """
+
+    # The above macros is Python 3.7+ so we use these instead
+    object __Pyx_TimeZone_FromOffset(object offset)
+
+
+cdef inline object timezone_new(object offset):
+    if PY_VERSION_HEX < 0x030700b1:
+        from datetime import timezone
+        return timezone(offset)
+    return __Pyx_TimeZone_FromOffset(offset)

asynctnt/iproto/response.pyx

@@ -199,6 +199,7 @@ cdef object _decode_obj(const char ** p, bytes encoding):
         object map_key
 
         int8_t ext_type
+        IProtoDateTime dt
 
     obj_type = mp_typeof(p[0][0])
     if obj_type == MP_UINT:
@@ -273,8 +274,13 @@ cdef object _decode_obj(const char ** p, bytes encoding):
             return uuid_decode(p, s_len)
         elif ext_type == tarantool.MP_ERROR:
             return parse_iproto_error(p, encoding)
+        elif ext_type == tarantool.MP_DATETIME:
+            datetime_zero(&dt)
+            datetime_decode(p, s_len, &dt)
+            return datetime_to_py(&dt)
         else:  # pragma: nocover
             logger.warning('Unexpected ext type: %d', ext_type)
+            p += s_len  # skip unknown ext
             return None
     else:  # pragma: nocover
         mp_next(p)

asynctnt/iproto/tarantool.pxd

@@ -102,6 +102,7 @@ cdef enum mp_extension_type:
     MP_DECIMAL = 1
     MP_UUID = 2
     MP_ERROR = 3
+    MP_DATETIME = 4
 
 cdef enum iproto_features:
     IPROTO_FEATURE_STREAMS = 0

requirements.txt

@@ -8,3 +8,5 @@ Sphinx
 sphinx_rtd_theme
 sphinxcontrib-asyncio
 coverage
+pytz
+python-dateutil