Skip to content

Commit cd2aafe

Browse files
committed
feat(auth): automatically determine region
1 parent d83e55f commit cd2aafe

File tree

7 files changed

+93
-9
lines changed

7 files changed

+93
-9
lines changed

redshift_connector/iam_helper.py

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -119,16 +119,21 @@ def set_iam_properties(info: RedshiftProperty) -> RedshiftProperty:
119119
"Please upgrade the installed version of boto3 to use this functionality."
120120
)
121121

122+
# consider overridden connection parameters
122123
if info.is_serverless_host:
123-
# consider overridden connection parameters
124-
if not info.region:
125-
info.set_region_from_host()
126124
if not info.serverless_acct_id:
127125
info.set_serverless_acct_id()
128126
if not info.serverless_work_group:
129127
info.set_serverless_work_group_from_host()
128+
if not info.region:
129+
info.set_region_from_host()
130130

131131
if info.iam is True:
132+
133+
if info.region is None:
134+
_logger.debug("Setting region via DNS lookup as region was not provided in connection parameters")
135+
info.set_region_from_endpoint_lookup()
136+
132137
if info.cluster_identifier is None and not info._is_serverless and not info.is_cname:
133138
raise InterfaceError(
134139
"Invalid connection property setting. cluster_identifier must be provided when IAM is enabled"

redshift_connector/redshift_property.py

Lines changed: 42 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -223,17 +223,56 @@ def set_serverless_acct_id(self: "RedshiftProperty") -> None:
223223

224224
def set_region_from_host(self: "RedshiftProperty") -> None:
225225
"""
226-
Sets the AWS region as parsed from the Redshift serverless endpoint.
226+
Sets the AWS region as parsed from the Redshift instance endpoint.
227227
"""
228228
import re
229229

230-
for serverless_pattern in (SERVERLESS_WITH_WORKGROUP_HOST_PATTERN, SERVERLESS_HOST_PATTERN):
231-
m2 = re.fullmatch(pattern=serverless_pattern, string=self.host)
230+
if self.is_serverless_host:
231+
patterns: typing.Tuple[str, ...] = (SERVERLESS_WITH_WORKGROUP_HOST_PATTERN, SERVERLESS_HOST_PATTERN)
232+
else:
233+
patterns = (PROVISIONED_HOST_PATTERN,)
234+
235+
for host_pattern in patterns:
236+
m2 = re.fullmatch(pattern=host_pattern, string=self.host)
232237

233238
if m2:
234239
self.put(key="region", value=m2.group(typing.cast(int, m2.lastindex)))
240+
_logger.debug("region set to %s", self.region)
235241
break
236242

243+
def set_region_from_endpoint_lookup(self: "RedshiftProperty") -> None:
244+
"""
245+
Sets the AWS region as determined from a DNS lookup of the Redshift instance endpoint.
246+
"""
247+
import socket
248+
249+
_logger.debug("set_region_from_endpoint_lookup")
250+
251+
if not all((self.host, self.port)):
252+
_logger.debug("host and port were unspecified, exiting set_region_from_endpoint_lookup")
253+
return
254+
try:
255+
addr_response: typing.List[
256+
typing.Tuple[
257+
socket.AddressFamily,
258+
socket.SocketKind,
259+
int,
260+
str,
261+
typing.Union[typing.Tuple[str, int], typing.Tuple[str, int, int, int]],
262+
]
263+
] = socket.getaddrinfo(host=self.host, port=self.port, family=socket.AF_INET)
264+
_logger.debug("%s", addr_response)
265+
host_response: typing.Tuple[str, typing.List, typing.List] = socket.gethostbyaddr(addr_response[0][4][0])
266+
ec2_instance_host: str = host_response[0]
267+
_logger.debug("underlying ec2 instance host %s", ec2_instance_host)
268+
ec2_region: str = ec2_instance_host.split(".")[1]
269+
self.put(key="region", value=ec2_region)
270+
except:
271+
msg: str = "Unable to automatically determine AWS region from host {} port {}. Please check host and port connection parameters are correct.".format(
272+
self.host, self.port
273+
)
274+
_logger.debug(msg)
275+
237276
def set_serverless_work_group_from_host(self: "RedshiftProperty") -> None:
238277
"""
239278
Sets the work_group as parsed from the Redshift serverless endpoint.

redshift_connector/utils/logging_utils.py

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -84,6 +84,9 @@ def is_populated(field: typing.Optional[str]):
8484
if parameter in logging_allow_list:
8585
temp.put(parameter, value)
8686
elif is_populated(value):
87-
temp.put(parameter, "***")
87+
try:
88+
temp.put(parameter, "***")
89+
except AttributeError:
90+
pass
8891

8992
return temp

test/conftest.py

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -98,7 +98,6 @@ def serverless_iam_db_kwargs() -> typing.Dict[str, typing.Union[str, bool]]:
9898
"access_key_id": conf.get("redshift-serverless", "access_key_id", fallback="mock_access_key_id"),
9999
"secret_access_key": conf.get("redshift-serverless", "secret_access_key", fallback="mock_secret_access_key"),
100100
"session_token": conf.get("redshift-serverless", "session_token", fallback="mock_session_token"),
101-
"region": conf.get("redshift-serverless", "region", fallback="mock_region"),
102101
"host": conf.get(
103102
"redshift-serverless", "host", fallback="testwg1.012345678901.us-east-2.redshift-serverless.amazonaws.com"
104103
),
Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,37 @@
1+
import logging
2+
3+
import pytest
4+
5+
from redshift_connector import InterfaceError, RedshiftProperty
6+
7+
LOGGER = logging.getLogger(__name__)
8+
LOGGER.propagate = True
9+
10+
11+
def test_set_region_from_endpoint_lookup(db_kwargs):
12+
rp: RedshiftProperty = RedshiftProperty()
13+
rp.put(key="host", value=db_kwargs["host"])
14+
rp.put(key="port", value=db_kwargs["port"])
15+
rp.set_region_from_host()
16+
17+
expected_region = rp.region
18+
rp.region = None
19+
20+
rp.set_region_from_endpoint_lookup()
21+
assert rp.region == expected_region
22+
23+
24+
@pytest.mark.parametrize("host, port", [("x", 1000), ("amazon.com", -1), ("-o", 5439)])
25+
def test_set_region_from_endpoint_lookup_raises(host, port, caplog):
26+
import logging
27+
28+
rp: RedshiftProperty = RedshiftProperty()
29+
rp.put(key="host", value=host)
30+
rp.put(key="port", value=port)
31+
expected_msg: str = "Unable to automatically determine AWS region from host {} port {}. Please check host and port connection parameters are correct.".format(
32+
host, port
33+
)
34+
35+
with caplog.at_level(logging.DEBUG):
36+
rp.set_region_from_endpoint_lookup()
37+
assert expected_msg in caplog.text

test/manual/test_redshift_custom_domain.py

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -50,7 +50,6 @@ def test_iam_connect(provisioned_cname_db_kwargs, sslmode):
5050
provisioned_cname_db_kwargs["iam"] = True
5151
provisioned_cname_db_kwargs["profile"] = "default"
5252
provisioned_cname_db_kwargs["auto_create"] = True
53-
provisioned_cname_db_kwargs["region"] = "eu-north-1"
5453
provisioned_cname_db_kwargs["ssl"] = True
5554
provisioned_cname_db_kwargs["sslmode"] = sslmode.value
5655
with redshift_connector.connect(**provisioned_cname_db_kwargs):

test/unit/test_redshift_property.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,8 @@ def test_set_serverless_acct_id_from_host(host, exp_account_id):
3838
("testwg1.012345678901.us-east-2.redshift-serverless.amazonaws.com", "us-east-2"),
3939
("123456789012.us-south-1.redshift-serverless.amazonaws.com", "us-south-1"),
4040
("testwg2.012345678901.ap-northeast-3.redshift-serverless.amazonaws.com", "ap-northeast-3"),
41+
("redshift-cluster-1.aaaaaaaaaaaa.us-east-2.redshift.amazonaws.com", "us-east-2"),
42+
("mylongredshiftclustername.aaaaaaaaaaaa.ap-northeast-1.redshift.amazonaws.com", "ap-northeast-1"),
4143
],
4244
)
4345
def test_set_region_from_host(host, exp_region):

0 commit comments

Comments
 (0)