[GitHub] [airflow] vincbeck commented on a diff in pull request #29911: custom waiters with dynamic values, applied to appflow

["vincbeck (via GitHub)" <gi...@apache.org>]

vincbeck commented on code in PR #29911:
URL: https://github.com/apache/airflow/pull/29911#discussion_r1126630419


##########
airflow/providers/amazon/aws/hooks/base_aws.py:
##########
@@ -794,25 +794,66 @@ def waiter_path(self) -> PathLike[str] | None:
         path = Path(__file__).parents[1].joinpath(f"waiters/{self.client_type}.json").resolve()
         return path if path.exists() else None
 
-    def get_waiter(self, waiter_name: str) -> Waiter:
+    def get_waiter(self, waiter_name: str, parameters: dict[str, str] | None = None) -> Waiter:
         """
         First checks if there is a custom waiter with the provided waiter_name and
         uses that if it exists, otherwise it will check the service client for a
         waiter that matches the name and pass that through.
 
         :param waiter_name: The name of the waiter.  The name should exactly match the
             name of the key in the waiter model file (typically this is CamelCase).
+        :param parameters: will scan the waiter config for the keys of that dict, and replace them with the
+            corresponding value. If a custom waiter has such keys to be expanded, they need to be provided
+            here.
         """
         if self.waiter_path and (waiter_name in self._list_custom_waiters()):
             # Technically if waiter_name is in custom_waiters then self.waiter_path must
             # exist but MyPy doesn't like the fact that self.waiter_path could be None.
             with open(self.waiter_path) as config_file:
-                config = json.load(config_file)
-                return BaseBotoWaiter(client=self.conn, model_config=config).waiter(waiter_name)
+                config_text = config_file.read()
+
+            if parameters:  # expand some variables in the config if needed
+                config_text = self._apply_parameters_value(config_text, parameters)
+            config = json.loads(config_text)
+            self._check_remaining_params(config, waiter_name)
+            return BaseBotoWaiter(client=self.conn, model_config=config).waiter(waiter_name)
         # If there is no custom waiter found for the provided name,
         # then try checking the service's official waiters.
         return self.conn.get_waiter(waiter_name)
 
+    waiter_config_key_name_format = re.compile("^[A-Z0-9_]+$")
+    waiter_config_key_placeholder_format = re.compile("<[A-Z0-9_]+>")
+
+    def _apply_parameters_value(self, config: str, parameters: dict[str, str]) -> str:
+        """replaces the given parameters in the config with their value"""
+        for k, v in parameters.items():
+            if not self.waiter_config_key_name_format.match(k):
+                raise AirflowException(
+                    f"{k} is not an accepted key for waiter configuration templatization. "
+                    "Use only capitals, digits and underscores."
+                )
+            if config.find(k) == -1:

Review Comment:
   We should try to find `f"<{k}>"` instead of `k`. There is a hole here, if a user forget to surround the variable with `<>`, it will not fail here but at line `841` while doing the replace



##########
airflow/providers/amazon/aws/hooks/base_aws.py:
##########
@@ -794,25 +794,66 @@ def waiter_path(self) -> PathLike[str] | None:
         path = Path(__file__).parents[1].joinpath(f"waiters/{self.client_type}.json").resolve()
         return path if path.exists() else None
 
-    def get_waiter(self, waiter_name: str) -> Waiter:
+    def get_waiter(self, waiter_name: str, parameters: dict[str, str] | None = None) -> Waiter:
         """
         First checks if there is a custom waiter with the provided waiter_name and
         uses that if it exists, otherwise it will check the service client for a
         waiter that matches the name and pass that through.
 
         :param waiter_name: The name of the waiter.  The name should exactly match the
             name of the key in the waiter model file (typically this is CamelCase).
+        :param parameters: will scan the waiter config for the keys of that dict, and replace them with the
+            corresponding value. If a custom waiter has such keys to be expanded, they need to be provided
+            here.
         """
         if self.waiter_path and (waiter_name in self._list_custom_waiters()):
             # Technically if waiter_name is in custom_waiters then self.waiter_path must
             # exist but MyPy doesn't like the fact that self.waiter_path could be None.
             with open(self.waiter_path) as config_file:
-                config = json.load(config_file)
-                return BaseBotoWaiter(client=self.conn, model_config=config).waiter(waiter_name)
+                config_text = config_file.read()
+
+            if parameters:  # expand some variables in the config if needed
+                config_text = self._apply_parameters_value(config_text, parameters)
+            config = json.loads(config_text)
+            self._check_remaining_params(config, waiter_name)
+            return BaseBotoWaiter(client=self.conn, model_config=config).waiter(waiter_name)
         # If there is no custom waiter found for the provided name,
         # then try checking the service's official waiters.
         return self.conn.get_waiter(waiter_name)
 
+    waiter_config_key_name_format = re.compile("^[A-Z0-9_]+$")
+    waiter_config_key_placeholder_format = re.compile("<[A-Z0-9_]+>")

Review Comment:
   nit. In terms of convention, we do want to put them here or at the top with the other instance variables?



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org