diff --git a/gcloud/bigtable/happybase/table.py b/gcloud/bigtable/happybase/table.py index c650245bb1d8..b9e377a448bc 100644 --- a/gcloud/bigtable/happybase/table.py +++ b/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 '' % (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 `, + :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 diff --git a/gcloud/bigtable/happybase/test_table.py b/gcloud/bigtable/happybase/test_table.py index 1343b7fb20f0..395b68501b62 100644 --- a/gcloud/bigtable/happybase/test_table.py +++ b/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