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)