Merge pull request #1511 from dhermes/happybase-table-families

Implementing HappyBase Table.families().

Author: dhermes

gcloud/bigtable/happybase/table.py

@@ -15,9 +15,18 @@
 """Google Cloud Bigtable HappyBase table module."""
 
 
+import six
+
+from gcloud._helpers import _total_seconds
+from gcloud.bigtable.column_family import GCRuleIntersection
+from gcloud.bigtable.column_family import MaxAgeGCRule
+from gcloud.bigtable.column_family import MaxVersionsGCRule
 from gcloud.bigtable.table import Table as _LowLevelTable
 
 
+_SIMPLE_GC_RULES = (MaxAgeGCRule, MaxVersionsGCRule)
+
+
 def make_row(cell_map, include_timestamp):
     """Make a row dict for a Thrift cell mapping.
 
@@ -92,6 +101,19 @@ def __init__(self, name, connection):
     def __repr__(self):
         return '<table.Table name=%r>' % (self.name,)
 
+    def families(self):
+        """Retrieve the column families for this table.
+
+        :rtype: dict
+        :returns: Mapping from column family name to garbage collection rule
+                  for a column family.
+        """
+        column_family_map = self._low_level_table.list_column_families()
+        result = {}
+        for col_fam, col_fam_obj in six.iteritems(column_family_map):
+            result[col_fam] = _gc_rule_to_dict(col_fam_obj.gc_rule)
+        return result
+
     def regions(self):
         """Retrieve the regions for this table.
 
