Merge pull request twisted#23 from jdavisp3/timestamp-2

Add full Timestamp support.

tests/test_objects.py

@@ -12,17 +12,19 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import time
 from StringIO import StringIO
 
 import txmongo
 from txmongo import database
 from txmongo import collection
 from txmongo import gridfs
-from txmongo._pymongo import objectid
-from txmongo._gridfs import GridIn, GridOut
+from txmongo import filter as qf
+from txmongo._pymongo import objectid, timestamp
+from txmongo._gridfs import GridIn
 from twisted.trial import unittest
 from twisted.trial import runner
-from twisted.internet import base, defer, reactor
+from twisted.internet import base, defer
 
 mongo_host="localhost"
 mongo_port=27017
@@ -73,6 +75,50 @@ def test_MongoOperations(self):
         # disconnect
         yield conn.disconnect()
 
+    @defer.inlineCallbacks
+    def test_Timestamps(self):
+        """Tests mongo operations with Timestamps"""
+        conn = yield txmongo.MongoConnection(mongo_host, mongo_port)
+        test = conn.foo.test_ts
+
+        # insert with specific timestamp
+        doc1 = {'_id':objectid.ObjectId(),
+                'ts':timestamp.Timestamp(1, 2)}
+        yield test.insert(doc1, safe=True)
+
+        result = yield test.find_one(doc1)
+        self.assertEqual(result.get('ts').time, 1)
+        self.assertEqual(result.get('ts').inc, 2)
+
+        # insert with specific timestamp
+        doc2 = {'_id':objectid.ObjectId(),
+                'ts':timestamp.Timestamp(2, 1)}
+        yield test.insert(doc2, safe=True)
+
+        # the objects come back sorted by ts correctly.
+        # (test that we stored inc/time in the right fields)
+        result = yield test.find(filter=qf.sort(qf.ASCENDING('ts')))
+        self.assertEqual(result[0]['_id'], doc1['_id'])
+        self.assertEqual(result[1]['_id'], doc2['_id'])
+
+        # insert with null timestamp
+        doc3 = {'_id':objectid.ObjectId(),
+                'ts':timestamp.Timestamp(0, 0)}
+        yield test.insert(doc3, safe=True)
+
+        # time field loaded correctly
+        result = yield test.find_one(doc3['_id'])
+        now = time.time()
+        self.assertTrue(now - 2 <= result['ts'].time <= now)
+
+        # delete
+        yield test.remove(doc1["_id"], safe=True)
+        yield test.remove(doc2["_id"], safe=True)
+        yield test.remove(doc3["_id"], safe=True)
+
+        # disconnect
+        yield conn.disconnect()
+
 
 class TestGridFsObjects(unittest.TestCase):
     """ Test the GridFS operations from txmongo._gridfs """

txmongo/_pymongo/_cbsonmodule.c

@@ -43,6 +43,7 @@ static PyObject* ObjectId;
 static PyObject* DBRef;
 static PyObject* RECompile;
 static PyObject* UUID;
+static PyObject* Timestamp;
 
 #if PY_VERSION_HEX < 0x02050000 && !defined(PY_SSIZE_T_MIN)
 typedef int Py_ssize_t;
