Repository: dbt-labs/dbt-redshift
Pull request #421 (2023-04-28 23:42:42)

[Backport 1.5.latest] Fix for #419: dbt-redshift 1.5 does not work with non-standard redshift hostnames
Backport b644963a6bf7d3e64bd71e6ee97a205900089db1 from #420.

leahwicz: @nssalian is this a regression or just a bug fix we want to get in? How risky would you say this change is since it would go straight to production runs 
nssalian: > @nssalian is this a regression or just a bug fix we want to get in? How risky would you say this change is since it would go straight to production runs

This is a fix for the regression reported #419. I reviewed it with the Redshift team @sathiish-kumar was helpful to review here and we don't see a risk for runs. The ConnectionManager should at most be called once so the region calculation should be once. And if it fails, there's a fallback to pass through so things aren't blocked.

Patch:
diff --git a/.changes/unreleased/Fixes-20230428-142321.yaml b/.changes/unreleased/Fixes-20230428-142321.yaml
new file mode 100644
index 00000000..c7fefda4
--- /dev/null
+++ b/.changes/unreleased/Fixes-20230428-142321.yaml
@@ -0,0 +1,6 @@
+kind: Fixes
+body: Adding region as independent param in profiles
+time: 2023-04-28T14:23:21.041865-07:00
+custom:
+  Author: nssalian
+  Issue: "419"
diff --git a/dbt/adapters/redshift/connections.py b/dbt/adapters/redshift/connections.py
index 9e830cb5..59058f53 100644
--- a/dbt/adapters/redshift/connections.py
+++ b/dbt/adapters/redshift/connections.py
@@ -7,6 +7,8 @@ from dataclasses import dataclass, field
 import agate
 import sqlparse
 import redshift_connector
+import urllib.request
+import json
 from redshift_connector.utils.oids import get_datatype_name
 
 from dbt.adapters.sql import SQLConnectionManager
@@ -17,13 +19,10 @@ import dbt.flags
 from dbt.dataclass_schema import FieldEncoder, dbtClassMixin, StrEnum
 from dbt.helper_types import Port
 
-
 logger = AdapterLogger("Redshift")
 
-
 drop_lock: Lock = dbt.flags.MP_CONTEXT.Lock()  # type: ignore
 
-
 IAMDuration = NewType("IAMDuration", int)
 
 
@@ -36,6 +35,23 @@ class IAMDurationEncoder(FieldEncoder):
 dbtClassMixin.register_field_encoders({IAMDuration: IAMDurationEncoder()})
 
 
+def _get_aws_regions():
+    # Extract the prefixes from the AWS IP ranges JSON to determine the available regions
+    url = "https://ip-ranges.amazonaws.com/ip-ranges.json"
+    response = urllib.request.urlopen(url)
+    data = json.loads(response.read().decode())
+    regions = set()
+
+    for prefix in data["prefixes"]:
+        if prefix["service"] == "AMAZON":
+            regions.add(prefix["region"])
+
+    return regions
+
+
+_AVAILABLE_AWS_REGIONS = _get_aws_regions()
+
+
 class RedshiftConnectionMethod(StrEnum):
     DATABASE = "database"
     IAM = "iam"
@@ -60,6 +76,7 @@ class RedshiftCredentials(Credentials):
     role: Optional[str] = None
     sslmode: Optional[str] = None
     retries: int = 1
+    region: Optional[str] = None  # if not provided, will be determined from host
 
     _ALIASES = {"dbname": "database", "pass": "password"}
 
@@ -78,6 +95,7 @@ class RedshiftCredentials(Credentials):
             "cluster_id",
             "iam_profile",
             "sslmode",
+            "region",
         )
 
     @property
@@ -85,6 +103,13 @@ class RedshiftCredentials(Credentials):
         return self.host
 
 
+def _is_valid_region(region):
+    if region is None or len(region) == 0:
+        logger.warning("Couldn't determine AWS regions. Skipping validation to avoid blocking.")
+        return True
+    return region in _AVAILABLE_AWS_REGIONS
+
+
 class RedshiftConnectMethodFactory:
     credentials: RedshiftCredentials
 
@@ -99,9 +124,27 @@ class RedshiftConnectMethodFactory:
             "port": self.credentials.port if self.credentials.port else 5439,
             "auto_create": self.credentials.autocreate,
             "db_groups": self.credentials.db_groups,
-            "region": self.credentials.host.split(".")[2],
+            "region": self.credentials.region,
             "timeout": self.credentials.connect_timeout,
         }
+        if kwargs["region"] is None:
+            logger.debug("No region provided, attempting to determine from host.")
+            try:
+                region_value = self.credentials.host.split(".")[2]
+            except IndexError:
+                raise dbt.exceptions.FailedToConnectError(
+                    "No region provided and unable to determine region from host: "
+                    "{}".format(self.credentials.host)
+                )
+
+            kwargs["region"] = region_value
+
+        # Validate the set region
+        if not _is_valid_region(kwargs["region"]):
+            raise dbt.exceptions.FailedToConnectError(
+                "Invalid region provided: {}".format(kwargs["region"])
+            )
+
         if self.credentials.sslmode:
             kwargs["sslmode"] = self.credentials.sslmode
 
