def __init__(
self,
project_id: str,
region: str,
operation_name: str,
gcp_conn_id: str = "google_cloud_default",
impersonation_chain: Optional[Union[str, Sequence[str]]] = None,
delegate_to: Optional[str] = None,
pooling_period_seconds: int = 30,
):
super().__init__()
self.project_id = project_id
self.region = region
self.operation_name = operation_name
self.gcp_conn_id = gcp_conn_id
self.impersonation_chain = impersonation_chain
self.delegate_to = delegate_to
self.pooling_period_seconds = pooling_period_seconds
self.gcp_hook = CloudComposerHook(
gcp_conn_id=self.gcp_conn_id,
impersonation_chain=self.impersonation_chain,
delegate_to=self.delegate_to,
)
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.delete_environment(
project_id=self.project_id,
region=self.region,
environment_id=self.environment_id,
retry=self.retry,
timeout=self.timeout,
metadata=self.metadata,
)
if not self.deferrable:
hook.wait_for_operation(timeout=self.timeout, operation=result)
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,
)
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.list_image_versions(
project_id=self.project_id,
region=self.region,
page_size=self.page_size,
page_token=self.page_token,
include_past_releases=self.include_past_releases,
retry=self.retry,
timeout=self.timeout,
metadata=self.metadata,
)
return [ImageVersion.to_dict(image) for image in result]
def execute(self, context: 'Context'):
hook = CloudComposerHook(
gcp_conn_id=self.gcp_conn_id,
impersonation_chain=self.impersonation_chain,
delegate_to=self.delegate_to,
)
CloudComposerEnvironmentsLink.persist(operator_instance=self, context=context)
result = hook.list_environments(
project_id=self.project_id,
region=self.region,
page_size=self.page_size,
page_token=self.page_token,
retry=self.retry,
timeout=self.timeout,
metadata=self.metadata,
)
return [Environment.to_dict(env) for env in result]
class CloudComposerExecutionTrigger(BaseTrigger):
"""The trigger handles the async communication with the Google Cloud Composer"""
def __init__(
self,
project_id: str,
region: str,
operation_name: str,
gcp_conn_id: str = "google_cloud_default",
impersonation_chain: Optional[Union[str, Sequence[str]]] = None,
delegate_to: Optional[str] = None,
pooling_period_seconds: int = 30,
):
super().__init__()
self.project_id = project_id
self.region = region
self.operation_name = operation_name
self.gcp_conn_id = gcp_conn_id
self.impersonation_chain = impersonation_chain
self.delegate_to = delegate_to
self.pooling_period_seconds = pooling_period_seconds
self.gcp_hook = CloudComposerHook(
gcp_conn_id=self.gcp_conn_id,
impersonation_chain=self.impersonation_chain,
delegate_to=self.delegate_to,
)
def serialize(self) -> Tuple[str, Dict[str, Any]]:
return (
'airflow.providers.google.cloud.triggers.cloud_composer.CloudComposerExecutionTrigger',
{
"project_id": self.project_id,
"region": self.region,
"operation_name": self.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,
},
)
async def run(self):
while True:
operation = self.gcp_hook.get_operation(
operation_name=self.operation_name)
if operation.done:
break
elif operation.error.message:
raise AirflowException(
f"Cloud Composer Environment error: {operation.error.message}"
)
await asyncio.sleep(self.pooling_period_seconds)
yield TriggerEvent({
"operation_name": operation.name,
"operation_done": operation.done,
})
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']}")
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)