@@ -461,6 +462,32 @@ static int write_element_to_buffer(bson_buffer* buffer, int type_byte, PyObject*
         Py_DECREF(as_doc);
         *(buffer->buffer + type_byte) = 0x03;
         return 1;
+    } else if (PyObject_IsInstance(value, Timestamp)) {
+        PyObject* obj;
+        long i;
+
+        obj = PyObject_GetAttrString(value, "inc");
+        if (!obj) {
+            return 0;
+        }
+        i = PyInt_AsLong(obj);
+        Py_DECREF(obj);
+        if (!buffer_write_bytes(buffer, (const char*)&i, 4)) {
+            return 0;
+        }
+
+        obj = PyObject_GetAttrString(value, "time");
+        if (!obj) {
+            return 0;
+        }
+        i = PyInt_AsLong(obj);
+        Py_DECREF(obj);
+        if (!buffer_write_bytes(buffer, (const char*)&i, 4)) {
+            return 0;
+        }
+
+        *(buffer->buffer + type_byte) = 0x11;
+        return 1;
     }
     else if (PyObject_HasAttrString(value, "pattern") &&
              PyObject_HasAttrString(value, "flags")) { /* TODO just a proxy for checking if it is a compiled re */
@@ -588,11 +615,19 @@ static int write_element_to_buffer(bson_buffer* buffer, int type_byte, PyObject*
             Py_DECREF(module);
         }
 
+        module = PyImport_ImportModule("txmongo._pymongo.timestamp");
+        if (!module) {
+            return 0;
+        }
+        Timestamp = PyObject_GetAttrString(module, "Timestamp");
+        Py_DECREF(module);
+
         if (PyObject_IsInstance(value, Binary) ||
             PyObject_IsInstance(value, Code) ||
             PyObject_IsInstance(value, ObjectId) ||
             PyObject_IsInstance(value, DBRef) ||
-            (UUID && PyObject_IsInstance(value, UUID))) {
+            (UUID && PyObject_IsInstance(value, UUID)) ||
+            PyObject_IsInstance(value, Timestamp)) {
 
             PyErr_SetString(PyExc_RuntimeError,
                             "A python module was reloaded without the C extension being reloaded.\n"
@@ -1275,11 +1310,10 @@ static PyObject* get_value(const char* buffer, int* position, int type) {
         }
     case 17:
         {
-            int i,
-                j;
-            memcpy(&i, buffer + *position, 4);
-            memcpy(&j, buffer + *position + 4, 4);
-            value = Py_BuildValue("(ii)", i, j);
+            unsigned int time, inc;
+            memcpy(&inc, buffer + *position, 4);
+            memcpy(&time, buffer + *position + 4, 4);
+            value = PyObject_CallFunction(Timestamp, "II", time, inc);
             if (!value) {
                 return NULL;
             }
@@ -1483,4 +1517,11 @@ PyMODINIT_FUNC init_cbson(void) {
         UUID = PyObject_GetAttrString(module, "UUID");
         Py_DECREF(module);
     }
+
+    module = PyImport_ImportModule("txmongo._pymongo.timestamp");
+    if (!module) {
+        return;
+    }
+    Timestamp = PyObject_GetAttrString(module, "Timestamp");
+    Py_DECREF(module);
 }

txmongo/_pymongo/bson.py

@@ -27,6 +27,7 @@
 from txmongo._pymongo.code import Code
 from txmongo._pymongo.objectid import ObjectId
 from txmongo._pymongo.son import SON
+from txmongo._pymongo.timestamp import Timestamp
 from txmongo._pymongo.errors import InvalidBSON, InvalidDocument
 from txmongo._pymongo.errors import InvalidName, InvalidStringData
 
@@ -325,9 +326,9 @@ def _get_ref(data):
 
 
 def _get_timestamp(data):
-    (timestamp, data) = _get_int(data)
     (inc, data) = _get_int(data)
-    return ((timestamp, inc), data)
+    (timestamp, data) = _get_int(data)
+    return (Timestamp(timestamp, inc), data)
 
 
 def _get_long(data):
@@ -466,6 +467,10 @@ def _element_to_bson(key, value, check_keys):
         return "\x0B" + name + _make_c_string(pattern, True) + _make_c_string(flags)
     if isinstance(value, DBRef):
         return _element_to_bson(key, value.as_doc(), False)
+    if isinstance(value, Timestamp):
+        inc = struct.pack("<i", value.inc)
+        timestamp = struct.pack("<i", value.time)
+        return "\x11" + name + inc + timestamp
 
     raise InvalidDocument("cannot convert value of type %s to bson" %
                           type(value))

txmongo/_pymongo/timestamp.py

@@ -0,0 +1,96 @@
+# Copyright 2010 10gen, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tools for representing MongoDB internal Timestamps.
+"""
+
+import calendar
+import datetime
+
+from tz_util import utc
+
+
+class Timestamp(object):
+    """MongoDB internal timestamps used in the opLog.
+    """
+
+    def __init__(self, time, inc):
+        """Create a new :class:`Timestamp`.
+
+        This class is only for use with the MongoDB opLog. If you need
+        to store a regular timestamp, please use a
+        :class:`~datetime.datetime`.
+
+        Raises :class:`TypeError` if `time` is not an instance of
+        :class: `int` or :class:`~datetime.datetime`, or `inc` is not
+        an instance of :class:`int`. Raises :class:`ValueError` if
+        `time` or `inc` is not in [0, 2**32).
+
+        :Parameters:
+          - `time`: time in seconds since epoch UTC, or a naive UTC
+            :class:`~datetime.datetime`, or an aware
+            :class:`~datetime.datetime`
+          - `inc`: the incrementing counter
+
+        .. versionchanged:: 1.7
+           `time` can now be a :class:`~datetime.datetime` instance.
+        """
+        if isinstance(time, datetime.datetime):
+            if time.utcoffset() is not None:
+                time = time - time.utcoffset()
+            time = int(calendar.timegm(time.timetuple()))
+        if not isinstance(time, (int, long)):
+            raise TypeError("time must be an instance of int")
+        if not isinstance(inc, (int, long)):
+            raise TypeError("inc must be an instance of int")
+        if not 0 <= time < 2 ** 32:
+            raise ValueError("time must be contained in [0, 2**32)")
+        if not 0 <= inc < 2 ** 32:
+            raise ValueError("inc must be contained in [0, 2**32)")
+
+        self.__time = time
+        self.__inc = inc
+
+    @property
+    def time(self):
+        """Get the time portion of this :class:`Timestamp`.
+        """
+        return self.__time
+
+    @property
+    def inc(self):
+        """Get the inc portion of this :class:`Timestamp`.
+        """
+        return self.__inc
+
+    def __eq__(self, other):
+        if isinstance(other, Timestamp):
+            return (self.__time == other.time and self.__inc == other.inc)
+        else:
+            return NotImplemented
+
+    def __ne__(self, other):
+        return not self == other
+
+    def __repr__(self):
+        return "Timestamp(%s, %s)" % (self.__time, self.__inc)
+
+    def as_datetime(self):
+        """Return a :class:`~datetime.datetime` instance corresponding
+        to the time portion of this :class:`Timestamp`.
+
+        .. versionchanged:: 1.8
+           The returned datetime is now timezone aware.
+        """
+        return datetime.datetime.fromtimestamp(self.__time, utc)

txmongo/_pymongo/tz_util.py

@@ -0,0 +1,45 @@
+# Copyright 2010 10gen, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Timezone related utilities for BSON."""
+
+from datetime import (timedelta,
+                      tzinfo)
+
+ZERO = timedelta(0)
+
+
+class FixedOffset(tzinfo):
+    """Fixed offset timezone, in minutes east from UTC.
+
+    Implementation from the Python `standard library documentation
+    <http://docs.python.org/library/datetime.html#tzinfo-objects>`_.
+    """
+
+    def __init__(self, offset, name):
+        self.__offset = timedelta(minutes=offset)
+        self.__name = name
+
+    def utcoffset(self, dt):
+        return self.__offset
+
+    def tzname(self, dt):
+        return self.__name
+
+    def dst(self, dt):
+        return ZERO
+
+
+utc = FixedOffset(0, "UTC")
+"""Fixed offset timezone representing UTC."""