Add ts_convert_timestamp transport option

Author: Luke Bakken

riak/codecs/pbuf.py

@@ -773,7 +773,8 @@ def encode_timeseries_query(self, table, query, interpolations=None):
         rc = riak.pb.messages.MSG_CODE_TS_QUERY_RESP
         return Msg(mc, req.SerializeToString(), rc)
 
-    def decode_timeseries(self, resp, tsobj):
+    def decode_timeseries(self, resp, tsobj,
+            convert_timestamp=False):
         """
         Fills an TsObject with the appropriate data and
         metadata from a TsGetResp / TsQueryResp.
@@ -783,6 +784,8 @@ def decode_timeseries(self, resp, tsobj):
                     riak.pb.riak_ts_pb2.TsGetResp
         :param tsobj: a TsObject
         :type tsobj: TsObject
+        :param convert_timestamp: Convert timestamps to datetime objects
+        :type tsobj: boolean
         """
         if resp.columns is not None:
             col_names = []
@@ -798,7 +801,7 @@ def decode_timeseries(self, resp, tsobj):
             for row in resp.rows:
                 tsobj.rows.append(
                     self.decode_timeseries_row(
-                        row, resp.columns))
+                        row, resp.columns, convert_timestamp))
 
     def decode_timeseries_col_type(self, col_type):
         # NB: these match the atom names for column types
@@ -816,7 +819,8 @@ def decode_timeseries_col_type(self, col_type):
             msg = 'could not decode column type: {}'.format(col_type)
             raise RiakError(msg)
 
-    def decode_timeseries_row(self, tsrow, tscols=None):
+    def decode_timeseries_row(self, tsrow, tscols=None,
+            convert_timestamp=False):
         """
         Decodes a TsRow into a list
 
@@ -850,8 +854,10 @@ def decode_timeseries_row(self, tsrow, tscols=None):
                 if col and col.type != TsColumnType.Value('TIMESTAMP'):
                     raise TypeError('expected TIMESTAMP column')
                 else:
-                    dt = datetime_from_unix_time_millis(
-                        cell.timestamp_value)
+                    dt = cell.timestamp_value
+                    if convert_timestamp:
+                        dt = datetime_from_unix_time_millis(
+                            cell.timestamp_value)
                     row.append(dt)
             elif cell.HasField('boolean_value'):
                 if col and col.type != TsColumnType.Value('BOOLEAN'):

riak/codecs/ttb.py

@@ -138,7 +138,8 @@ def encode_timeseries_query(self, table, query, interpolations=None):
         rc = MSG_CODE_TS_TTB_MSG
         return Msg(mc, encode(req), rc)
 
-    def decode_timeseries(self, resp_ttb, tsobj):
+    def decode_timeseries(self, resp_ttb, tsobj,
+            convert_timestamp=False):
         """
         Fills an TsObject with the appropriate data and
         metadata from a TTB-encoded TsGetResp / TsQueryResp.
@@ -147,6 +148,8 @@ def decode_timeseries(self, resp_ttb, tsobj):
         :type resp_ttb: TTB-encoded tsqueryrsp or tsgetresp
         :param tsobj: a TsObject
         :type tsobj: TsObject
+        :param convert_timestamp: Convert timestamps to datetime objects
+        :type tsobj: boolean
         """
         if resp_ttb is None:
             return tsobj
@@ -169,7 +172,8 @@ def decode_timeseries(self, resp_ttb, tsobj):
                 tsobj.rows = []
                 for resp_row in resp_rows:
                     tsobj.rows.append(
-                        self.decode_timeseries_row(resp_row, resp_coltypes))
+                        self.decode_timeseries_row(resp_row, resp_coltypes,
+                            convert_timestamp))
             else:
                 raise RiakError(
                     "Expected 3-tuple in response, got: {}".format(resp_data))
@@ -181,14 +185,16 @@ def decode_timeseries_cols(self, cnames, ctypes):
         ctypes = [str(ctype) for ctype in ctypes]
         return TsColumns(cnames, ctypes)
 
-    def decode_timeseries_row(self, tsrow, tsct):
+    def decode_timeseries_row(self, tsrow, tsct, convert_timestamp=False):
         """
         Decodes a TTB-encoded TsRow into a list
 
         :param tsrow: the TTB decoded TsRow to decode.
         :type tsrow: TTB dncoded row
         :param tsct: the TTB decoded column types (atoms).
         :type tsct: list
