Improve Connection.cancel() readme and tests (#418)

README.md

@@ -549,6 +549,8 @@ def cancel_query(connection, timeout=5):
     time.sleep(timeout)
     connection.cancel()
 
+# Example 1: Cancel the query before Cursor.execute() return.
+#            The query stops executing in a shorter time after the cancel message is sent.
 with vertica_python.connect(**conn_info) as conn:
     cur = conn.cursor()
 
@@ -562,9 +564,26 @@ with vertica_python.connect(**conn_info) as conn:
         pass
 
     p1.join()
+
+# Example 2: Cancel the query after Cursor.execute() return.
+#            Less number of rows read after the cancel message is sent.
+with vertica_python.connect(**conn_info) as conn:
+    cur = conn.cursor()
+    cur.execute("SELECT id, time FROM large_table")
+    nCount = 0
+    try:
+        while cur.fetchone():
+            nCount += 1
+            if nCount == 100:
+                conn.cancel()
+    except vertica_python.errors.QueryCanceled as e:
+        pass
+        # nCount is less than the number of rows in large_table
+
 ```
 
 
+
 ## Rowcount oddities
 
 vertica_python behaves a bit differently than dbapi when returning rowcounts.

vertica_python/tests/integration_tests/test_cancel.py

@@ -55,5 +55,49 @@ def cancel_query(conn, delay=5):
                 cur.execute(long_running_query)
             p1.join()
 
+            # Must be able to successfully run next query
+            cur.execute("SELECT 1")
+            res = cur.fetchall()
+            self.assertListOfListsEqual(res, [[1]])
+
+    def test_connection_cancel_returned_query(self):
+        with self._connect() as conn:
+            cur = conn.cursor()
+            cur.execute("DROP TABLE IF EXISTS vptest")
+            try:
+                # Creating and loading table
+                cur.execute("CREATE TABLE vptest(id INTEGER, time TIMESTAMP)")
+                cur.execute("INSERT INTO vptest"
+                        " SELECT row_number() OVER(), slice_time"
+                        "    FROM("
+                        "         SELECT slice_time FROM("
+                        "         SELECT '2021-01-01'::timestamp s UNION ALL SELECT '2022-01-01'::timestamp s"
+                        "          ) sq TIMESERIES slice_time AS '1 second' OVER(ORDER BY s)"
+                        "    ) sq2")
+
+                # This query returns over 30,000,000 rows. We cancel the command after
+                # reading 100 of them, and then continue reading results. This quickly
+                # results in an exception being thrown due to the cancel having taken effect.
+                cur.execute("SELECT id, time FROM vptest")
+
+                nCount = 0
+                with self.assertRaises(errors.QueryCanceled):
+                    while cur.fetchone():
+                        nCount += 1
+                        if nCount == 100:
+                            conn.cancel()
+
+                # The number of rows read after the cancel message is sent to the server can vary.
+                # 100,000 seems to leave a safe margin while still falling well short of
+                # the 30,000,000+ rows we'd have read if the cancel didn't work.
+                self.assertTrue(100 <= nCount < 100000)
+
+                # Must be able to successfully run next query
+                cur.execute("SELECT 1")
+                res = cur.fetchall()
+                self.assertListOfListsEqual(res, [[1]])
+            finally:
+                cur.execute("DROP TABLE IF EXISTS vptest")
+
 
 exec(CancelTestCase.createPrepStmtClass())