Skip to content

Commit

Permalink
fork abbywh branch back into repo
Browse files Browse the repository at this point in the history
  • Loading branch information
mikealfare committed Apr 23, 2024
2 parents 9cc7a7c + d574ef4 commit b8e2185
Show file tree
Hide file tree
Showing 5 changed files with 98 additions and 1 deletion.
6 changes: 6 additions & 0 deletions .changes/unreleased/Features-20230921-153707.yaml
Original file line number Diff line number Diff line change
@@ -0,0 +1,6 @@
kind: Features
body: Support storing test failures as views
time: 2023-09-21T15:37:07.970722-04:00
custom:
Author: mikealfare
Issue: "6914"
6 changes: 6 additions & 0 deletions .changes/unreleased/Fixes-20230807-174409.yaml
Original file line number Diff line number Diff line change
@@ -0,0 +1,6 @@
kind: Fixes
body: Use the PID to terminate the session
time: 2023-08-07T17:44:09.15097-06:00
custom:
Author: dbeatty10
Issue: "553"
31 changes: 30 additions & 1 deletion dbt/adapters/redshift/connections.py
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,7 @@ def get_message(self) -> str:
class RedshiftConnectionMethod(StrEnum):
DATABASE = "database"
IAM = "iam"
IAMR = "iamr"


class UserSSLMode(StrEnum):
Expand Down Expand Up @@ -102,9 +103,9 @@ def parse(cls, user_sslmode: UserSSLMode) -> "RedshiftSSLConfig":
@dataclass
class RedshiftCredentials(Credentials):
host: str
user: str
port: Port
method: str = RedshiftConnectionMethod.DATABASE # type: ignore
user: Optional[str] = None # type: ignore
password: Optional[str] = None # type: ignore
cluster_id: Optional[str] = field(
default=None,
Expand Down Expand Up @@ -187,6 +188,9 @@ def get_connect_method(self):
"'password' field is required for 'database' credentials"
)

if self.credentials.user is None:
raise FailedToConnectError("'user' field is required for 'database' credentials")

def connect():
logger.debug("Connecting to redshift with username/password based auth...")
c = redshift_connector.connect(
Expand All @@ -207,6 +211,9 @@ def connect():
"'host' must be provided for serverless endpoint."
)

if self.credentials.user is None:
raise FailedToConnectError("'user' field is required for 'iam' credentials")

def connect():
logger.debug("Connecting to redshift with IAM based auth...")
c = redshift_connector.connect(
Expand All @@ -224,6 +231,28 @@ def connect():
c.cursor().execute("set role {}".format(self.credentials.role))
return c

elif method == RedshiftConnectionMethod.IAMR:
if not self.credentials.cluster_id and "serverless" not in self.credentials.host:
raise FailedToConnectError(
"Failed to use IAM method. 'cluster_id' must be provided for provisioned cluster. "
"'host' must be provided for serverless endpoint."
)

def connect():
logger.debug("Connecting to redshift with IAM based auth...")
c = redshift_connector.connect(
iam=True,
cluster_identifier=self.credentials.cluster_id,
profile=self.credentials.iam_profile,
group_federation=True,
**kwargs,
)
if self.credentials.autocommit:
c.autocommit = True
if self.credentials.role:
c.cursor().execute("set role {}".format(self.credentials.role))
return c

else:
raise FailedToConnectError("Invalid 'method' in profile: '{}'".format(method))

Expand Down
2 changes: 2 additions & 0 deletions dbt/include/redshift/profile_template.yml
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,8 @@ prompts:
hide_input: true
iam:
_fixed_method: iam
iamr:
_fixed_method: iamr
dbname:
hint: 'default database that dbt will build objects in'
schema:
Expand Down
54 changes: 54 additions & 0 deletions tests/unit/test_redshift_adapter.py
Original file line number Diff line number Diff line change
Expand Up @@ -202,6 +202,60 @@ def test_explicit_iam_serverless_with_profile(self):
)

@mock.patch("redshift_connector.connect", MagicMock())
def test_explicit_iamr_conn_without_profile(self):
self.config.credentials = self.config.credentials.replace(
method="iamr",
cluster_id="my_redshift",
host="thishostshouldnotexist.test.us-east-1",
user=None,
)
connection = self.adapter.acquire_connection("dummy")
connection.handle
redshift_connector.connect.assert_called_once_with(
iam=True,
host="thishostshouldnotexist.test.us-east-1",
database="redshift",
cluster_identifier="my_redshift",
region=None,
timeout=None,
auto_create=False,
db_groups=[],
profile=None,
port=5439,
group_federation=True,
**DEFAULT_SSL_CONFIG,
)

@mock.patch("redshift_connector.connect", Mock())
@mock.patch("boto3.Session", Mock())
def test_explicit_iamr_conn_with_profile(self):
self.config.credentials = self.config.credentials.replace(
method="iamr",
cluster_id="my_redshift",
iam_profile="test",
host="thishostshouldnotexist.test.us-east-1",
user=None,
)
connection = self.adapter.acquire_connection("dummy")
connection.handle

redshift_connector.connect.assert_called_once_with(
iam=True,
host="thishostshouldnotexist.test.us-east-1",
database="redshift",
cluster_identifier="my_redshift",
region=None,
auto_create=False,
db_groups=[],
profile="test",
timeout=None,
port=5439,
group_federation=True,
**DEFAULT_SSL_CONFIG,
)

@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(
Expand Down

0 comments on commit b8e2185

Please sign in to comment.