Added counter column support to xget

ljsking · May 31, 2012 · 4b90354 · 4b90354
1 parent 1fc6e2f
commit 4b90354
Show file tree

Hide file tree

Showing 2 changed files with 42 additions and 4 deletions.
diff --git a/pycassa/columnfamily.py b/pycassa/columnfamily.py
@@ -541,10 +541,16 @@ def xget(self, key, column_start="", column_finish="", column_reversed=False,
                     continue
 
                 if self.super:
-                    scol = cosc.super_column
+                    if self._have_counters:
+                        scol = cosc.counter_super_column
+                    else:
+                        scol = cosc.super_column
                     yield (self._unpack_name(scol.name, True), self._scol_to_dict(scol, include_timestamp))
                 else:
-                    col = cosc.column
+                    if self._have_counters:
+                        col = cosc.counter_column
+                    else:
+                        col = cosc.column
                     yield (self._unpack_name(col.name, False), self._col_to_dict(col, include_timestamp))
 
                 count += 1
@@ -555,9 +561,15 @@ def xget(self, key, column_start="", column_finish="", column_reversed=False,
                 return
 
             if self.super:
-                last_name = list_cosc[-1].super_column.name
+                if self._have_counters:
+                    last_name = list_cosc[-1].counter_super_column.name
+                else:
+                    last_name = list_cosc[-1].super_column.name
             else:
-                last_name = list_cosc[-1].column.name
+                if self._have_counters:
+                    last_name = list_cosc[-1].counter_column.name
+                else:
+                    last_name = list_cosc[-1].column.name
             i += 1
 
     def get(self, key, columns=None, column_start="", column_finish="",

diff --git a/tests/test_columnfamily.py b/tests/test_columnfamily.py
@@ -410,6 +410,19 @@ def test_xget(self):
             assert_equal(len(res), 200)
             assert_equal(res, [(str(i), str(i)) for i in range(100, 300)])
 
+    def test_xget_counter(self):
+        if not have_counters:
+            raise SkipTest('Cassandra 0.7 does not support counters')
+
+        key = 'test_xget_counter'
+        counter_cf.insert(key, {'col1': 1})
+        res = list(counter_cf.xget(key))
+        assert_equal(res, [('col1', 1)])
+
+        counter_cf.insert(key, {'col1': 1, 'col2': 1})
+        res = list(counter_cf.xget(key))
+        assert_equal(res, [('col1', 2), ('col2', 1)])
+
 class TestSuperColumnFamily(unittest.TestCase):
 
     def tearDown(self):
@@ -632,3 +645,16 @@ def test_remove_counter(self):
 
         counter_scf.remove_counter(key, 'col', super_column='scol')
         assert_raises(NotFoundException, scf.get, key)
+
+    def test_xget_counter(self):
+        if not have_counters:
+            raise SkipTest('Cassandra 0.7 does not support counters')
+
+        key = 'test_xget_counter'
+        counter_scf.insert(key, {'scol': {'col1': 1}})
+        res = list(counter_scf.xget(key))
+        assert_equal(res, [('scol', {'col1': 1})])
+
+        counter_scf.insert(key, {'scol': {'col1': 1, 'col2': 1}})
+        res = list(counter_scf.xget(key))
+        assert_equal(res, [('scol', {'col1': 2, 'col2': 1})])