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] 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 `(<something>,)`
* 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)"),
                 ]
             )