From d4e05f605a3b89485d84b263149834d1d6217cf5 Mon Sep 17 00:00:00 2001 From: Abby Whittier Date: Mon, 16 Oct 2023 10:48:16 -0400 Subject: [PATCH 1/5] adding SSO support for redshift Committer: Abby Whittier --- dbt/adapters/redshift/connections.py | 24 +++++++++- dbt/include/redshift/profile_template.yml | 2 + tests/unit/test_redshift_adapter.py | 54 ++++++++++++++++++++++- 3 files changed, 78 insertions(+), 2 deletions(-) diff --git a/dbt/adapters/redshift/connections.py b/dbt/adapters/redshift/connections.py index 83d05b587..055cee662 100644 --- a/dbt/adapters/redshift/connections.py +++ b/dbt/adapters/redshift/connections.py @@ -39,6 +39,7 @@ def get_message(self) -> str: class RedshiftConnectionMethod(StrEnum): DATABASE = "database" IAM = "iam" + IAMR = "iamr" class UserSSLMode(StrEnum): @@ -104,9 +105,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, @@ -226,6 +227,27 @@ 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 dbt.exceptions.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, + **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 dbt.exceptions.FailedToConnectError( "Invalid 'method' in profile: '{}'".format(method) diff --git a/dbt/include/redshift/profile_template.yml b/dbt/include/redshift/profile_template.yml index 41f33e87e..5f6b0a91a 100644 --- a/dbt/include/redshift/profile_template.yml +++ b/dbt/include/redshift/profile_template.yml @@ -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: diff --git a/tests/unit/test_redshift_adapter.py b/tests/unit/test_redshift_adapter.py index c31366a1e..a531035ab 100644 --- a/tests/unit/test_redshift_adapter.py +++ b/tests/unit/test_redshift_adapter.py @@ -12,7 +12,10 @@ ) from dbt.clients import agate_helper from dbt.exceptions import FailedToConnectError -from dbt.adapters.redshift.connections import RedshiftConnectMethodFactory, RedshiftSSLConfig +from dbt.adapters.redshift.connections import ( + RedshiftConnectMethodFactory, + RedshiftSSLConfig, +) from .utils import ( config_from_parts_or_dicts, mock_connection, @@ -199,6 +202,55 @@ def test_explicit_iam_serverless_with_profile(self): **DEFAULT_SSL_CONFIG, ) + @mock.patch("redshift_connector.connect", Mock()) + 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", + ) + 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, + **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", + ) + 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, + **DEFAULT_SSL_CONFIG, + ) + @mock.patch("redshift_connector.connect", Mock()) @mock.patch("boto3.Session", Mock()) def test_explicit_region(self): From a8ceb848531d61c7a46abe3269cecf8fd825c5b5 Mon Sep 17 00:00:00 2001 From: Mike Alfare <13974384+mikealfare@users.noreply.github.com> Date: Tue, 10 Oct 2023 19:05:42 -0400 Subject: [PATCH 2/5] ADAP-891: Support test results as views (#614) * implement store-failures-as tests --- .../unreleased/Features-20230921-153707.yaml | 6 +++++ .../adapter/test_store_test_failures.py | 27 ++++++++++++++++++- 2 files changed, 32 insertions(+), 1 deletion(-) create mode 100644 .changes/unreleased/Features-20230921-153707.yaml diff --git a/.changes/unreleased/Features-20230921-153707.yaml b/.changes/unreleased/Features-20230921-153707.yaml new file mode 100644 index 000000000..cf2c60ad5 --- /dev/null +++ b/.changes/unreleased/Features-20230921-153707.yaml @@ -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" diff --git a/tests/functional/adapter/test_store_test_failures.py b/tests/functional/adapter/test_store_test_failures.py index 5d6b70fbb..7f591654e 100644 --- a/tests/functional/adapter/test_store_test_failures.py +++ b/tests/functional/adapter/test_store_test_failures.py @@ -1,7 +1,32 @@ +from dbt.tests.adapter.store_test_failures_tests import basic from dbt.tests.adapter.store_test_failures_tests.test_store_test_failures import ( TestStoreTestFailures, ) -class RedshiftTestStoreTestFailures(TestStoreTestFailures): +class TestRedshiftTestStoreTestFailures(TestStoreTestFailures): + pass + + +class TestStoreTestFailuresAsInteractions(basic.StoreTestFailuresAsInteractions): + pass + + +class TestStoreTestFailuresAsProjectLevelOff(basic.StoreTestFailuresAsProjectLevelOff): + pass + + +class TestStoreTestFailuresAsProjectLevelView(basic.StoreTestFailuresAsProjectLevelView): + pass + + +class TestStoreTestFailuresAsGeneric(basic.StoreTestFailuresAsGeneric): + pass + + +class TestStoreTestFailuresAsProjectLevelEphemeral(basic.StoreTestFailuresAsProjectLevelEphemeral): + pass + + +class TestStoreTestFailuresAsExceptions(basic.StoreTestFailuresAsExceptions): pass From 48128e2fbd198b6b60a22e4ddf444831cf5feae6 Mon Sep 17 00:00:00 2001 From: Doug Beatty <44704949+dbeatty10@users.noreply.github.com> Date: Tue, 10 Oct 2023 19:15:20 -0600 Subject: [PATCH 3/5] Use the PID to terminate the session (#568) * The first element of the result is the PID * Debug-level logging of high-level message + SQL * Using redshift_connector `cursor.fetchone()` returns `(,)` * Use cursor to call `select pg_terminate_backend({pid})` directly rather than using the `SQLConnectionManager` --------- Co-authored-by: Mike Alfare <13974384+mikealfare@users.noreply.github.com> --- .changes/unreleased/Fixes-20230807-174409.yaml | 6 ++++++ dbt/adapters/redshift/connections.py | 10 ++++++---- tests/unit/test_redshift_adapter.py | 3 +-- 3 files changed, 13 insertions(+), 6 deletions(-) create mode 100644 .changes/unreleased/Fixes-20230807-174409.yaml diff --git a/.changes/unreleased/Fixes-20230807-174409.yaml b/.changes/unreleased/Fixes-20230807-174409.yaml new file mode 100644 index 000000000..ad4f11b42 --- /dev/null +++ b/.changes/unreleased/Fixes-20230807-174409.yaml @@ -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" diff --git a/dbt/adapters/redshift/connections.py b/dbt/adapters/redshift/connections.py index 055cee662..375e2244a 100644 --- a/dbt/adapters/redshift/connections.py +++ b/dbt/adapters/redshift/connections.py @@ -262,8 +262,9 @@ class RedshiftConnectionManager(SQLConnectionManager): def _get_backend_pid(self): sql = "select pg_backend_pid()" _, cursor = self.add_query(sql) + res = cursor.fetchone() - return res + return res[0] def cancel(self, connection: Connection): try: @@ -275,9 +276,10 @@ def cancel(self, connection: Connection): raise sql = f"select pg_terminate_backend({pid})" - _, cursor = self.add_query(sql) - res = cursor.fetchone() - logger.debug(f"Cancel query '{connection.name}': {res}") + cursor = connection.handle.cursor() + logger.debug(f"Cancel query on: '{connection.name}' with PID: {pid}") + logger.debug(sql) + cursor.execute(sql) @classmethod def get_response(cls, cursor: redshift_connector.Cursor) -> AdapterResponse: diff --git a/tests/unit/test_redshift_adapter.py b/tests/unit/test_redshift_adapter.py index a531035ab..f1752dbd3 100644 --- a/tests/unit/test_redshift_adapter.py +++ b/tests/unit/test_redshift_adapter.py @@ -523,14 +523,13 @@ def test_cancel_open_connections_single(self): with mock.patch.object(self.adapter.connections, "add_query") as add_query: query_result = mock.MagicMock() cursor = mock.Mock() - cursor.fetchone.return_value = 42 + cursor.fetchone.return_value = (42,) add_query.side_effect = [(None, cursor), (None, query_result)] 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)"), ] ) From dcc87c4786dc84762742a2693522484e269cf023 Mon Sep 17 00:00:00 2001 From: Abby Whittier Date: Mon, 16 Oct 2023 13:03:36 -0400 Subject: [PATCH 4/5] added error checking for new optional user field --- dbt/adapters/redshift/connections.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/dbt/adapters/redshift/connections.py b/dbt/adapters/redshift/connections.py index 375e2244a..34952906e 100644 --- a/dbt/adapters/redshift/connections.py +++ b/dbt/adapters/redshift/connections.py @@ -190,6 +190,11 @@ def get_connect_method(self): "'password' field is required for 'database' credentials" ) + if self.credentials.user is None: + raise dbt.exceptions.FailedToConnectError( + "'user' field is required for 'database' credentials" + ) + def connect(): logger.debug("Connecting to redshift with username/password based auth...") c = redshift_connector.connect( @@ -210,6 +215,11 @@ def connect(): "'host' must be provided for serverless endpoint." ) + if self.credentials.user is None: + raise dbt.exceptions.FailedToConnectError( + "'user' field is required for 'iam' credentials" + ) + def connect(): logger.debug("Connecting to redshift with IAM based auth...") c = redshift_connector.connect( From a60743c88c72f9373f07922cd3efbadd499e4f67 Mon Sep 17 00:00:00 2001 From: Abby Whittier Date: Thu, 19 Oct 2023 10:26:34 -0400 Subject: [PATCH 5/5] black formatting --- dbt/adapters/redshift/connections.py | 1 + tests/unit/test_redshift_adapter.py | 4 ++++ 2 files changed, 5 insertions(+) diff --git a/dbt/adapters/redshift/connections.py b/dbt/adapters/redshift/connections.py index 34952906e..5d7c9fe05 100644 --- a/dbt/adapters/redshift/connections.py +++ b/dbt/adapters/redshift/connections.py @@ -250,6 +250,7 @@ def connect(): iam=True, cluster_identifier=self.credentials.cluster_id, profile=self.credentials.iam_profile, + group_federation=True, **kwargs, ) if self.credentials.autocommit: diff --git a/tests/unit/test_redshift_adapter.py b/tests/unit/test_redshift_adapter.py index f1752dbd3..97e9e8b06 100644 --- a/tests/unit/test_redshift_adapter.py +++ b/tests/unit/test_redshift_adapter.py @@ -208,6 +208,7 @@ def test_explicit_iamr_conn_without_profile(self): method="iamr", cluster_id="my_redshift", host="thishostshouldnotexist.test.us-east-1", + user=None, ) connection = self.adapter.acquire_connection("dummy") connection.handle @@ -222,6 +223,7 @@ def test_explicit_iamr_conn_without_profile(self): db_groups=[], profile=None, port=5439, + group_federation=True, **DEFAULT_SSL_CONFIG, ) @@ -233,6 +235,7 @@ def test_explicit_iamr_conn_with_profile(self): cluster_id="my_redshift", iam_profile="test", host="thishostshouldnotexist.test.us-east-1", + user=None, ) connection = self.adapter.acquire_connection("dummy") connection.handle @@ -248,6 +251,7 @@ def test_explicit_iamr_conn_with_profile(self): profile="test", timeout=None, port=5439, + group_federation=True, **DEFAULT_SSL_CONFIG, )