@@ -104,3 +126,48 @@ def regions(self):
         """
         raise NotImplementedError('The Cloud Bigtable API does not have a '
                                   'concept of splitting a table into regions.')
+
+
+def _gc_rule_to_dict(gc_rule):
+    """Converts garbage collection rule to dictionary if possible.
+
+    This is in place to support dictionary values as was done
+    in HappyBase, which has somewhat different garbage collection rule
+    settings for column families.
+
+    Only does this if the garbage collection rule is:
+
+    * :class:`.MaxAgeGCRule`
+    * :class:`.MaxVersionsGCRule`
+    * Composite :class:`.GCRuleIntersection` with two rules, one each
+      of type :class:`.MaxAgeGCRule` and :class:`.MaxVersionsGCRule`
+
+    Otherwise, just returns the input without change.
+
+    :type gc_rule: :data:`NoneType <types.NoneType>`,
+                   :class:`.GarbageCollectionRule`
+    :param gc_rule: A garbage collection rule to convert to a dictionary
+                    (if possible).
+
+    :rtype: dict or :class:`.GarbageCollectionRule`
+    :returns: The converted garbage collection rule.
+    """
+    result = gc_rule
+    if gc_rule is None:
+        result = {}
+    elif isinstance(gc_rule, MaxAgeGCRule):
+        result = {'time_to_live': _total_seconds(gc_rule.max_age)}
+    elif isinstance(gc_rule, MaxVersionsGCRule):
+        result = {'max_versions': gc_rule.max_num_versions}
+    elif isinstance(gc_rule, GCRuleIntersection):
+        if len(gc_rule.rules) == 2:
+            rule1, rule2 = gc_rule.rules
+            if (isinstance(rule1, _SIMPLE_GC_RULES) and
+                    isinstance(rule2, _SIMPLE_GC_RULES)):
+                rule1 = _gc_rule_to_dict(rule1)
+                rule2 = _gc_rule_to_dict(rule2)
+                key1, = rule1.keys()
+                key2, = rule2.keys()
+                if key1 != key2:
+                    result = {key1: rule1[key1], key2: rule2[key2]}
+    return result

gcloud/bigtable/happybase/test_table.py

@@ -78,6 +78,36 @@ def test_constructor_null_connection(self):
         self.assertEqual(table.connection, connection)
         self.assertEqual(table._low_level_table, None)
 
+    def test_families(self):
+        from gcloud._testing import _Monkey
+        from gcloud.bigtable.happybase import table as MUT
+
+        name = 'table-name'
+        connection = None
+        table = self._makeOne(name, connection)
+        table._low_level_table = _MockLowLevelTable()
+
+        # Mock the column families to be returned.
+        col_fam_name = 'fam'
+        gc_rule = object()
+        col_fam = _MockLowLevelColumnFamily(col_fam_name, gc_rule=gc_rule)
+        col_fams = {col_fam_name: col_fam}
+        table._low_level_table.column_families = col_fams
+
+        to_dict_result = object()
+        to_dict_calls = []
+
+        def mock_gc_rule_to_dict(gc_rule):
+            to_dict_calls.append(gc_rule)
+            return to_dict_result
+
+        with _Monkey(MUT, _gc_rule_to_dict=mock_gc_rule_to_dict):
+            result = table.families()
+
+        self.assertEqual(result, {col_fam_name: to_dict_result})
+        self.assertEqual(table._low_level_table.list_column_families_calls, 1)
+        self.assertEqual(to_dict_calls, [gc_rule])
+
     def test___repr__(self):
         name = 'table-name'
         table = self._makeOne(name, None)
@@ -92,14 +122,116 @@ def test_regions(self):
             table.regions()
 
 
+class Test__gc_rule_to_dict(unittest2.TestCase):
+
+    def _callFUT(self, *args, **kwargs):
+        from gcloud.bigtable.happybase.table import _gc_rule_to_dict
+        return _gc_rule_to_dict(*args, **kwargs)
+
+    def test_with_null(self):
+        gc_rule = None
+        result = self._callFUT(gc_rule)
+        self.assertEqual(result, {})
+
+    def test_with_max_versions(self):
+        from gcloud.bigtable.column_family import MaxVersionsGCRule
+
+        max_versions = 2
+        gc_rule = MaxVersionsGCRule(max_versions)
+        result = self._callFUT(gc_rule)
+        expected_result = {'max_versions': max_versions}
+        self.assertEqual(result, expected_result)
+
+    def test_with_max_age(self):
+        import datetime
+        from gcloud.bigtable.column_family import MaxAgeGCRule
+
+        time_to_live = 101
+        max_age = datetime.timedelta(seconds=time_to_live)
+        gc_rule = MaxAgeGCRule(max_age)
+        result = self._callFUT(gc_rule)
+        expected_result = {'time_to_live': time_to_live}
+        self.assertEqual(result, expected_result)
+
+    def test_with_non_gc_rule(self):
+        gc_rule = object()
+        result = self._callFUT(gc_rule)
+        self.assertTrue(result is gc_rule)
+
+    def test_with_gc_rule_union(self):
+        from gcloud.bigtable.column_family import GCRuleUnion
+
+        gc_rule = GCRuleUnion(rules=[])
+        result = self._callFUT(gc_rule)
+        self.assertTrue(result is gc_rule)
+
+    def test_with_intersection_other_than_two(self):
+        from gcloud.bigtable.column_family import GCRuleIntersection
+
+        gc_rule = GCRuleIntersection(rules=[])
+        result = self._callFUT(gc_rule)
+        self.assertTrue(result is gc_rule)
+
+    def test_with_intersection_two_max_num_versions(self):
+        from gcloud.bigtable.column_family import GCRuleIntersection
+        from gcloud.bigtable.column_family import MaxVersionsGCRule
+
+        rule1 = MaxVersionsGCRule(1)
+        rule2 = MaxVersionsGCRule(2)
+        gc_rule = GCRuleIntersection(rules=[rule1, rule2])
+        result = self._callFUT(gc_rule)
+        self.assertTrue(result is gc_rule)
+
+    def test_with_intersection_two_rules(self):
+        import datetime
+        from gcloud.bigtable.column_family import GCRuleIntersection
+        from gcloud.bigtable.column_family import MaxAgeGCRule
+        from gcloud.bigtable.column_family import MaxVersionsGCRule
+
+        time_to_live = 101
+        max_age = datetime.timedelta(seconds=time_to_live)
+        rule1 = MaxAgeGCRule(max_age)
+        max_versions = 2
+        rule2 = MaxVersionsGCRule(max_versions)
+        gc_rule = GCRuleIntersection(rules=[rule1, rule2])
+        result = self._callFUT(gc_rule)
+        expected_result = {
+            'max_versions': max_versions,
+            'time_to_live': time_to_live,
+        }
+        self.assertEqual(result, expected_result)
+
+    def test_with_intersection_two_nested_rules(self):
+        from gcloud.bigtable.column_family import GCRuleIntersection
+
+        rule1 = GCRuleIntersection(rules=[])
+        rule2 = GCRuleIntersection(rules=[])
+        gc_rule = GCRuleIntersection(rules=[rule1, rule2])
+        result = self._callFUT(gc_rule)
+        self.assertTrue(result is gc_rule)
+
+
 class _Connection(object):
 
     def __init__(self, cluster):
         self._cluster = cluster
 
 
+class _MockLowLevelColumnFamily(object):
+
+    def __init__(self, column_family_id, gc_rule=None):
+        self.column_family_id = column_family_id
+        self.gc_rule = gc_rule
+
+
 class _MockLowLevelTable(object):
 
     def __init__(self, *args, **kwargs):
         self.args = args
         self.kwargs = kwargs
+        self.list_column_families_calls = 0
+        self.column_families = {}
+
+    def list_column_families(self):
+        self.list_column_families_calls += 1
+        return self.column_families