vandonr-amz commented on code in PR #29911:
URL: https://github.com/apache/airflow/pull/29911#discussion_r1134386280


##########
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:
   I wanted to have the variable close to the usage so that when reading the 
function you can just raise your eyes to see the variable. And static checks 
are not complaining, so I guess all good ?



##########
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:
   yes good catch



-- 
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: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]

Reply via email to