Skip to content

Commit

Permalink
Cache Hook when initializing CloudFormationCreateStackSensor (#14638)
Browse files Browse the repository at this point in the history
- `CloudFormationCreateStackSensor` created `AWSCloudFormationHook` in its init method
- This commit delays creation of hook until it is needed
  • Loading branch information
kaxil committed Mar 7, 2021
1 parent 6d087fb commit b896aee
Showing 1 changed file with 16 additions and 9 deletions.
25 changes: 16 additions & 9 deletions airflow/providers/amazon/aws/sensors/cloud_formation.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,11 @@
"""This module contains sensors for AWS CloudFormation."""
from typing import Optional

try:
from functools import cached_property
except ImportError:
from cached_property import cached_property

from airflow.providers.amazon.aws.hooks.cloud_formation import AWSCloudFormationHook
from airflow.sensors.base import BaseSensorOperator
from airflow.utils.decorators import apply_defaults
Expand All @@ -43,7 +48,8 @@ class CloudFormationCreateStackSensor(BaseSensorOperator):
def __init__(self, *, stack_name, aws_conn_id='aws_default', region_name=None, **kwargs):
super().__init__(**kwargs)
self.stack_name = stack_name
self.hook = AWSCloudFormationHook(aws_conn_id=aws_conn_id, region_name=region_name)
self.aws_conn_id = aws_conn_id
self.region_name = region_name

def poke(self, context):
stack_status = self.hook.get_stack_status(self.stack_name)
Expand All @@ -53,6 +59,11 @@ def poke(self, context):
return False
raise ValueError(f'Stack {self.stack_name} in bad state: {stack_status}')

@cached_property
def hook(self) -> AWSCloudFormationHook:
"""Create and return an AWSCloudFormationHook"""
return AWSCloudFormationHook(aws_conn_id=self.aws_conn_id, region_name=self.region_name)


class CloudFormationDeleteStackSensor(BaseSensorOperator):
"""
Expand Down Expand Up @@ -83,20 +94,16 @@ def __init__(
self.aws_conn_id = aws_conn_id
self.region_name = region_name
self.stack_name = stack_name
self.hook: Optional[AWSCloudFormationHook] = None

def poke(self, context):
stack_status = self.get_hook().get_stack_status(self.stack_name)
stack_status = self.hook.get_stack_status(self.stack_name)
if stack_status in ('DELETE_COMPLETE', None):
return True
if stack_status == 'DELETE_IN_PROGRESS':
return False
raise ValueError(f'Stack {self.stack_name} in bad state: {stack_status}')

def get_hook(self) -> AWSCloudFormationHook:
@cached_property
def hook(self) -> AWSCloudFormationHook:
"""Create and return an AWSCloudFormationHook"""
if self.hook:
return self.hook

self.hook = AWSCloudFormationHook(aws_conn_id=self.aws_conn_id, region_name=self.region_name)
return self.hook
return AWSCloudFormationHook(aws_conn_id=self.aws_conn_id, region_name=self.region_name)

0 comments on commit b896aee

Please sign in to comment.