+        :param convert_timestamp: Convert timestamps to datetime objects
+        :type tsobj: boolean
         :rtype list
         """
         row = []
@@ -198,7 +204,7 @@ def decode_timeseries_row(self, tsrow, tsct):
             elif isinstance(cell, list) and len(cell) == 0:
                 row.append(None)
             else:
-                if tsct[i] == timestamp_a:
+                if convert_timestamp and tsct[i] == timestamp_a:
                     row.append(datetime_from_unix_time_millis(cell))
                 else:
                     row.append(cell)

riak/tests/test_timeseries_pbuf.py

@@ -161,7 +161,7 @@ def test_decode_data_from_query(self):
 
         tsobj = TsObject(None, self.table)
         c = PbufCodec()
-        c.decode_timeseries(tqr, tsobj)
+        c.decode_timeseries(tqr, tsobj, True)
 
         self.assertEqual(len(tsobj.rows), len(self.rows))
         self.assertEqual(len(tsobj.columns.names), len(tqr.columns))
@@ -197,7 +197,8 @@ def test_decode_data_from_query(self):
 @unittest.skipUnless(is_timeseries_supported() and RUN_TIMESERIES,
                      'Timeseries not supported or RUN_TIMESERIES is 0')
 class TimeseriesPbufTests(IntegrationTestBase, unittest.TestCase):
-    client_options = {'transport_options': {'use_ttb': False}}
+    client_options = {'transport_options':
+            {'use_ttb': False, 'ts_convert_timestamp': True}}
 
     @classmethod
     def setUpClass(cls):

riak/tests/test_timeseries_ttb.py

@@ -92,8 +92,9 @@ def test_decode_data_from_get(self):
             self.assertEqual(r[0], dr[0].encode('utf-8'))
             self.assertEqual(r[1], dr[1])
             self.assertEqual(r[2], dr[2])
-            dt = datetime_from_unix_time_millis(dr[3])
-            self.assertEqual(r[3], dt)
+            # NB *not* decoding timestamps
+            # dt = datetime_from_unix_time_millis(dr[3])
+            self.assertEqual(r[3], dr[3])
             if i == 0:
                 self.assertEqual(r[4], True)
             else:
@@ -123,7 +124,8 @@ def test_encode_data_for_put(self):
 @unittest.skipUnless(is_timeseries_supported() and RUN_TIMESERIES,
                      'Timeseries not supported or RUN_TIMESERIES is 0')
 class TimeseriesTtbTests(IntegrationTestBase, unittest.TestCase):
-    client_options = {'transport_options': {'use_ttb': True}}
+    client_options = {'transport_options':
+            {'use_ttb': True, 'ts_convert_timestamp': True}}
 
     @classmethod
     def setUpClass(cls):

riak/transports/tcp/stream.py

@@ -174,6 +174,10 @@ class PbufTsKeyStream(PbufStream, TtbCodec):
 
     _expect = riak.pb.messages.MSG_CODE_TS_LIST_KEYS_RESP
 
+    def __init__(self, transport, codec, convert_timestamp=False):
+        super(PbufTsKeyStream, self).__init__(transport, codec)
+        self._convert_timestamp = convert_timestamp
+
     def next(self):
         response = super(PbufTsKeyStream, self).next()
 
@@ -182,7 +186,8 @@ def next(self):
 
         keys = []
         for tsrow in response.keys:
-            keys.append(self.codec.decode_timeseries_row(tsrow))
+            keys.append(self.codec.decode_timeseries_row(tsrow,
+                convert_timestamp=self._convert_timestamp))
 
         return keys
 

riak/transports/tcp/transport.py

@@ -38,6 +38,8 @@ def __init__(self,
         self._pbuf_c = None
         self._ttb_c = None
         self._use_ttb = kwargs.get('use_ttb', True)
+        self._ts_convert_timestamp = \
+                kwargs.get('ts_convert_timestamp', False)
 
     def _get_pbuf_codec(self):
         if not self._pbuf_c:
@@ -147,7 +149,8 @@ def ts_get(self, table, key):
         msg = codec.encode_timeseries_keyreq(table, key)
         resp_code, resp = self._request(msg, codec)
         tsobj = TsObject(self._client, table)
-        codec.decode_timeseries(resp, tsobj)
+        codec.decode_timeseries(resp, tsobj,
+                self._ts_convert_timestamp)
         return tsobj
 
     def ts_put(self, tsobj):
@@ -173,7 +176,8 @@ def ts_query(self, table, query, interpolations=None):
         msg = codec.encode_timeseries_query(table, query, interpolations)
         resp_code, resp = self._request(msg, codec)
         tsobj = TsObject(self._client, table)
-        codec.decode_timeseries(resp, tsobj)
+        codec.decode_timeseries(resp, tsobj,
+                self._ts_convert_timestamp)
         return tsobj
 
     def ts_stream_keys(self, table, timeout=None):
@@ -185,7 +189,7 @@ def ts_stream_keys(self, table, timeout=None):
         codec = self._get_codec(msg_code)
         msg = codec.encode_timeseries_listkeysreq(table, timeout)
         self._send_msg(msg.msg_code, msg.data)
-        return PbufTsKeyStream(self, codec)
+        return PbufTsKeyStream(self, codec, self._ts_convert_timestamp)
 
     def delete(self, robj, rw=None, r=None, w=None, dw=None,
                pr=None, pw=None, timeout=None):