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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
