Moved cluster states to hook

apache · Jan 26, 2024 · 1cab491 · 1cab491
1 parent 1b20b51
commit 1cab491
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 5 deletions.
diff --git a/airflow/providers/amazon/aws/hooks/neptune.py b/airflow/providers/amazon/aws/hooks/neptune.py
@@ -32,6 +32,9 @@ class NeptuneHook(AwsBaseHook):
         - :class:`~airflow.providers.amazon.aws.hooks.base_aws.AwsBaseHook`
     """
 
+    AVAILABLE_STATES = ["available"]
+    STOPPED_STATES = ["stopped"]
+
     def __init__(self, *args, **kwargs):
         kwargs["client_type"] = "neptune"
         super().__init__(*args, **kwargs)

diff --git a/airflow/providers/amazon/aws/operators/neptune.py b/airflow/providers/amazon/aws/operators/neptune.py
@@ -31,9 +31,6 @@
 if TYPE_CHECKING:
     from airflow.utils.context import Context
 
-AVAILABLE_STATES = ["available"]
-STOPPED_STATES = ["stopped"]
-
 
 class NeptuneStartDbClusterOperator(AwsBaseOperator[NeptuneHook]):
     """Starts an Amazon Neptune DB cluster.
@@ -89,7 +86,7 @@ def execute(self, context: Context) -> dict[str, str]:
 
         # Check to make sure the cluster is not already available.
         status = self.hook.get_cluster_status(self.cluster_id)
-        if status.lower() in AVAILABLE_STATES:
+        if status.lower() in NeptuneHook.AVAILABLE_STATES:
             self.log.info("Neptune cluster %s is already available.", self.cluster_id)
             return {"db_cluster_id": self.cluster_id}
 
@@ -183,7 +180,7 @@ def execute(self, context: Context) -> dict[str, str]:
 
         # Check to make sure the cluster is not already stopped.
         status = self.hook.get_cluster_status(self.cluster_id)
-        if status.lower() in STOPPED_STATES:
+        if status.lower() in NeptuneHook.STOPPED_STATES:
             self.log.info("Neptune cluster %s is already stopped.", self.cluster_id)
             return {"db_cluster_id": self.cluster_id}