PYTHON-1824 Allow encoding/decoding out-of-range datetimes via DatetimeMS and datetime_conversion (#981)

https://jira.mongodb.org/browse/PYTHON-1824

Co-authored-by: Ben Warner <[email protected]>

Author: bcwarner

bson/__init__.py

@@ -54,7 +54,6 @@
    subtype 0. It will be decoded back to bytes.
 """
 
-import calendar
 import datetime
 import itertools
 import re
@@ -100,9 +99,18 @@
 from bson.codec_options import (
     DEFAULT_CODEC_OPTIONS,
     CodecOptions,
+    DatetimeConversionOpts,
     _DocumentType,
     _raw_document_class,
 )
+from bson.datetime_ms import (
+    EPOCH_AWARE,
+    EPOCH_NAIVE,
+    DatetimeMS,
+    _datetime_to_millis,
+    _millis_to_datetime,
+    utc,
+)
 from bson.dbref import DBRef
 from bson.decimal128 import Decimal128
 from bson.errors import InvalidBSON, InvalidDocument, InvalidStringData
@@ -113,7 +121,6 @@
 from bson.regex import Regex
 from bson.son import RE_TYPE, SON
 from bson.timestamp import Timestamp
-from bson.tz_util import utc
 
 # Import some modules for type-checking only.
 if TYPE_CHECKING:
@@ -187,12 +194,10 @@
     "is_valid",
     "BSON",
     "has_c",
+    "DatetimeConversionOpts",
+    "DatetimeMS",
 ]
 
-EPOCH_AWARE = datetime.datetime.fromtimestamp(0, utc)
-EPOCH_NAIVE = datetime.datetime.utcfromtimestamp(0)
-
-
 BSONNUM = b"\x01"  # Floating point
 BSONSTR = b"\x02"  # UTF-8 string
 BSONOBJ = b"\x03"  # Embedded document
@@ -413,7 +418,7 @@ def _get_boolean(
 
 def _get_date(
     data: Any, view: Any, position: int, dummy0: int, opts: CodecOptions, dummy1: Any
-) -> Tuple[datetime.datetime, int]:
+) -> Tuple[Union[datetime.datetime, DatetimeMS], int]:
     """Decode a BSON datetime to python datetime.datetime."""
     return _millis_to_datetime(_UNPACK_LONG_FROM(data, position)[0], opts), position + 8
 
@@ -724,6 +729,12 @@ def _encode_datetime(name: bytes, value: datetime.datetime, dummy0: Any, dummy1:
     return b"\x09" + name + _PACK_LONG(millis)
 
 
+def _encode_datetime_ms(name: bytes, value: DatetimeMS, dummy0: Any, dummy1: Any) -> bytes:
+    """Encode datetime.datetime."""
+    millis = int(value)
+    return b"\x09" + name + _PACK_LONG(millis)
+
+
 def _encode_none(name: bytes, dummy0: Any, dummy1: Any, dummy2: Any) -> bytes:
     """Encode python None."""
     return b"\x0A" + name
@@ -814,6 +825,7 @@ def _encode_maxkey(name: bytes, dummy0: Any, dummy1: Any, dummy2: Any) -> bytes:
     bool: _encode_bool,
     bytes: _encode_bytes,
     datetime.datetime: _encode_datetime,
+    DatetimeMS: _encode_datetime_ms,
     dict: _encode_mapping,
     float: _encode_float,
     int: _encode_int,
@@ -948,27 +960,6 @@ def _dict_to_bson(doc: Any, check_keys: bool, opts: CodecOptions, top_level: boo
     _dict_to_bson = _cbson._dict_to_bson  # noqa: F811
 
 
-def _millis_to_datetime(millis: int, opts: CodecOptions) -> datetime.datetime:
-    """Convert milliseconds since epoch UTC to datetime."""
-    diff = ((millis % 1000) + 1000) % 1000
-    seconds = (millis - diff) // 1000
-    micros = diff * 1000
-    if opts.tz_aware:
-        dt = EPOCH_AWARE + datetime.timedelta(seconds=seconds, microseconds=micros)
-        if opts.tzinfo:
-            dt = dt.astimezone(opts.tzinfo)
-        return dt
-    else:
-        return EPOCH_NAIVE + datetime.timedelta(seconds=seconds, microseconds=micros)
-
-
-def _datetime_to_millis(dtm: datetime.datetime) -> int:
-    """Convert datetime to milliseconds since epoch UTC."""
-    if dtm.utcoffset() is not None:
-        dtm = dtm - dtm.utcoffset()  # type: ignore
-    return int(calendar.timegm(dtm.timetuple()) * 1000 + dtm.microsecond // 1000)
-
-
 _CODEC_OPTIONS_TYPE_ERROR = TypeError("codec_options must be an instance of CodecOptions")
 
 

bson/_cbsonmodule.c

@@ -52,6 +52,9 @@ struct module_state {
     PyObject* BSONInt64;
     PyObject* Decimal128;
     PyObject* Mapping;
+    PyObject* DatetimeMS;
+    PyObject* _min_datetime_ms;
+    PyObject* _max_datetime_ms;
 };
 
 #define GETSTATE(m) ((struct module_state*)PyModule_GetState(m))
@@ -72,6 +75,12 @@ struct module_state {
 /* The smallest possible BSON document, i.e. "{}" */
 #define BSON_MIN_SIZE 5
 
+/* Datetime codec options */
+#define DATETIME 1
+#define DATETIME_CLAMP 2
+#define DATETIME_MS 3
+#define DATETIME_AUTO 4
+
 /* Get an error class from the bson.errors module.
  *
  * Returns a new ref */
@@ -179,6 +188,45 @@ static long long millis_from_datetime(PyObject* datetime) {
     return millis;
 }
 
+/* Extended-range datetime, returns a DatetimeMS object with millis */
+static PyObject* datetime_ms_from_millis(PyObject* self, long long millis){
+    // Allocate a new DatetimeMS object.
+    struct module_state *state = GETSTATE(self);
+
+    PyObject* dt;
+    PyObject* ll_millis;
+
+    if (!(ll_millis = PyLong_FromLongLong(millis))){
+        return NULL;
+    }
+    dt = PyObject_CallFunctionObjArgs(state->DatetimeMS, ll_millis, NULL);
+    Py_DECREF(ll_millis);
+    return dt;
+}
+
+/* Extended-range datetime, takes a DatetimeMS object and extracts the long long value. */
+static int millis_from_datetime_ms(PyObject* dt, long long* out){
+    PyObject* ll_millis;
+    long long millis;
+
+    if (!(ll_millis = PyNumber_Long(dt))){
+        if (PyErr_Occurred()) { // TypeError
+            return 0;
+        }
+    }
+
+    if ((millis = PyLong_AsLongLong(ll_millis)) == -1){
+        if (PyErr_Occurred()) { /* Overflow */
+            PyErr_SetString(PyExc_OverflowError,
+                            "MongoDB datetimes can only handle up to 8-byte ints");
+            return 0;
+        }
+    }
+    Py_DECREF(ll_millis);
+    *out = millis;
+    return 1;
+}
+
 /* Just make this compatible w/ the old API. */
 int buffer_write_bytes(buffer_t buffer, const char* data, int size) {
     if (pymongo_buffer_write(buffer, data, size)) {
@@ -342,7 +390,10 @@ static int _load_python_objects(PyObject* module) {
         _load_object(&state->BSONInt64, "bson.int64", "Int64") ||
         _load_object(&state->Decimal128, "bson.decimal128", "Decimal128") ||
         _load_object(&state->UUID, "uuid", "UUID") ||
-        _load_object(&state->Mapping, "collections.abc", "Mapping")) {
+        _load_object(&state->Mapping, "collections.abc", "Mapping") ||
+        _load_object(&state->DatetimeMS, "bson.datetime_ms", "DatetimeMS") ||
+        _load_object(&state->_min_datetime_ms, "bson.datetime_ms", "_min_datetime_ms") ||
+        _load_object(&state->_max_datetime_ms, "bson.datetime_ms", "_max_datetime_ms")) {
         return 1;
     }
     /* Reload our REType hack too. */
@@ -466,13 +517,14 @@ int convert_codec_options(PyObject* options_obj, void* p) {
 
     options->unicode_decode_error_handler = NULL;
 
-    if (!PyArg_ParseTuple(options_obj, "ObbzOO",
+    if (!PyArg_ParseTuple(options_obj, "ObbzOOb",
                           &options->document_class,
                           &options->tz_aware,
                           &options->uuid_rep,
                           &options->unicode_decode_error_handler,
                           &options->tzinfo,
-                          &type_registry_obj))
+                          &type_registry_obj,
+                          &options->datetime_conversion))
         return 0;
 
     type_marker = _type_marker(options->document_class);
@@ -1049,6 +1101,13 @@ static int _write_element_to_buffer(PyObject* self, buffer_t buffer,
         }
         *(pymongo_buffer_get_buffer(buffer) + type_byte) = 0x09;
         return buffer_write_int64(buffer, (int64_t)millis);
+    } else if (PyObject_TypeCheck(value, (PyTypeObject *) state->DatetimeMS)) {
+        long long millis;
+        if (!millis_from_datetime_ms(value, &millis)) {
+            return 0;
+        }
+        *(pymongo_buffer_get_buffer(buffer) + type_byte) = 0x09;
+        return buffer_write_int64(buffer, (int64_t)millis);
     } else if (PyObject_TypeCheck(value, state->REType)) {
         return _write_regex_to_buffer(buffer, type_byte, value);
     }
@@ -1854,8 +1913,79 @@ static PyObject* get_value(PyObject* self, PyObject* name, const char* buffer,
             }
             memcpy(&millis, buffer + *position, 8);
             millis = (int64_t)BSON_UINT64_FROM_LE(millis);
-            naive = datetime_from_millis(millis);
             *position += 8;
+
+            if (options->datetime_conversion == DATETIME_MS){
+                value = datetime_ms_from_millis(self, millis);
+                break;
+            }
+
+            int dt_clamp = options->datetime_conversion == DATETIME_CLAMP;
+            int dt_auto = options->datetime_conversion == DATETIME_AUTO;
+
+
+            if (dt_clamp || dt_auto){
+                PyObject *min_millis_fn = _get_object(state->_min_datetime_ms, "bson.datetime_ms", "_min_datetime_ms");
+                PyObject *max_millis_fn = _get_object(state->_max_datetime_ms, "bson.datetime_ms", "_max_datetime_ms");
+                PyObject *min_millis_fn_res;
+                PyObject *max_millis_fn_res;
+                int64_t min_millis;
+                int64_t max_millis;
+
+                if (min_millis_fn == NULL || max_millis_fn == NULL) {
+                    Py_XDECREF(min_millis_fn);
+                    Py_XDECREF(max_millis_fn);
+                    goto invalid;
+                }
+
+                if (options->tz_aware){
+                    PyObject* tzinfo = options->tzinfo;
+                    if (tzinfo == Py_None) {
+                        // Default to UTC.
+                        utc_type = _get_object(state->UTC, "bson.tz_util", "utc");
+                        tzinfo = utc_type;
+                    }
+                    min_millis_fn_res = PyObject_CallFunctionObjArgs(min_millis_fn, tzinfo, NULL);
+                    max_millis_fn_res = PyObject_CallFunctionObjArgs(max_millis_fn, tzinfo, NULL);
+                } else {
+                    min_millis_fn_res = PyObject_CallObject(min_millis_fn, NULL);
+                    max_millis_fn_res = PyObject_CallObject(max_millis_fn, NULL);
+                }
+
+                Py_DECREF(min_millis_fn);
+                Py_DECREF(max_millis_fn);
+
+                if (!min_millis_fn_res || !max_millis_fn_res){
+                    Py_XDECREF(min_millis_fn_res);
+                    Py_XDECREF(max_millis_fn_res);
+                    goto invalid;
+                }
+
+                min_millis = PyLong_AsLongLong(min_millis_fn_res);
+                max_millis = PyLong_AsLongLong(max_millis_fn_res);
+
+                if ((min_millis == -1 || max_millis == -1) && PyErr_Occurred())
+                {
+                    // min/max_millis check
+                    goto invalid;
+                }
+
+                if (dt_clamp) {
+                    if (millis < min_millis) {
+                        millis = min_millis;
+                    } else if (millis > max_millis) {
+                        millis = max_millis;
+                    }
+                    // Continues from here to return a datetime.
+                } else if (dt_auto) {
+                    if (millis < min_millis || millis > max_millis){
+                        value = datetime_ms_from_millis(self, millis);
+                        break; // Out-of-range so done.
+                    }
+                }
+            }
+
+            naive = datetime_from_millis(millis);
             if (!options->tz_aware) { /* In the naive case, we're done here. */
                 value = naive;
                 break;

bson/_cbsonmodule.h

@@ -62,6 +62,7 @@ typedef struct codec_options_t {
     char* unicode_decode_error_handler;
     PyObject* tzinfo;
     type_registry_t type_registry;
+    unsigned char datetime_conversion;
     PyObject* options_obj;
     unsigned char is_raw_bson;
 } codec_options_t;