feat(auth): automatically determine region

Author: Brooke-white

redshift_connector/iam_helper.py

@@ -119,16 +119,21 @@ def set_iam_properties(info: RedshiftProperty) -> RedshiftProperty:
                     "Please upgrade the installed version of boto3 to use this functionality."
                 )
 
+        # consider overridden connection parameters
         if info.is_serverless_host:
-            # consider overridden connection parameters
-            if not info.region:
-                info.set_region_from_host()
             if not info.serverless_acct_id:
                 info.set_serverless_acct_id()
             if not info.serverless_work_group:
                 info.set_serverless_work_group_from_host()
+        if not info.region:
+            info.set_region_from_host()
 
         if info.iam is True:
+
+            if info.region is None:
+                _logger.debug("Setting region via DNS lookup as region was not provided in connection parameters")
+                info.set_region_from_endpoint_lookup()
+
             if info.cluster_identifier is None and not info._is_serverless and not info.is_cname:
                 raise InterfaceError(
                     "Invalid connection property setting. cluster_identifier must be provided when IAM is enabled"

redshift_connector/redshift_property.py

@@ -223,17 +223,56 @@ def set_serverless_acct_id(self: "RedshiftProperty") -> None:
 
     def set_region_from_host(self: "RedshiftProperty") -> None:
         """
-        Sets the AWS region as parsed from the Redshift serverless endpoint.
+        Sets the AWS region as parsed from the Redshift instance endpoint.
         """
         import re
 
-        for serverless_pattern in (SERVERLESS_WITH_WORKGROUP_HOST_PATTERN, SERVERLESS_HOST_PATTERN):
-            m2 = re.fullmatch(pattern=serverless_pattern, string=self.host)
+        if self.is_serverless_host:
+            patterns: typing.Tuple[str, ...] = (SERVERLESS_WITH_WORKGROUP_HOST_PATTERN, SERVERLESS_HOST_PATTERN)
+        else:
+            patterns = (PROVISIONED_HOST_PATTERN,)
+
+        for host_pattern in patterns:
+            m2 = re.fullmatch(pattern=host_pattern, string=self.host)
 
             if m2:
                 self.put(key="region", value=m2.group(typing.cast(int, m2.lastindex)))
+                _logger.debug("region set to %s", self.region)
                 break
 
+    def set_region_from_endpoint_lookup(self: "RedshiftProperty") -> None:
+        """
+        Sets the AWS region as determined from a DNS lookup of the Redshift instance endpoint.
+        """
+        import socket
+
+        _logger.debug("set_region_from_endpoint_lookup")
+
+        if not all((self.host, self.port)):
+            _logger.debug("host and port were unspecified, exiting set_region_from_endpoint_lookup")
+            return
+        try:
+            addr_response: typing.List[
+                typing.Tuple[
+                    socket.AddressFamily,
+                    socket.SocketKind,
+                    int,
+                    str,
+                    typing.Union[typing.Tuple[str, int], typing.Tuple[str, int, int, int]],
+                ]
+            ] = socket.getaddrinfo(host=self.host, port=self.port, family=socket.AF_INET)
+            _logger.debug("%s", addr_response)
+            host_response: typing.Tuple[str, typing.List, typing.List] = socket.gethostbyaddr(addr_response[0][4][0])
+            ec2_instance_host: str = host_response[0]
+            _logger.debug("underlying ec2 instance host %s", ec2_instance_host)
+            ec2_region: str = ec2_instance_host.split(".")[1]
+            self.put(key="region", value=ec2_region)
+        except:
+            msg: str = "Unable to automatically determine AWS region from host {} port {}. Please check host and port connection parameters are correct.".format(
+                self.host, self.port
+            )
+            _logger.debug(msg)
+
     def set_serverless_work_group_from_host(self: "RedshiftProperty") -> None:
         """
         Sets the work_group as parsed from the Redshift serverless endpoint.

redshift_connector/utils/logging_utils.py

@@ -84,6 +84,9 @@ def is_populated(field: typing.Optional[str]):
         if parameter in logging_allow_list:
             temp.put(parameter, value)
         elif is_populated(value):
-            temp.put(parameter, "***")
+            try:
+                temp.put(parameter, "***")
+            except AttributeError:
+                pass
 
     return temp

test/conftest.py

@@ -98,7 +98,6 @@ def serverless_iam_db_kwargs() -> typing.Dict[str, typing.Union[str, bool]]:
         "access_key_id": conf.get("redshift-serverless", "access_key_id", fallback="mock_access_key_id"),
         "secret_access_key": conf.get("redshift-serverless", "secret_access_key", fallback="mock_secret_access_key"),
         "session_token": conf.get("redshift-serverless", "session_token", fallback="mock_session_token"),
-        "region": conf.get("redshift-serverless", "region", fallback="mock_region"),
         "host": conf.get(
             "redshift-serverless", "host", fallback="testwg1.012345678901.us-east-2.redshift-serverless.amazonaws.com"
         ),

test/integration/test_redshift_property.py

@@ -0,0 +1,37 @@
+import logging
+
+import pytest
+
+from redshift_connector import InterfaceError, RedshiftProperty
+
+LOGGER = logging.getLogger(__name__)
+LOGGER.propagate = True
+
+
+def test_set_region_from_endpoint_lookup(db_kwargs):
+    rp: RedshiftProperty = RedshiftProperty()
+    rp.put(key="host", value=db_kwargs["host"])
+    rp.put(key="port", value=db_kwargs["port"])
+    rp.set_region_from_host()
+
+    expected_region = rp.region
+    rp.region = None
+
+    rp.set_region_from_endpoint_lookup()
+    assert rp.region == expected_region
+
+
+@pytest.mark.parametrize("host, port", [("x", 1000), ("amazon.com", -1), ("-o", 5439)])
+def test_set_region_from_endpoint_lookup_raises(host, port, caplog):
+    import logging
+
+    rp: RedshiftProperty = RedshiftProperty()
+    rp.put(key="host", value=host)
+    rp.put(key="port", value=port)
+    expected_msg: str = "Unable to automatically determine AWS region from host {} port {}. Please check host and port connection parameters are correct.".format(
+        host, port
+    )
+
+    with caplog.at_level(logging.DEBUG):
+        rp.set_region_from_endpoint_lookup()
+    assert expected_msg in caplog.text

test/manual/test_redshift_custom_domain.py

@@ -50,7 +50,6 @@ def test_iam_connect(provisioned_cname_db_kwargs, sslmode):
     provisioned_cname_db_kwargs["iam"] = True
     provisioned_cname_db_kwargs["profile"] = "default"
     provisioned_cname_db_kwargs["auto_create"] = True
-    provisioned_cname_db_kwargs["region"] = "eu-north-1"
     provisioned_cname_db_kwargs["ssl"] = True
     provisioned_cname_db_kwargs["sslmode"] = sslmode.value
     with redshift_connector.connect(**provisioned_cname_db_kwargs):

test/unit/test_redshift_property.py

@@ -38,6 +38,8 @@ def test_set_serverless_acct_id_from_host(host, exp_account_id):
         ("testwg1.012345678901.us-east-2.redshift-serverless.amazonaws.com", "us-east-2"),
         ("123456789012.us-south-1.redshift-serverless.amazonaws.com", "us-south-1"),
         ("testwg2.012345678901.ap-northeast-3.redshift-serverless.amazonaws.com", "ap-northeast-3"),
+        ("redshift-cluster-1.aaaaaaaaaaaa.us-east-2.redshift.amazonaws.com", "us-east-2"),
+        ("mylongredshiftclustername.aaaaaaaaaaaa.ap-northeast-1.redshift.amazonaws.com", "ap-northeast-1"),
     ],
 )
 def test_set_region_from_host(host, exp_region):