@@ -117,7 +160,9 @@ class RedshiftConnectMethodFactory:
             def connect():
                 logger.debug("Connecting to redshift with username/password based auth...")
                 c = redshift_connector.connect(
-                    user=self.credentials.user, password=self.credentials.password, **kwargs
+                    user=self.credentials.user,
+                    password=self.credentials.password,
+                    **kwargs,
                 )
                 if self.credentials.role:
                     c.cursor().execute("set role {}".format(self.credentials.role))


Test patch:
diff --git a/tests/unit/test_redshift_adapter.py b/tests/unit/test_redshift_adapter.py
index 27bcd98f..25f3b3d7 100644
--- a/tests/unit/test_redshift_adapter.py
+++ b/tests/unit/test_redshift_adapter.py
@@ -97,7 +97,9 @@ class TestRedshiftAdapter(unittest.TestCase):
     @mock.patch("redshift_connector.connect", Mock())
     def test_explicit_iam_conn_without_profile(self):
         self.config.credentials = self.config.credentials.replace(
-            method="iam", cluster_id="my_redshift", host="thishostshouldnotexist.test.us-east-1"
+            method="iam",
+            cluster_id="my_redshift",
+            host="thishostshouldnotexist.test.us-east-1",
         )
         connection = self.adapter.acquire_connection("dummy")
         connection.handle
@@ -171,6 +173,92 @@ class TestRedshiftAdapter(unittest.TestCase):
             port=5439,
         )
 
+    @mock.patch("redshift_connector.connect", Mock())
+    @mock.patch("boto3.Session", Mock())
+    def test_explicit_region(self):
+        # Successful test
+        self.config.credentials = self.config.credentials.replace(
+            method="iam",
+            iam_profile="test",
+            host="doesnotexist.1233.redshift-serverless.amazonaws.com",
+            region="us-east-2",
+        )
+        connection = self.adapter.acquire_connection("dummy")
+        connection.handle
+        redshift_connector.connect.assert_called_once_with(
+            iam=True,
+            host="doesnotexist.1233.redshift-serverless.amazonaws.com",
+            database="redshift",
+            cluster_identifier=None,
+            region="us-east-2",
+            auto_create=False,
+            db_groups=[],
+            db_user="root",
+            password="",
+            user="",
+            profile="test",
+            timeout=30,
+            port=5439,
+        )
+
+    @mock.patch("redshift_connector.connect", Mock())
+    @mock.patch("boto3.Session", Mock())
+    def test_explicit_region_failure(self):
+        # Failure test with no region
+        self.config.credentials = self.config.credentials.replace(
+            method="iam",
+            iam_profile="test",
+            host="doesnotexist.1233_no_region",
+            region=None,
+        )
+
+        with self.assertRaises(dbt.exceptions.FailedToConnectError):
+            connection = self.adapter.acquire_connection("dummy")
+            connection.handle
+            redshift_connector.connect.assert_called_once_with(
+                iam=True,
+                host="doesnotexist.1233_no_region",
+                database="redshift",
+                cluster_identifier=None,
+                auto_create=False,
+                db_groups=[],
+                db_user="root",
+                password="",
+                user="",
+                profile="test",
+                timeout=30,
+                port=5439,
+            )
+
+    @mock.patch("redshift_connector.connect", Mock())
+    @mock.patch("boto3.Session", Mock())
+    def test_explicit_invalid_region(self):
+        # Invalid region test
+        self.config.credentials = self.config.credentials.replace(
+            method="iam",
+            iam_profile="test",
+            host="doesnotexist.1233_no_region.us-not-a-region-1",
+            region=None,
+        )
+
+        with self.assertRaises(dbt.exceptions.FailedToConnectError):
+            connection = self.adapter.acquire_connection("dummy")
+            connection.handle
+            redshift_connector.connect.assert_called_once_with(
+                iam=True,
+                host="doesnotexist.1233_no_region",
+                database="redshift",
+                cluster_identifier=None,
+                auto_create=False,
+                db_groups=[],
+                db_user="root",
+                password="",
+                user="",
+                profile="test",
+                timeout=30,
+                port=5439,
+            )
+
     @mock.patch("redshift_connector.connect", Mock())
     @mock.patch("boto3.Session", Mock())
     def test_serverless_iam_failure(self):
@@ -263,7 +351,10 @@ class TestRedshiftAdapter(unittest.TestCase):
 
             self.assertEqual(len(list(self.adapter.cancel_open_connections())), 1)
             add_query.assert_has_calls(
-                [call("select pg_backend_pid()"), call("select pg_terminate_backend(42)")]
+                [
+                    call("select pg_backend_pid()"),
+                    call("select pg_terminate_backend(42)"),
+                ]
             )
 
         master.handle.get_backend_pid.assert_not_called()