Fix query exception handling (microsoft#244)

* added new abstract driver property query_canceled_error to access psycopg2.extensions.QueryCanceledError
* altered query exception handling to use abstracted driver properties
* fixed batch execute function to avoid "cursor does not exist" operational error

Author: KevinMicrosoft

pgsqltoolsservice/driver/types/driver.py

@@ -73,6 +73,13 @@ def database_error(self) -> str:
     @property
     @abstractmethod
     def transaction_in_error(self) -> bool:
+        """Returns bool indicating if transaction is in error"""
+        pass
+
+    @property
+    @abstractmethod
+    def query_canceled_error(self) -> Exception:
+        """Returns driver query canceled error"""
         pass
 
     @property

pgsqltoolsservice/driver/types/psycopg_driver.py

@@ -131,13 +131,19 @@ def default_database(self):
 
     @property
     def database_error(self):
-        """ Returns the type of database error this connection throws"""
+        """Returns the type of database error this connection throws"""
         return self._database_error
 
     @property
     def transaction_in_error(self) -> bool:
+        """Returns bool indicating if transaction is in error"""
         return self._conn.get_transaction_status() is psycopg2.extensions.TRANSACTION_STATUS_INERROR
 
+    @property
+    def query_canceled_error(self) -> Exception:
+        """Returns driver query canceled error"""
+        return psycopg2.extensions.QueryCanceledError
+
     @property
     def cancellation_query(self) -> str:
         backend_pid = self._conn.get_backend_pid()

pgsqltoolsservice/driver/types/pymysql_driver.py

@@ -166,6 +166,11 @@ def database_error(self):
     def transaction_in_error(self) -> int:
         pass
 
+    @property
+    def query_canceled_error(self) -> Exception:
+        """Returns query canceled error type"""
+        pass
+
     @property
     def cancellation_query(self) -> str:
         # TODO generate a query that kills the current query process

pgsqltoolsservice/query/batch.py

@@ -118,33 +118,33 @@ def execute(self, conn: ServerConnection) -> None:
         if self._batch_events and self._batch_events._on_execution_started:
             self._batch_events._on_execution_started(self)
 
-        with self.get_cursor(conn) as cursor:
-            try:
-                cursor.execute(self.batch_text)
-                # Commit the transaction if autocommit is True
-                if conn.autocommit:
-                    conn.commit()
-
-                self.after_execute(cursor)
-            except conn.database_error as error:
-                self._has_error = True
-                # We just raise the error with primary message and not the cursor stacktrace
-                raise error
-            finally:
-                # We are doing this because when the execute fails for named cursors
-                # cursor is not activated on the server which results in failure on close
-                # Hence we are checking if the cursor was really executed for us to close it
-                if cursor and cursor.rowcount != -1 and cursor.rowcount is not None:
-                    cursor.close()
-                self._has_executed = True
-                self._execution_end_time = datetime.now()
-                
-                # TODO: PyMySQL doesn't support notices from a connection
-                # self._notices = _get_notices(cursor.connection).notices
-                # cursor.connection.notices = []
-
-                if self._batch_events and self._batch_events._on_execution_completed:
-                    self._batch_events._on_execution_completed(self)
+        try:
+            cursor = self.get_cursor(conn)
+            cursor.execute(self.batch_text)
+            # Commit the transaction if autocommit is True
+            if conn.autocommit:
+                conn.commit()
+
+            self.after_execute(cursor)
+        except conn.database_error as error:
+            self._has_error = True
+            # We just raise the error with primary message and not the cursor stacktrace
+            raise conn.database_error(error.diag.message_primary) from error
+        finally:
+            # We are doing this because when the execute fails for named cursors
+            # cursor is not activated on the server which results in failure on close
+            # Hence we are checking if the cursor was really executed for us to close it
+            if cursor and cursor.rowcount != -1 and cursor.rowcount is not None:
+                cursor.close()
+            self._has_executed = True
+            self._execution_end_time = datetime.now()
+            
+            # TODO: PyMySQL doesn't support notices from a connection
+            # self._notices = _get_notices(cursor.connection).notices
+            # cursor.connection.notices = []
+
+            if self._batch_events and self._batch_events._on_execution_completed:
+                self._batch_events._on_execution_completed(self)
 
     def after_execute(self, cursor) -> None:
         if cursor.description is not None:

pgsqltoolsservice/query_execution/query_execution_service.py

@@ -417,7 +417,7 @@ def _get_query_text_from_execute_params(self, params: ExecuteRequestParamsBase):
 
     def _resolve_query_exception(self, e: Exception, query: Query, worker_args: ExecuteRequestWorkerArgs, is_rollback_error=False):
         utils.log.log_debug(self._service_provider.logger, f'Query execution failed for following query: {query.query_text}\n {e}')
-        if isinstance(e, conn.database_error) or isinstance(e, RuntimeError):
+        if isinstance(e, worker_args.connection.database_error) or isinstance(e, RuntimeError) or isinstance(e, worker_args.connection.query_canceled_error):
             error_message = str(e)
         else:
             error_message = 'Unhandled exception while executing query: {}'.format(str(e))  # TODO: Localize
@@ -434,7 +434,7 @@ def _resolve_query_exception(self, e: Exception, query: Query, worker_args: Exec
 
         # If there was a failure in the middle of a transaction, roll it back.
         # Note that conn.rollback() won't work since the connection is in autocommit mode
-        if not is_rollback_error and worker_args.connection.get_transaction_status() is psycopg2.extensions.TRANSACTION_STATUS_INERROR:
+        if not is_rollback_error and worker_args.connection.transaction_in_error:
             rollback_query = Query(query.owner_uri, 'ROLLBACK', QueryExecutionSettings(ExecutionPlanOptions(), None), QueryEvents())
             try:
                 rollback_query.execute(worker_args.connection)