def execute(self, context):
hook = BigtableHook(gcp_conn_id=self.gcp_conn_id)
instance = hook.get_instance(project_id=self.project_id,
instance_id=self.instance_id)
if instance:
# Based on Instance.__eq__ instance with the same ID and client is
# considered as equal.
self.log.info(
"The instance '%s' already exists in this project. "
"Consider it as created", self.instance_id)
return
try:
hook.create_instance(
project_id=self.project_id,
instance_id=self.instance_id,
main_cluster_id=self.main_cluster_id,
main_cluster_zone=self.main_cluster_zone,
replica_cluster_id=self.replica_cluster_id,
replica_cluster_zone=self.replica_cluster_zone,
instance_display_name=self.instance_display_name,
instance_type=self.instance_type,
instance_labels=self.instance_labels,
cluster_nodes=self.cluster_nodes,
cluster_storage_type=self.cluster_storage_type,
timeout=self.timeout,
)
except google.api_core.exceptions.GoogleAPICallError as e:
self.log.error('An error occurred. Exiting.')
raise e
def poke(self, context):
hook = BigtableHook(gcp_conn_id=self.gcp_conn_id)
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)
return True
def __init__(self, # pylint: disable=too-many-arguments
instance_id,
main_cluster_id,
main_cluster_zone,
project_id=None,
replica_cluster_id=None,
replica_cluster_zone=None,
instance_display_name=None,
instance_type=None,
instance_labels=None,
cluster_nodes=None,
cluster_storage_type=None,
timeout=None,
gcp_conn_id='google_cloud_default',
*args, **kwargs):
self.project_id = project_id
self.instance_id = instance_id
self.main_cluster_id = main_cluster_id
self.main_cluster_zone = main_cluster_zone
self.replica_cluster_id = replica_cluster_id
self.replica_cluster_zone = replica_cluster_zone
self.instance_display_name = instance_display_name
self.instance_type = instance_type
self.instance_labels = instance_labels
self.cluster_nodes = cluster_nodes
self.cluster_storage_type = cluster_storage_type
self.timeout = timeout
self._validate_inputs()
self.hook = BigtableHook(gcp_conn_id=gcp_conn_id)
super().__init__(*args, **kwargs)
def __init__(self,
instance_id,
project_id=None,
gcp_conn_id='google_cloud_default',
*args, **kwargs):
self.project_id = project_id
self.instance_id = instance_id
self._validate_inputs()
self.hook = BigtableHook(gcp_conn_id=gcp_conn_id)
super().__init__(*args, **kwargs)
def execute(self, context):
hook = BigtableHook(gcp_conn_id=self.gcp_conn_id)
try:
hook.delete_instance(project_id=self.project_id,
instance_id=self.instance_id)
except google.api_core.exceptions.NotFound:
self.log.info(
"The instance '%s' does not exist in project '%s'. "
"Consider it as deleted", self.instance_id, self.project_id)
except google.api_core.exceptions.GoogleAPICallError as e:
self.log.error('An error occurred. Exiting.')
raise e
def __init__(self,
instance_id,
table_id,
project_id=None,
initial_split_keys=None,
column_families=None,
gcp_conn_id='google_cloud_default',
*args, **kwargs):
self.project_id = project_id
self.instance_id = instance_id
self.table_id = table_id
self.initial_split_keys = initial_split_keys or list()
self.column_families = column_families or dict()
self._validate_inputs()
self.hook = BigtableHook(gcp_conn_id=gcp_conn_id)
self.instance = None
super().__init__(*args, **kwargs)
def execute(self, context):
hook = BigtableHook(gcp_conn_id=self.gcp_conn_id)
instance = hook.get_instance(project_id=self.project_id,
instance_id=self.instance_id)
if not instance:
raise AirflowException(
"Dependency: instance '{}' does not exist.".format(
self.instance_id))
try:
hook.update_cluster(instance=instance,
cluster_id=self.cluster_id,
nodes=self.nodes)
except google.api_core.exceptions.NotFound:
raise AirflowException(
"Dependency: cluster '{}' does not exist for instance '{}'.".
format(self.cluster_id, self.instance_id))
except google.api_core.exceptions.GoogleAPICallError as e:
self.log.error('An error occurred. Exiting.')
raise e
def execute(self, context):
hook = BigtableHook(gcp_conn_id=self.gcp_conn_id)
instance = hook.get_instance(project_id=self.project_id,
instance_id=self.instance_id)
if not instance:
raise AirflowException(
"Dependency: instance '{}' does not exist in project '{}'.".
format(self.instance_id, self.project_id))
try:
hook.create_table(instance=instance,
table_id=self.table_id,
initial_split_keys=self.initial_split_keys,
column_families=self.column_families)
except google.api_core.exceptions.AlreadyExists:
if not self._compare_column_families(hook, instance):
raise AirflowException(
"Table '{}' already exists with different Column Families."
.format(self.table_id))
self.log.info(
"The table '%s' already exists. Consider it as created",
self.table_id)
def execute(self, context):
hook = BigtableHook(gcp_conn_id=self.gcp_conn_id)
instance = hook.get_instance(project_id=self.project_id,
instance_id=self.instance_id)
if not instance:
raise AirflowException("Dependency: instance '{}' does not exist.".format(
self.instance_id))
try:
hook.delete_table(
project_id=self.project_id,
instance_id=self.instance_id,
table_id=self.table_id,
)
except google.api_core.exceptions.NotFound:
# It's OK if table doesn't exists.
self.log.info("The table '%s' no longer exists. Consider it as deleted",
self.table_id)
except google.api_core.exceptions.GoogleAPICallError as e:
self.log.error('An error occurred. Exiting.')
raise e
def setUp(self):
with mock.patch(
'airflow.contrib.hooks.gcp_api_base_hook.GoogleCloudBaseHook.__init__',
new=mock_base_gcp_hook_no_default_project_id):
self.bigtable_hook_no_default_project_id = BigtableHook(
gcp_conn_id='test')
def setUp(self):
with mock.patch('airflow.gcp.hooks.base.CloudBaseHook.__init__',
new=mock_base_gcp_hook_default_project_id):
self.bigtable_hook_default_project_id = BigtableHook(
gcp_conn_id='test')