def test_creating_table_that_exists_with_different_column_families_gc_rule_in__table(
self, mock_hook):
op = BigtableCreateTableOperator(
project_id=PROJECT_ID,
instance_id=INSTANCE_ID,
table_id=TABLE_ID,
initial_split_keys=INITIAL_SPLIT_KEYS,
column_families={"cf-id": MaxVersionsGCRule(1)},
task_id="id",
gcp_conn_id=GCP_CONN_ID,
impersonation_chain=IMPERSONATION_CHAIN,
)
cf_mock = mock.Mock()
cf_mock.gc_rule = mock.Mock(return_value=MaxVersionsGCRule(2))
mock_hook.return_value.get_column_families_for_table.return_value = {
"cf-id": cf_mock
}
mock_hook.return_value.create_table.side_effect = mock.Mock(
side_effect=google.api_core.exceptions.AlreadyExists(
"Table already exists."))
with self.assertRaises(AirflowException) as e:
op.execute(None)
err = e.exception
self.assertEqual(
str(err),
"Table '{}' already exists with different Column Families.".format(
TABLE_ID))
mock_hook.assert_called_once_with(
gcp_conn_id=GCP_CONN_ID,
impersonation_chain=IMPERSONATION_CHAIN,
)
def test_to_pb_nested(self):
import datetime
from google.protobuf import duration_pb2
from google.cloud.bigtable.column_family import MaxAgeGCRule
from google.cloud.bigtable.column_family import MaxVersionsGCRule
max_num_versions1 = 42
rule1 = MaxVersionsGCRule(max_num_versions1)
pb_rule1 = _GcRulePB(max_num_versions=max_num_versions1)
max_age = datetime.timedelta(seconds=1)
rule2 = MaxAgeGCRule(max_age)
pb_rule2 = _GcRulePB(max_age=duration_pb2.Duration(seconds=1))
rule3 = self._make_one(rules=[rule1, rule2])
pb_rule3 = _GcRulePB(intersection=_GcRuleIntersectionPB(
rules=[pb_rule1, pb_rule2]))
max_num_versions2 = 1337
rule4 = MaxVersionsGCRule(max_num_versions2)
pb_rule4 = _GcRulePB(max_num_versions=max_num_versions2)
rule5 = self._make_one(rules=[rule3, rule4])
pb_rule5 = _GcRulePB(intersection=_GcRuleIntersectionPB(
rules=[pb_rule3, pb_rule4]))
gc_rule_pb = rule5.to_pb()
self.assertEqual(gc_rule_pb, pb_rule5)
def test_max_num_versions(self):
from google.cloud.bigtable.column_family import MaxVersionsGCRule
orig_rule = MaxVersionsGCRule(1)
gc_rule_pb = orig_rule.to_pb()
result = self._callFUT(gc_rule_pb)
self.assertTrue(isinstance(result, MaxVersionsGCRule))
self.assertEqual(result, orig_rule)
def test_max_num_versions(self):
from google.cloud.bigtable.column_family import MaxVersionsGCRule
orig_rule = MaxVersionsGCRule(1)
gc_rule_pb = orig_rule.to_pb()
result = self._call_fut(gc_rule_pb)
self.assertIsInstance(result, MaxVersionsGCRule)
self.assertEqual(result, orig_rule)
def test_to_pb_with_rule(self):
from google.cloud.bigtable.column_family import MaxVersionsGCRule
gc_rule = MaxVersionsGCRule(1)
column_family = self._make_one("column_family_id", None, gc_rule=gc_rule)
pb_val = column_family.to_pb()
expected = _ColumnFamilyPB(gc_rule=gc_rule.to_pb())
self.assertEqual(pb_val, expected)
def test_with_intersection_two_max_num_versions(self):
from google.cloud.bigtable.column_family import GCRuleIntersection
from google.cloud.bigtable.column_family import MaxVersionsGCRule
rule1 = MaxVersionsGCRule(1)
rule2 = MaxVersionsGCRule(2)
gc_rule = GCRuleIntersection(rules=[rule1, rule2])
result = self._call_fut(gc_rule)
self.assertTrue(result is gc_rule)
def test__gc_rule_from_pb_max_num_versions():
from google.cloud.bigtable.column_family import _gc_rule_from_pb
from google.cloud.bigtable.column_family import MaxVersionsGCRule
orig_rule = MaxVersionsGCRule(1)
gc_rule_pb = orig_rule.to_pb()
result = _gc_rule_from_pb(gc_rule_pb)
assert isinstance(result, MaxVersionsGCRule)
assert result == orig_rule
def test_column_family_to_pb_with_rule():
from google.cloud.bigtable.column_family import MaxVersionsGCRule
gc_rule = MaxVersionsGCRule(1)
column_family = _make_column_family("column_family_id",
None,
gc_rule=gc_rule)
pb_val = column_family.to_pb()
expected = _ColumnFamilyPB(gc_rule=gc_rule.to_pb())
assert pb_val == expected
def test_with_max_versions(self):
from google.cloud.bigtable.column_family import MaxVersionsGCRule
max_versions = 2
gc_rule = MaxVersionsGCRule(max_versions)
result = self._call_fut(gc_rule)
expected_result = {"max_versions": max_versions}
self.assertEqual(result, expected_result)
def test_dictionary_versions_key(self):
from google.cloud.bigtable.column_family import MaxVersionsGCRule
versions = 42
option = {"max_versions": versions}
result = self._call_fut(option)
gc_rule = MaxVersionsGCRule(versions)
self.assertEqual(result, gc_rule)
def test_intersection(self):
import datetime
from google.cloud.bigtable.column_family import GCRuleIntersection
from google.cloud.bigtable.column_family import MaxAgeGCRule
from google.cloud.bigtable.column_family import MaxVersionsGCRule
rule1 = MaxVersionsGCRule(1)
rule2 = MaxAgeGCRule(datetime.timedelta(seconds=1))
orig_rule = GCRuleIntersection([rule1, rule2])
gc_rule_pb = orig_rule.to_pb()
result = self._call_fut(gc_rule_pb)
self.assertIsInstance(result, GCRuleIntersection)
self.assertEqual(result, orig_rule)
def test_create_with_column_families(self):
from google.cloud.bigtable.column_family import ColumnFamily
from google.cloud.bigtable.column_family import MaxVersionsGCRule
cf_id1 = 'col-fam-id1'
cf1 = ColumnFamily(cf_id1, None)
cf_id2 = 'col-fam-id2'
gc_rule = MaxVersionsGCRule(42)
cf2 = ColumnFamily(cf_id2, None, gc_rule=gc_rule)
initial_split_keys = None
column_families = [cf1, cf2]
self._create_test_helper(initial_split_keys,
column_families=column_families)
def test__gc_rule_from_pb_intersection():
import datetime
from google.cloud.bigtable.column_family import _gc_rule_from_pb
from google.cloud.bigtable.column_family import GCRuleIntersection
from google.cloud.bigtable.column_family import MaxAgeGCRule
from google.cloud.bigtable.column_family import MaxVersionsGCRule
rule1 = MaxVersionsGCRule(1)
rule2 = MaxAgeGCRule(datetime.timedelta(seconds=1))
orig_rule = GCRuleIntersection([rule1, rule2])
gc_rule_pb = orig_rule.to_pb()
result = _gc_rule_from_pb(gc_rule_pb)
assert isinstance(result, GCRuleIntersection)
assert result == orig_rule
def test_create_table_with_families(self):
temp_table_id = "test-create-table-with-failies"
temp_table = Config.INSTANCE.table(temp_table_id)
gc_rule = MaxVersionsGCRule(1)
temp_table.create(column_families={COLUMN_FAMILY_ID1: gc_rule})
self.tables_to_delete.append(temp_table)
col_fams = temp_table.list_column_families()
self.assertEqual(len(col_fams), 1)
retrieved_col_fam = col_fams[COLUMN_FAMILY_ID1]
self.assertIs(retrieved_col_fam._table, temp_table)
self.assertEqual(retrieved_col_fam.column_family_id, COLUMN_FAMILY_ID1)
self.assertEqual(retrieved_col_fam.gc_rule, gc_rule)
def _parse_family_option(option):
"""Parses a column family option into a garbage collection rule.
.. note::
If ``option`` is not a dictionary, the type is not checked.
If ``option`` is :data:`None`, there is nothing to do, since this
is the correct output.
:type option: :class:`dict`,
:data:`NoneType <types.NoneType>`,
:class:`.GarbageCollectionRule`
:param option: A column family option passes as a dictionary value in
:meth:`Connection.create_table`.
:rtype: :class:`.GarbageCollectionRule`
:returns: A garbage collection rule parsed from the input.
"""
result = option
if isinstance(result, dict):
if not set(result.keys()) <= set(['max_versions', 'time_to_live']):
all_keys = ', '.join(repr(key) for key in result.keys())
warning_msg = ('Cloud Bigtable only supports max_versions and '
'time_to_live column family settings. '
'Received: %s' % (all_keys, ))
_WARN(warning_msg)
max_num_versions = result.get('max_versions')
max_age = None
if 'time_to_live' in result:
max_age = datetime.timedelta(seconds=result['time_to_live'])
versions_rule = age_rule = None
if max_num_versions is not None:
versions_rule = MaxVersionsGCRule(max_num_versions)
if max_age is not None:
age_rule = MaxAgeGCRule(max_age)
if versions_rule is None:
result = age_rule
else:
if age_rule is None:
result = versions_rule
else:
result = GCRuleIntersection(rules=[age_rule, versions_rule])
return result
def test_dictionary_both_keys(self):
import datetime
from google.cloud.bigtable.column_family import GCRuleIntersection
from google.cloud.bigtable.column_family import MaxAgeGCRule
from google.cloud.bigtable.column_family import MaxVersionsGCRule
versions = 42
time_to_live = 24 * 60 * 60
option = {"max_versions": versions, "time_to_live": time_to_live}
result = self._call_fut(option)
max_age = datetime.timedelta(days=1)
# NOTE: This relies on the order of the rules in the method we are
# calling matching this order here.
gc_rule1 = MaxAgeGCRule(max_age)
gc_rule2 = MaxVersionsGCRule(versions)
gc_rule = GCRuleIntersection(rules=[gc_rule1, gc_rule2])
self.assertEqual(result, gc_rule)
def test_with_intersection_two_rules(self):
import datetime
from google.cloud.bigtable.column_family import GCRuleIntersection
from google.cloud.bigtable.column_family import MaxAgeGCRule
from google.cloud.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._call_fut(gc_rule)
expected_result = {
"max_versions": max_versions,
"time_to_live": time_to_live
}
self.assertEqual(result, expected_result)
def test_to_pb(self):
import datetime
from google.protobuf import duration_pb2
from google.cloud.bigtable.column_family import MaxAgeGCRule
from google.cloud.bigtable.column_family import MaxVersionsGCRule
max_num_versions = 42
rule1 = MaxVersionsGCRule(max_num_versions)
pb_rule1 = _GcRulePB(max_num_versions=max_num_versions)
max_age = datetime.timedelta(seconds=1)
rule2 = MaxAgeGCRule(max_age)
pb_rule2 = _GcRulePB(max_age=duration_pb2.Duration(seconds=1))
rule3 = self._make_one(rules=[rule1, rule2])
pb_rule3 = _GcRulePB(union=_GcRuleUnionPB(rules=[pb_rule1, pb_rule2]))
gc_rule_pb = rule3.to_pb()
self.assertEqual(gc_rule_pb, pb_rule3)
def setup_column_family(self, table_name, column_family, silent=False):
if not self.admin or self.read_only:
raise RuntimeError("admin operations not allowed")
ad = self.get_admin_connection()
full_table_name = self.get_full_table_name(table_name)
t = ad.table(full_table_name)
families_before = t.list_column_families()
logger.debug(
"CREATE CF: Existing Families: {}".format(families_before))
if silent and column_family in families_before:
logger.warning("CREATE CF: Ignoring existing family: {}".format(
column_family))
return
cf1 = t.column_family(column_family, gc_rule=MaxVersionsGCRule(1))
cf1.create()
time.sleep(0.25)
logger.warning("CREATE CF: Created Family: {}".format(column_family))
return t
def test_create_column_family(self):
temp_table_id = "test-create-column-family"
temp_table = Config.INSTANCE.table(temp_table_id)
temp_table.create()
self.tables_to_delete.append(temp_table)
self.assertEqual(temp_table.list_column_families(), {})
gc_rule = MaxVersionsGCRule(1)
column_family = temp_table.column_family(COLUMN_FAMILY_ID1,
gc_rule=gc_rule)
column_family.create()
col_fams = temp_table.list_column_families()
self.assertEqual(len(col_fams), 1)
retrieved_col_fam = col_fams[COLUMN_FAMILY_ID1]
self.assertIs(retrieved_col_fam._table, column_family._table)
self.assertEqual(retrieved_col_fam.column_family_id,
column_family.column_family_id)
self.assertEqual(retrieved_col_fam.gc_rule, gc_rule)
def _check_and_create_table(self):
""" Checks if table exists and creates new one if necessary """
table_ids = [t.table_id for t in self.instance.list_tables()]
if not self.table_id in table_ids:
self.table.create()
f_data = self.table.column_family(self.data_family_id)
f_data.create()
f_inc = self.table.column_family(self.incrementer_family_id,
gc_rule=MaxVersionsGCRule(1))
f_inc.create()
f_map = self.table.column_family(self.mapping_family_id)
f_map.create()
f_log = self.table.column_family(self.log_family_id)
f_log.create()
print("Table created")
def test_gc_rule_intersection_to_pb():
import datetime
from google.protobuf import duration_pb2
from google.cloud.bigtable.column_family import MaxAgeGCRule
from google.cloud.bigtable.column_family import MaxVersionsGCRule
max_num_versions = 42
rule1 = MaxVersionsGCRule(max_num_versions)
pb_rule1 = _GcRulePB(max_num_versions=max_num_versions)
max_age = datetime.timedelta(seconds=1)
rule2 = MaxAgeGCRule(max_age)
pb_rule2 = _GcRulePB(max_age=duration_pb2.Duration(seconds=1))
rule3 = _make_gc_rule_intersection(rules=[rule1, rule2])
pb_rule3 = _GcRulePB(intersection=_GcRuleIntersectionPB(
rules=[pb_rule1, pb_rule2]))
gc_rule_pb = rule3.to_pb()
assert gc_rule_pb == pb_rule3
def test_table_create_w_families(
data_instance_populated,
tables_to_delete,
):
from google.cloud.bigtable.column_family import MaxVersionsGCRule
temp_table_id = "test-create-table-with-failies"
column_family_id = "col-fam-id1"
temp_table = data_instance_populated.table(temp_table_id)
gc_rule = MaxVersionsGCRule(1)
temp_table.create(column_families={column_family_id: gc_rule})
tables_to_delete.append(temp_table)
col_fams = temp_table.list_column_families()
assert len(col_fams) == 1
retrieved_col_fam = col_fams[column_family_id]
assert retrieved_col_fam._table is temp_table
assert retrieved_col_fam.column_family_id == column_family_id
assert retrieved_col_fam.gc_rule == gc_rule
def test_update_column_family(self):
temp_table_id = "test-update-column-family"
temp_table = Config.INSTANCE_DATA.table(temp_table_id)
temp_table.create()
self.tables_to_delete.append(temp_table)
gc_rule = MaxVersionsGCRule(1)
column_family = temp_table.column_family(COLUMN_FAMILY_ID1, gc_rule=gc_rule)
column_family.create()
# Check that our created table is as expected.
col_fams = temp_table.list_column_families()
self.assertEqual(col_fams, {COLUMN_FAMILY_ID1: column_family})
# Update the column family's GC rule and then try to update.
column_family.gc_rule = None
column_family.update()
# Check that the update has propagated.
col_fams = temp_table.list_column_families()
self.assertIsNone(col_fams[COLUMN_FAMILY_ID1].gc_rule)
def test_column_family_create(data_instance_populated, tables_to_delete):
from google.cloud.bigtable.column_family import MaxVersionsGCRule
temp_table_id = "test-create-column-family"
temp_table = data_instance_populated.table(temp_table_id)
temp_table.create()
tables_to_delete.append(temp_table)
assert temp_table.list_column_families() == {}
column_family_id = "col-fam-id1"
gc_rule = MaxVersionsGCRule(1)
column_family = temp_table.column_family(column_family_id, gc_rule=gc_rule)
column_family.create()
col_fams = temp_table.list_column_families()
assert len(col_fams) == 1
retrieved_col_fam = col_fams[column_family_id]
assert retrieved_col_fam._table is column_family._table
assert retrieved_col_fam.column_family_id == column_family.column_family_id
assert retrieved_col_fam.gc_rule == gc_rule
def test_column_family_update(data_instance_populated, tables_to_delete):
from google.cloud.bigtable.column_family import MaxVersionsGCRule
temp_table_id = "test-update-column-family"
temp_table = data_instance_populated.table(temp_table_id)
temp_table.create()
tables_to_delete.append(temp_table)
column_family_id = "col-fam-id1"
gc_rule = MaxVersionsGCRule(1)
column_family = temp_table.column_family(column_family_id, gc_rule=gc_rule)
column_family.create()
# Check that our created table is as expected.
col_fams = temp_table.list_column_families()
assert col_fams == {column_family_id: column_family}
# Update the column family's GC rule and then try to update.
column_family.gc_rule = None
column_family.update()
# Check that the update has propagated.
col_fams = temp_table.list_column_families()
assert col_fams[column_family_id].gc_rule is None
def test_column_family_create_with_gc_rule():
from google.cloud.bigtable.column_family import MaxVersionsGCRule
gc_rule = MaxVersionsGCRule(1337)
_create_test_helper(gc_rule=gc_rule)
def _make_max_versions_gc_rule(*args, **kwargs):
from google.cloud.bigtable.column_family import MaxVersionsGCRule
return MaxVersionsGCRule(*args, **kwargs)
def test_update_with_gc_rule(self):
from google.cloud.bigtable.column_family import MaxVersionsGCRule
gc_rule = MaxVersionsGCRule(1337)
self._update_test_helper(gc_rule=gc_rule)