def poke(self, context):
instance = self.hook.get_instance(project_id=self.project_id,
instance_id=self.instance_id)
if not instance:
self.log.info("Dependency: instance '%s' does not exist.",
self.instance_id)
return False
try:
cluster_states = self.hook.get_cluster_states_for_table(
instance=instance, table_id=self.table_id)
except google.api_core.exceptions.NotFound:
self.log.info(
"Dependency: table '%s' does not exist in instance '%s'.",
self.table_id, self.instance_id)
return False
ready_state = ClusterState(
enums.Table.ClusterState.ReplicationState.READY)
is_table_replicated = True
for cluster_id in cluster_states.keys():
if cluster_states[cluster_id] != ready_state:
self.log.info(
"Table '%s' is not yet replicated on cluster '%s'.",
self.table_id, cluster_id)
is_table_replicated = False
if not is_table_replicated:
return False
self.log.info("Table '%s' is replicated.", self.table_id)
return True
def test_wait_ready(self, mock_hook):
mock_hook.return_value.get_instance.return_value = mock.Mock(Instance)
mock_hook.return_value.get_cluster_states_for_table.return_value = {
"cl-id": ClusterState(4)
}
op = BigtableTableWaitForReplicationSensor(project_id=PROJECT_ID,
instance_id=INSTANCE_ID,
table_id=TABLE_ID,
task_id="id")
self.assertTrue(op.poke(None))
mock_hook.assert_called_once_with()
def test_wait_not_ready(self, mock_hook):
mock_hook.return_value.get_instance.return_value = mock.Mock(Instance)
mock_hook.return_value.get_cluster_states_for_table.return_value = {
"cl-id": ClusterState(0)
}
op = BigtableTableReplicationCompletedSensor(project_id=PROJECT_ID,
instance_id=INSTANCE_ID,
table_id=TABLE_ID,
task_id="id",
gcp_conn_id=GCP_CONN_ID)
self.assertFalse(op.poke(None))
mock_hook.assert_called_once_with(gcp_conn_id=GCP_CONN_ID)
def test_wait_ready(self, mock_hook):
mock_hook.return_value.get_instance.return_value = mock.Mock(Instance)
mock_hook.return_value.get_cluster_states_for_table.return_value = {
"cl-id": ClusterState(4)
}
op = BigtableTableReplicationCompletedSensor(
project_id=PROJECT_ID,
instance_id=INSTANCE_ID,
table_id=TABLE_ID,
task_id="id",
gcp_conn_id=GCP_CONN_ID,
impersonation_chain=IMPERSONATION_CHAIN,
)
assert op.poke(None)
mock_hook.assert_called_once_with(
gcp_conn_id=GCP_CONN_ID,
impersonation_chain=IMPERSONATION_CHAIN,
)
def poke(self, context: 'Context') -> bool:
hook = BigtableHook(
gcp_conn_id=self.gcp_conn_id,
impersonation_chain=self.impersonation_chain,
)
instance = hook.get_instance(project_id=self.project_id,
instance_id=self.instance_id)
if not instance:
self.log.info("Dependency: instance '%s' does not exist.",
self.instance_id)
return False
try:
cluster_states = hook.get_cluster_states_for_table(
instance=instance, table_id=self.table_id)
except google.api_core.exceptions.NotFound:
self.log.info(
"Dependency: table '%s' does not exist in instance '%s'.",
self.table_id, self.instance_id)
return False
ready_state = ClusterState(
enums.Table.ClusterState.ReplicationState.READY)
is_table_replicated = True
for cluster_id in cluster_states.keys():
if cluster_states[cluster_id] != ready_state:
self.log.info(
"Table '%s' is not yet replicated on cluster '%s'.",
self.table_id, cluster_id)
is_table_replicated = False
if not is_table_replicated:
return False
self.log.info("Table '%s' is replicated.", self.table_id)
BigtableTablesLink.persist(context=context, task_instance=self)
return True
def test_create_instance_w_two_clusters(self):
from google.cloud.bigtable import enums
from google.cloud.bigtable.table import ClusterState
_PRODUCTION = enums.Instance.Type.PRODUCTION
ALT_INSTANCE_ID = "dif" + unique_resource_id("-")
instance = Config.CLIENT.instance(ALT_INSTANCE_ID,
instance_type=_PRODUCTION,
labels=LABELS)
ALT_CLUSTER_ID_1 = ALT_INSTANCE_ID + "-c1"
ALT_CLUSTER_ID_2 = ALT_INSTANCE_ID + "-c2"
LOCATION_ID_2 = "us-central1-f"
STORAGE_TYPE = enums.StorageType.HDD
cluster_1 = instance.cluster(
ALT_CLUSTER_ID_1,
location_id=LOCATION_ID,
serve_nodes=SERVE_NODES,
default_storage_type=STORAGE_TYPE,
)
cluster_2 = instance.cluster(
ALT_CLUSTER_ID_2,
location_id=LOCATION_ID_2,
serve_nodes=SERVE_NODES,
default_storage_type=STORAGE_TYPE,
)
operation = instance.create(clusters=[cluster_1, cluster_2])
# We want to make sure the operation completes.
operation.result(timeout=10)
# Make sure this instance gets deleted after the test case.
self.instances_to_delete.append(instance)
# Create a new instance instance and make sure it is the same.
instance_alt = Config.CLIENT.instance(ALT_INSTANCE_ID)
instance_alt.reload()
self.assertEqual(instance, instance_alt)
self.assertEqual(instance.display_name, instance_alt.display_name)
self.assertEqual(instance.type_, instance_alt.type_)
clusters, failed_locations = instance_alt.list_clusters()
self.assertEqual(failed_locations, [])
clusters.sort(key=lambda x: x.name)
alt_cluster_1, alt_cluster_2 = clusters
self.assertEqual(cluster_1.location_id, alt_cluster_1.location_id)
self.assertEqual(alt_cluster_1.state, enums.Cluster.State.READY)
self.assertEqual(cluster_1.serve_nodes, alt_cluster_1.serve_nodes)
self.assertEqual(cluster_1.default_storage_type,
alt_cluster_1.default_storage_type)
self.assertEqual(cluster_2.location_id, alt_cluster_2.location_id)
self.assertEqual(alt_cluster_2.state, enums.Cluster.State.READY)
self.assertEqual(cluster_2.serve_nodes, alt_cluster_2.serve_nodes)
self.assertEqual(cluster_2.default_storage_type,
alt_cluster_2.default_storage_type)
# Test list clusters in project via 'client.list_clusters'
clusters, failed_locations = Config.CLIENT.list_clusters()
self.assertFalse(failed_locations)
found = set([cluster.name for cluster in clusters])
self.assertTrue(
{alt_cluster_1.name, alt_cluster_2.name,
Config.CLUSTER.name}.issubset(found))
temp_table_id = "test-get-cluster-states"
temp_table = instance.table(temp_table_id)
temp_table.create()
result = temp_table.get_cluster_states()
ReplicationState = enums.Table.ReplicationState
expected_results = [
ClusterState(ReplicationState.STATE_NOT_KNOWN),
ClusterState(ReplicationState.INITIALIZING),
ClusterState(ReplicationState.PLANNED_MAINTENANCE),
ClusterState(ReplicationState.UNPLANNED_MAINTENANCE),
ClusterState(ReplicationState.READY),
]
cluster_id_list = result.keys()
self.assertEqual(len(cluster_id_list), 2)
self.assertIn(ALT_CLUSTER_ID_1, cluster_id_list)
self.assertIn(ALT_CLUSTER_ID_2, cluster_id_list)
for clusterstate in result.values():
self.assertIn(clusterstate, expected_results)
# Test create app profile with multi_cluster_routing policy
app_profiles_to_delete = []
description = "routing policy-multy"
app_profile_id_1 = "app_profile_id_1"
routing = enums.RoutingPolicyType.ANY
self._test_create_app_profile_helper(
app_profile_id_1,
instance,
routing_policy_type=routing,
description=description,
ignore_warnings=True,
)
app_profiles_to_delete.append(app_profile_id_1)
# Test list app profiles
self._test_list_app_profiles_helper(instance, [app_profile_id_1])
# Test modify app profile app_profile_id_1
# routing policy to single cluster policy,
# cluster -> ALT_CLUSTER_ID_1,
# allow_transactional_writes -> disallowed
# modify description
description = "to routing policy-single"
routing = enums.RoutingPolicyType.SINGLE
self._test_modify_app_profile_helper(
app_profile_id_1,
instance,
routing_policy_type=routing,
description=description,
cluster_id=ALT_CLUSTER_ID_1,
allow_transactional_writes=False,
)
# Test modify app profile app_profile_id_1
# cluster -> ALT_CLUSTER_ID_2,
# allow_transactional_writes -> allowed
self._test_modify_app_profile_helper(
app_profile_id_1,
instance,
routing_policy_type=routing,
description=description,
cluster_id=ALT_CLUSTER_ID_2,
allow_transactional_writes=True,
ignore_warnings=True,
)
# Test create app profile with single cluster routing policy
description = "routing policy-single"
app_profile_id_2 = "app_profile_id_2"
routing = enums.RoutingPolicyType.SINGLE
self._test_create_app_profile_helper(
app_profile_id_2,
instance,
routing_policy_type=routing,
description=description,
cluster_id=ALT_CLUSTER_ID_2,
allow_transactional_writes=False,
)
app_profiles_to_delete.append(app_profile_id_2)
# Test list app profiles
self._test_list_app_profiles_helper(
instance, [app_profile_id_1, app_profile_id_2])
# Test modify app profile app_profile_id_2 to
# allow transactional writes
# Note: no need to set ``ignore_warnings`` to True
# since we are not restrictings anything with this modification.
self._test_modify_app_profile_helper(
app_profile_id_2,
instance,
routing_policy_type=routing,
description=description,
cluster_id=ALT_CLUSTER_ID_2,
allow_transactional_writes=True,
)
# Test modify app profile app_profile_id_2 routing policy
# to multi_cluster_routing policy
# modify description
description = "to routing policy-multy"
routing = enums.RoutingPolicyType.ANY
self._test_modify_app_profile_helper(
app_profile_id_2,
instance,
routing_policy_type=routing,
description=description,
allow_transactional_writes=False,
ignore_warnings=True,
)
# Test delete app profiles
for app_profile_id in app_profiles_to_delete:
self._test_delete_app_profile_helper(app_profile_id, instance)
def test_instance_create_w_two_clusters(
admin_client,
unique_suffix,
admin_instance_populated,
admin_cluster,
location_id,
instance_labels,
instances_to_delete,
skip_on_emulator,
):
alt_instance_id = f"dif{unique_suffix}"
instance = admin_client.instance(
alt_instance_id,
instance_type=enums.Instance.Type.PRODUCTION,
labels=instance_labels,
)
serve_nodes = 1
alt_cluster_id_1 = f"{alt_instance_id}-c1"
cluster_1 = instance.cluster(
alt_cluster_id_1,
location_id=location_id,
serve_nodes=serve_nodes,
default_storage_type=enums.StorageType.HDD,
)
alt_cluster_id_2 = f"{alt_instance_id}-c2"
location_id_2 = "us-central1-f"
cluster_2 = instance.cluster(
alt_cluster_id_2,
location_id=location_id_2,
serve_nodes=serve_nodes,
default_storage_type=enums.StorageType.HDD,
)
operation = instance.create(clusters=[cluster_1, cluster_2])
instances_to_delete.append(instance)
operation.result(timeout=120) # Ensure the operation completes.
# Create a new instance instance and make sure it is the same.
instance_alt = admin_client.instance(alt_instance_id)
instance_alt.reload()
assert instance == instance_alt
assert instance.display_name == instance_alt.display_name
assert instance.type_ == instance_alt.type_
clusters, failed_locations = instance_alt.list_clusters()
assert failed_locations == []
alt_cluster_1, alt_cluster_2 = sorted(clusters, key=lambda x: x.name)
assert cluster_1.location_id == alt_cluster_1.location_id
assert alt_cluster_1.state == enums.Cluster.State.READY
assert cluster_1.serve_nodes == alt_cluster_1.serve_nodes
assert cluster_1.default_storage_type == alt_cluster_1.default_storage_type
assert cluster_2.location_id == alt_cluster_2.location_id
assert alt_cluster_2.state == enums.Cluster.State.READY
assert cluster_2.serve_nodes == alt_cluster_2.serve_nodes
assert cluster_2.default_storage_type == alt_cluster_2.default_storage_type
# Test list clusters in project via 'client.list_clusters'
clusters, failed_locations = admin_client.list_clusters()
assert not failed_locations
found = set([cluster.name for cluster in clusters])
expected = {alt_cluster_1.name, alt_cluster_2.name, admin_cluster.name}
assert expected.issubset(found)
temp_table_id = "test-get-cluster-states"
temp_table = instance.table(temp_table_id)
_helpers.retry_grpc_unavailable(temp_table.create)()
EncryptionType = enums.EncryptionInfo.EncryptionType
encryption_info = temp_table.get_encryption_info()
assert (encryption_info[alt_cluster_id_1][0].encryption_type ==
EncryptionType.GOOGLE_DEFAULT_ENCRYPTION)
assert (encryption_info[alt_cluster_id_2][0].encryption_type ==
EncryptionType.GOOGLE_DEFAULT_ENCRYPTION)
c_states = temp_table.get_cluster_states()
cluster_ids = set(c_states.keys())
assert cluster_ids == {alt_cluster_id_1, alt_cluster_id_2}
ReplicationState = enums.Table.ReplicationState
expected_results = [
ClusterState(ReplicationState.STATE_NOT_KNOWN),
ClusterState(ReplicationState.INITIALIZING),
ClusterState(ReplicationState.PLANNED_MAINTENANCE),
ClusterState(ReplicationState.UNPLANNED_MAINTENANCE),
ClusterState(ReplicationState.READY),
]
for clusterstate in c_states.values():
assert clusterstate in expected_results
# Test create app profile with multi_cluster_routing policy
app_profiles_to_delete = []
description = "routing policy-multy"
app_profile_id_1 = "app_profile_id_1"
routing = enums.RoutingPolicyType.ANY
app_profile_1 = _create_app_profile_helper(
app_profile_id_1,
instance,
routing_policy_type=routing,
description=description,
ignore_warnings=True,
)
app_profiles_to_delete.append(app_profile_1)
# Test list app profiles
_list_app_profiles_helper(instance, [app_profile_id_1])
# Test modify app profile app_profile_id_1
# routing policy to single cluster policy,
# cluster -> alt_cluster_id_1,
# allow_transactional_writes -> disallowed
# modify description
description = "to routing policy-single"
routing = enums.RoutingPolicyType.SINGLE
_modify_app_profile_helper(
app_profile_id_1,
instance,
routing_policy_type=routing,
description=description,
cluster_id=alt_cluster_id_1,
allow_transactional_writes=False,
)
# Test modify app profile app_profile_id_1
# cluster -> alt_cluster_id_2,
# allow_transactional_writes -> allowed
_modify_app_profile_helper(
app_profile_id_1,
instance,
routing_policy_type=routing,
description=description,
cluster_id=alt_cluster_id_2,
allow_transactional_writes=True,
ignore_warnings=True,
)
# Test create app profile with single cluster routing policy
description = "routing policy-single"
app_profile_id_2 = "app_profile_id_2"
routing = enums.RoutingPolicyType.SINGLE
app_profile_2 = _create_app_profile_helper(
app_profile_id_2,
instance,
routing_policy_type=routing,
description=description,
cluster_id=alt_cluster_id_2,
allow_transactional_writes=False,
)
app_profiles_to_delete.append(app_profile_2)
# Test list app profiles
_list_app_profiles_helper(instance, [app_profile_id_1, app_profile_id_2])
# Test modify app profile app_profile_id_2 to
# allow transactional writes
# Note: no need to set ``ignore_warnings`` to True
# since we are not restrictings anything with this modification.
_modify_app_profile_helper(
app_profile_id_2,
instance,
routing_policy_type=routing,
description=description,
cluster_id=alt_cluster_id_2,
allow_transactional_writes=True,
)
# Test modify app profile app_profile_id_2 routing policy
# to multi_cluster_routing policy
# modify description
description = "to routing policy-multy"
routing = enums.RoutingPolicyType.ANY
_modify_app_profile_helper(
app_profile_id_2,
instance,
routing_policy_type=routing,
description=description,
allow_transactional_writes=False,
ignore_warnings=True,
)
# Test delete app profiles
for app_profile in app_profiles_to_delete:
_delete_app_profile_helper(app_profile)
def _test_state_helper(self, instance, clusterid1, clusterid2):
# test get_cluster_states for a table in instance
from google.cloud.bigtable.enums import Table as enum_table
from google.cloud.bigtable.table import ClusterState
STATE_NOT_KNOWN = enum_table.ReplicationState.STATE_NOT_KNOWN
INITIALIZING = enum_table.ReplicationState.INITIALIZING
PLANNED_MAINTENANCE = enum_table.ReplicationState. \
PLANNED_MAINTENANCE
UNPLANNED_MAINTENANCE = enum_table.ReplicationState. \
UNPLANNED_MAINTENANCE
READY = enum_table.ReplicationState.READY
temp_table_id = 'test-get-cluster-states'
temp_table = instance.table(temp_table_id)
temp_table.create()
result = temp_table.get_cluster_states()
expected_results = [
ClusterState(STATE_NOT_KNOWN),
ClusterState(INITIALIZING),
ClusterState(PLANNED_MAINTENANCE),
ClusterState(UNPLANNED_MAINTENANCE),
ClusterState(READY)
]
cluster_id_list = result.keys()
self.assertEqual(len(cluster_id_list), 2)
self.assertIn(clusterid1, cluster_id_list)
self.assertIn(clusterid2, cluster_id_list)
for clusterstate in result.values():
self.assertIn(clusterstate, expected_results)
# Test create app profile with multi_cluster_routing policy
app_profiles_to_delete = []
description = 'routing policy-multy'
app_profile_id_1 = 'app_profile_id_1'
routing = enums.RoutingPolicyType.ANY
self._test_create_app_profile_helper(app_profile_id_1,
instance,
routing_policy_type=routing,
description=description,
ignore_warnings=True)
app_profiles_to_delete.append(app_profile_id_1)
# Test list app profiles
self._test_list_app_profiles_helper(instance, [app_profile_id_1])
# Test modify app profile app_profile_id_1
# routing policy to single cluster policy,
# cluster -> ALT_CLUSTER_ID_1,
# allow_transactional_writes -> disallowed
# modify description
description = 'to routing policy-single'
routing = enums.RoutingPolicyType.SINGLE
self._test_modify_app_profile_helper(app_profile_id_1,
instance,
routing_policy_type=routing,
description=description,
cluster_id=ALT_CLUSTER_ID_1,
allow_transactional_writes=False)
# Test modify app profile app_profile_id_1
# cluster -> ALT_CLUSTER_ID_2,
# allow_transactional_writes -> allowed
self._test_modify_app_profile_helper(app_profile_id_1,
instance,
routing_policy_type=routing,
description=description,
cluster_id=ALT_CLUSTER_ID_2,
allow_transactional_writes=True,
ignore_warnings=True)
# Test create app profile with single cluster routing policy
description = 'routing policy-single'
app_profile_id_2 = 'app_profile_id_2'
routing = enums.RoutingPolicyType.SINGLE
self._test_create_app_profile_helper(app_profile_id_2,
instance,
routing_policy_type=routing,
description=description,
cluster_id=ALT_CLUSTER_ID_2,
allow_transactional_writes=False)
app_profiles_to_delete.append(app_profile_id_2)
# Test list app profiles
self._test_list_app_profiles_helper(
instance, [app_profile_id_1, app_profile_id_2])
# Test modify app profile app_profile_id_2 to
# allow transactional writes
# Note: no need to set ``ignore_warnings`` to True
# since we are not restrictings anything with this modification.
self._test_modify_app_profile_helper(app_profile_id_2,
instance,
routing_policy_type=routing,
description=description,
cluster_id=ALT_CLUSTER_ID_2,
allow_transactional_writes=True)
# Test modify app profile app_profile_id_2 routing policy
# to multi_cluster_routing policy
# modify description
description = 'to routing policy-multy'
routing = enums.RoutingPolicyType.ANY
self._test_modify_app_profile_helper(app_profile_id_2,
instance,
routing_policy_type=routing,
description=description,
allow_transactional_writes=False,
ignore_warnings=True)
# Test delete app profiles
for app_profile_id in app_profiles_to_delete:
self._test_delete_app_profile_helper(app_profile_id, instance)
