xinbinhuang commented on a change in pull request #15104:
URL: https://github.com/apache/airflow/pull/15104#discussion_r606105285
##########
File path: airflow/providers/amazon/aws/secrets/secrets_manager.py
##########
@@ -102,17 +102,66 @@ def client(self):
)
return session.client(service_name="secretsmanager", **self.kwargs)
- def get_conn_uri(self, conn_id: str) -> Optional[str]:
+ def _get_extra(self, secret, conn_string):
+ if 'extra' in secret:
+ extra_dict = ast.literal_eval(secret['extra'])
+ counter = 0
+ for key, value in extra_dict.items():
+ if counter == 0:
+ conn_string += f'?{key}={value}'
+ else:
+ conn_string += f'&{key}={value}'
+
+ counter += 1
+
+ return conn_string
+
+ def get_conn_uri(self, conn_id: str):
"""
Get Connection Value
:param conn_id: connection id
:type conn_id: str
"""
- if self.connections_prefix is None:
- return None
+ if self.connections_prefix and self.sep:
+ conn_id = self.build_path(self.connections_prefix, conn_id,
self.sep)
- return self._get_secret(self.connections_prefix, conn_id)
+ try:
+ secret_string = self._get_secret(conn_id)
+ secret = ast.literal_eval(secret_string) # json.loads gives error
+ except ValueError: # 'malformed node or string: ' error, for empty
conns
+ connection = None
+ secret = None
+
+ # These lines will check if we have with some denomination stored an
username, password and host
+ if secret:
+ possible_words_for_conn_fields = {
+ 'user': ['user', 'username', 'login', 'user_name'],
+ 'password': ['password', 'pass', 'key'],
+ 'host': ['host', 'remote_host', 'server'],
+ 'port': ['port'],
+ 'schema': ['database', 'schema'],
+ 'conn_type': ['conn_type', 'conn_id', 'connection_type',
'engine']
+ }
+
+ conn_d = {}
+ for conn_field, possible_words in
possible_words_for_conn_fields.items():
+ try:
+ conn_d[conn_field] = [v for k, v in secret.items() if k in
possible_words][0]
+ except IndexError:
+ conn_d[conn_field] = ''
+
+ conn_type = conn_d['conn_type']
+ user = conn_d['user']
+ password = conn_d['password']
+ host = conn_d['host']
+ port = conn_d['port']
+ schema = conn_d['schema']
+ conn_string =
f'{conn_type}://{user}:{password}@{host}:{port}/{schema}'
+
+ connection = self._get_extra(secret, conn_string)
Review comment:
Also, I would recommend separating this part of logic (`if secret: `)
into another method for easier testing
--
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.
For queries about this service, please contact Infrastructure at:
[email protected]
