def execute(self, context: 'Context'):
hook = CloudComposerHook(
gcp_conn_id=self.gcp_conn_id,
impersonation_chain=self.impersonation_chain,
delegate_to=self.delegate_to,
)
name = hook.get_environment_name(self.project_id, self.region,
self.environment_id)
if isinstance(self.environment, Environment):
self.environment.name = name
else:
self.environment["name"] = name
CloudComposerEnvironmentLink.persist(operator_instance=self,
context=context)
try:
result = hook.create_environment(
project_id=self.project_id,
region=self.region,
environment=self.environment,
retry=self.retry,
timeout=self.timeout,
metadata=self.metadata,
)
if not self.deferrable:
environment = hook.wait_for_operation(timeout=self.timeout,
operation=result)
return Environment.to_dict(environment)
else:
self.defer(
trigger=CloudComposerExecutionTrigger(
project_id=self.project_id,
region=self.region,
operation_name=result.operation.name,
gcp_conn_id=self.gcp_conn_id,
impersonation_chain=self.impersonation_chain,
delegate_to=self.delegate_to,
pooling_period_seconds=self.pooling_period_seconds,
),
method_name=GOOGLE_DEFAULT_DEFERRABLE_METHOD_NAME,
)
except AlreadyExists:
environment = hook.get_environment(
project_id=self.project_id,
region=self.region,
environment_id=self.environment_id,
retry=self.retry,
timeout=self.timeout,
metadata=self.metadata,
)
return Environment.to_dict(environment)
def execute(self, context: 'Context'):
hook = CloudComposerHook(
gcp_conn_id=self.gcp_conn_id,
impersonation_chain=self.impersonation_chain,
delegate_to=self.delegate_to,
)
result = hook.get_environment(
project_id=self.project_id,
region=self.region,
environment_id=self.environment_id,
retry=self.retry,
timeout=self.timeout,
metadata=self.metadata,
)
CloudComposerEnvironmentLink.persist(operator_instance=self, context=context)
return Environment.to_dict(result)
def execute_complete(self, context: "Context", event: dict):
if event["operation_done"]:
hook = CloudComposerHook(
gcp_conn_id=self.gcp_conn_id,
impersonation_chain=self.impersonation_chain,
delegate_to=self.delegate_to,
)
env = hook.get_environment(
project_id=self.project_id,
region=self.region,
environment_id=self.environment_id,
retry=self.retry,
timeout=self.timeout,
metadata=self.metadata,
)
return Environment.to_dict(env)
else:
raise AirflowException(f"Unexpected error in the operation: {event['operation_name']}")
