Merge pull request #511 from basho/bugs/am/ts-list-keys-table-string

Allow TS ListKeys to accept either a Table object or the table name as a String

Author: alexmoore

riak/client/operations.py

@@ -659,15 +659,18 @@ def ts_stream_keys(self, table, timeout=None):
             stream.close()
 
         :param table: the table from which to stream keys
-        :type table: Table
+        :type table: string or :class:`Table <riak.table.Table>`
         :param timeout: a timeout value in milliseconds
         :type timeout: int
         :rtype: iterator
         """
+        t = table
+        if isinstance(t, six.string_types):
+            t = Table(self, table)
         _validate_timeout(timeout)
         resource = self._acquire()
         transport = resource.object
-        stream = transport.ts_stream_keys(table, timeout)
+        stream = transport.ts_stream_keys(t, timeout)
         stream.attach(resource)
         try:
             for keylist in stream:

riak/tests/test_timeseries_pbuf.py

@@ -436,16 +436,25 @@ def test_stream_keys(self):
         table = Table(self.client, table_name)
         streamed_keys = []
         for keylist in table.stream_keys():
-            self.assertNotEqual([], keylist)
-            streamed_keys += keylist
-            for key in keylist:
-                self.assertIsInstance(key, list)
-                self.assertEqual(len(key), 3)
-                self.assertEqual(bytes_to_str(key[0]), 'hash1')
-                self.assertEqual(bytes_to_str(key[1]), 'user2')
-                self.assertIsInstance(key[2], datetime.datetime)
+            self.validate_keylist(streamed_keys, keylist)
         self.assertGreater(len(streamed_keys), 0)
 
+    def test_stream_keys_from_string_table(self):
+        streamed_keys = []
+        for keylist in self.client.ts_stream_keys(table_name):
+            self.validate_keylist(streamed_keys, keylist)
+        self.assertGreater(len(streamed_keys), 0)
+
+    def validate_keylist(self, streamed_keys, keylist):
+        self.assertNotEqual([], keylist)
+        streamed_keys += keylist
+        for key in keylist:
+            self.assertIsInstance(key, list)
+            self.assertEqual(len(key), 3)
+            self.assertEqual(bytes_to_str(key[0]), 'hash1')
+            self.assertEqual(bytes_to_str(key[1]), 'user2')
+            self.assertIsInstance(key[2], datetime.datetime)
+
     def test_delete_single_value(self):
         key = ['hash1', 'user2', self.twentyFiveMinsAgo]
         rslt = self.client.ts_delete(table_name, key)