Implement query cancellation support within MCPKit, and add it to the task functionality:

 - MCPKit now supports cancelling the active query; for MySQL servers >= 5.0.0 a query kill is attempted from a new connection, and if that fails or for MySQL < 5 a reconnect is triggered.
 - TableDocument now supports enabling a cancel task button on the task interface, including an optional callback
 - Implement query cancellation for custom queries.  This addresses Issue #86.
 - Implement query cancellation for table content loads, filters, and sorts.

6 files changed, 168 insertions, 39 deletions

diff --git a/Source/CustomQuery.m b/Source/CustomQuery.m
index f515f335..dc0d9b76 100644
--- a/Source/CustomQuery.m
+++ b/Source/CustomQuery.m
@@ -388,6 +388,7 @@
 	MCPStreamingResult	*streamingResult  = nil;
 	NSMutableString		*errors     = [NSMutableString string];
 	SEL					callbackMethod = NULL;
+	NSString			*taskButtonString;
 	
 	int i, totalQueriesRun = 0, totalAffectedRows = 0;
 	double executionTime = 0;
@@ -425,6 +426,13 @@
 	long queryCount = [queries count];
 	NSMutableArray *tempQueries = [NSMutableArray arrayWithCapacity:queryCount];
 
+	// Enable task cancellation
+	if (queryCount > 1)
+		taskButtonString = NSLocalizedString(@"Stop queries", @"Stop queries string");
+	else
+		taskButtonString = NSLocalizedString(@"Stop query", @"Stop query string");
+	[tableDocumentInstance enableTaskCancellationWithTitle:taskButtonString callbackObject:nil callbackFunction:NULL];
+
 	// Perform the supplied queries in series
 	for ( i = 0 ; i < queryCount ; i++ ) {
 
@@ -450,7 +458,7 @@
 
 		// If this is the last query, retrieve and store the result; otherwise,
 		// discard the result without fully loading.
-		if (totalQueriesRun == queryCount) {
+		if (totalQueriesRun == queryCount || [mySQLConnection queryCancelled]) {
 
 			// get column definitions for the result array
 			if (cqColumnDefinition) [cqColumnDefinition release];
@@ -507,7 +515,17 @@
 			totalAffectedRows += [streamingResult numOfRows];
 
 		// Store any error messages
-		if ( ![[mySQLConnection getLastErrorMessage] isEqualToString:@""] ) {
+		if ( ![[mySQLConnection getLastErrorMessage] isEqualToString:@""] || [mySQLConnection queryCancelled]) {
+
+			NSString *errorString;
+			if ([mySQLConnection queryCancelled]) {
+				if ([mySQLConnection queryCancellationUsedReconnect])
+					errorString = NSLocalizedString(@"Query cancelled.  Please note that to cancel the query the connection had to be reset; transactions and connection variables were reset.", @"Query cancel by resetting connection error");
+				else
+					errorString = NSLocalizedString(@"Query cancelled.", @"Query cancelled error");
+			} else {
+				errorString = [mySQLConnection getLastErrorMessage];
+			}
 
 			// If the query errored, append error to the error log for display at the end
 			if ( queryCount > 1 ) {
@@ -519,35 +537,37 @@
 					// Update error text for the user
 					[errors appendString:[NSString stringWithFormat:NSLocalizedString(@"[ERROR in query %d] %@\n", @"error text when multiple custom query failed"),
 										i+1,
-										[mySQLConnection getLastErrorMessage]]];
+										errorString]];
 					[errorText setStringValue:errors];
+
 					// ask the user to continue after detecting an error
-					NSAlert *alert = [[[NSAlert alloc] init] autorelease];
-					[alert addButtonWithTitle:NSLocalizedString(@"Run All", @"run all button")];
-					[alert addButtonWithTitle:NSLocalizedString(@"Continue", @"continue button")];
-					[alert addButtonWithTitle:NSLocalizedString(@"Stop", @"stop button")];
-					[alert setMessageText:NSLocalizedString(@"MySQL Error", @"mysql error message")];
-					[alert setInformativeText:[mySQLConnection getLastErrorMessage]];
-					[alert setAlertStyle:NSWarningAlertStyle];
-					int choice = [alert runModal];
-					switch (choice){
-						case NSAlertFirstButtonReturn:
-							suppressErrorSheet = YES;
-						case NSAlertSecondButtonReturn:
-							break;
-						default:
-							if(i < queryCount-1) // output that message only if it was not the last one
-								[errors appendString:NSLocalizedString(@"Execution stopped!\n", @"execution stopped message")];
-							i = queryCount; // break for loop; for safety reasons stop the execution of the following queries
+					if (![mySQLConnection queryCancelled]) {
+						NSAlert *alert = [[[NSAlert alloc] init] autorelease];
+						[alert addButtonWithTitle:NSLocalizedString(@"Run All", @"run all button")];
+						[alert addButtonWithTitle:NSLocalizedString(@"Continue", @"continue button")];
+						[alert addButtonWithTitle:NSLocalizedString(@"Stop", @"stop button")];
+						[alert setMessageText:NSLocalizedString(@"MySQL Error", @"mysql error message")];
+						[alert setInformativeText:[mySQLConnection getLastErrorMessage]];
+						[alert setAlertStyle:NSWarningAlertStyle];
+						int choice = [alert runModal];
+						switch (choice){
+							case NSAlertFirstButtonReturn:
+								suppressErrorSheet = YES;
+							case NSAlertSecondButtonReturn:
+								break;
+							default:
+								if(i < queryCount-1) // output that message only if it was not the last one
+									[errors appendString:NSLocalizedString(@"Execution stopped!\n", @"execution stopped message")];
+								i = queryCount; // break for loop; for safety reasons stop the execution of the following queries
+						}
 					}
-
 				} else {
 					[errors appendString:[NSString stringWithFormat:NSLocalizedString(@"[ERROR in query %d] %@\n", @"error text when multiple custom query failed"),
 											i+1,
-											[mySQLConnection getLastErrorMessage]]];
+											errorString]];
 				}
 			} else {
-				[errors setString:[mySQLConnection getLastErrorMessage]];
+				[errors setString:errorString];
 			}
 		} else {
 			// Check if table/db list needs an update
@@ -557,6 +577,9 @@
 			if(!databaseWasChanged && [query isMatchedByRegex:@"(?i)\\b(use|drop\\s+database|drop\\s+schema)\\b\\s+."])
 				databaseWasChanged = YES;
 		}
+
+		// If the query was cancelled, end all queries.
+		if ([mySQLConnection queryCancelled]) break;
 	}
 
 	// Reload table list if at least one query began with drop, alter, rename, or create
@@ -605,7 +628,7 @@
 	}
 
 	// Error checking
-	if ( [errors length] && !queryIsTableSorter ) {
+	if ( [mySQLConnection queryCancelled] || ([errors length] && !queryIsTableSorter)) {
 		// set the error text
 		[errorText setStringValue:errors];
 		// select the line x of the first error if error message contains "at line x"
@@ -656,7 +679,18 @@
 	}
 	
 	// Set up the status string
-	if ( totalQueriesRun > 1 ) {
+	if ( [mySQLConnection queryCancelled] ) {
+		if (totalQueriesRun > 1) {
+			[affectedRowsText setStringValue:[NSString stringWithFormat:NSLocalizedString(@"Cancelled in query %i, after %@", @"text showing multiple queries were cancelled"),
+                                              totalQueriesRun,
+                                              [NSString stringForTimeInterval:executionTime]
+                                              ]];
+		} else {
+			[affectedRowsText setStringValue:[NSString stringWithFormat:NSLocalizedString(@"Cancelled after %@", @"text showing a query was cancelled"),
+                                              [NSString stringForTimeInterval:executionTime]
+                                              ]];
+		}
+	} else if ( totalQueriesRun > 1 ) {
 		if (totalAffectedRows==1) {
 			[affectedRowsText setStringValue:[NSString stringWithFormat:NSLocalizedString(@"1 row affected in total, by %i queries taking %@", @"text showing one row has been affected by multiple queries"),
                                               totalQueriesRun,
diff --git a/Source/TableContent.h b/Source/TableContent.h
index 889a79e5..d0cc0ec3 100644
--- a/Source/TableContent.h
+++ b/Source/TableContent.h
@@ -65,7 +65,7 @@
 	NSString *compareType;
 	NSNumber *sortCol;
 	BOOL isEditingRow, isEditingNewRow, isSavingRow, isDesc, setLimit;
-	BOOL isFiltered, isLimited, maxNumRowsIsEstimate;
+	BOOL isFiltered, isLimited, isInterruptedLoad, maxNumRowsIsEstimate;
 	NSUserDefaults *prefs;
 	int currentlyEditingRow, maxNumRows;
 
diff --git a/Source/TableContent.m b/Source/TableContent.m
index 552fbd41..d5c9e8d1 100644
--- a/Source/TableContent.m
+++ b/Source/TableContent.m
@@ -86,6 +86,7 @@
 
 		isFiltered = NO;
 		isLimited = NO;
+		isInterruptedLoad = NO;
 		
 		prefs = [NSUserDefaults standardUserDefaults];
 		
@@ -532,7 +533,10 @@
 		rowsToLoad = rowsToLoad - ([limitRowsField intValue]-1);
 		if (rowsToLoad > [prefs integerForKey:SPLimitResultsValue]) rowsToLoad = [prefs integerForKey:SPLimitResultsValue];
 	}
-	
+
+	// If within a task, allow this query to be cancelled
+	[tableDocumentInstance enableTaskCancellationWithTitle:NSLocalizedString(@"Stop", @"Stop load") callbackObject:nil callbackFunction:NULL];
+
 	// Perform and process the query
 	[tableContentView performSelectorOnMainThread:@selector(noteNumberOfRowsChanged) withObject:nil waitUntilDone:YES]; 
 	[self setUsedQuery:queryString];
@@ -541,7 +545,7 @@
 	[streamingResult release];
 
 	// If the result is empty, and a limit is active, reset the limit
-	if ([prefs boolForKey:SPLimitResults] && queryStringBeforeLimit && !tableRowsCount) {
+	if ([prefs boolForKey:SPLimitResults] && queryStringBeforeLimit && !tableRowsCount && ![mySQLConnection queryCancelled]) {
 		[limitRowsField setStringValue:@"1"];
 		queryString = [NSMutableString stringWithFormat:@"%@ LIMIT 0,%d", queryStringBeforeLimit, [prefs integerForKey:SPLimitResultsValue]];
 		[self setUsedQuery:queryString];
@@ -549,6 +553,14 @@
 		[self processResultIntoDataStorage:streamingResult approximateRowCount:[prefs integerForKey:SPLimitResultsValue]];
 		[streamingResult release];
 	}
+	
+	if ([mySQLConnection queryCancelled] || ![[mySQLConnection getLastErrorMessage] isEqualToString:@""])
+		isInterruptedLoad = YES;
+	else 
+		isInterruptedLoad = NO;
+
+	// End cancellation ability
+	[tableDocumentInstance disableTaskCancellation];
 
 	if ([prefs boolForKey:SPLimitResults]
 		&& ([limitRowsField intValue] > 1
@@ -821,8 +833,15 @@
 	NSString *rowString;
 	NSMutableString *countString = [NSMutableString string];
 
+	// If the result is partial due to an error or query cancellation, show a very basic count
+	if (isInterruptedLoad) {
+		if (tableRowsCount == 1)
+			[countString appendFormat:NSLocalizedString(@"%d row in partial load", @"text showing a single row a partially loaded result"), tableRowsCount];
+		else
+			[countString appendFormat:NSLocalizedString(@"%d rows in partial load", @"text showing how many rows are in a partially loaded result"), tableRowsCount];
+		
 	// If no filter or limit is active, show just the count of rows in the table
-	if (!isFiltered && !isLimited) {
+	} else if (!isFiltered && !isLimited) {
 		if (tableRowsCount == 1)
 			[countString appendFormat:NSLocalizedString(@"%d row in table", @"text showing a single row in the result"), tableRowsCount];
 		else
@@ -1128,7 +1147,7 @@
 	
 	NSString *contextInfo = @"removerow";
 	
-	if (([tableContentView numberOfSelectedRows] == [tableContentView numberOfRows]) && !isFiltered && !isLimited) {
+	if (([tableContentView numberOfSelectedRows] == [tableContentView numberOfRows]) && !isFiltered && !isLimited && !isInterruptedLoad) {
 		
 		contextInfo = @"removeallrows";
 		
@@ -2288,7 +2307,7 @@
 	BOOL checkStatusCount = NO;
 
 	// For unfiltered and non-limited tables, use the result count - and update the status count
-	if (!isLimited && !isFiltered) {
+	if (!isLimited && !isFiltered && !isInterruptedLoad) {
 		maxNumRows = tableRowsCount;
 		maxNumRowsIsEstimate = NO;
 		[tableDataInstance setStatusValue:[NSString stringWithFormat:@"%d", maxNumRows] forKey:@"Rows"];
diff --git a/Source/TableDocument.h b/Source/TableDocument.h
index f9e96730..db4b48ae 100644
--- a/Source/TableDocument.h
+++ b/Source/TableDocument.h
@@ -138,6 +138,9 @@
 	float taskProgressValueDisplayInterval;
 	NSTimer *taskDrawTimer;
 	NSViewAnimation *taskFadeAnimator;
+	BOOL taskCanBeCancelled;
+	id taskCancellationCallbackObject;
+	SEL taskCancellationCallbackSelector;
 
 	NSToolbar *mainToolbar;
 	NSToolbarItem *chooseDatabaseToolbarItem;
@@ -186,6 +189,9 @@
 - (void) setTaskPercentage:(float)taskPercentage;
 - (void) setTaskProgressToIndeterminate;
 - (void) endTask;
+- (void) enableTaskCancellationWithTitle:(NSString *)buttonTitle callbackObject:(id)callbackObject callbackFunction:(SEL)callbackFunction;
+- (void) disableTaskCancellation;
+- (IBAction) cancelTask:(id)sender;
 - (BOOL) isWorking;
 - (void) setDatabaseListIsSelectable:(BOOL)isSelectable;
 - (void) centerTaskWindow;
diff --git a/Source/TableDocument.m b/Source/TableDocument.m
index fa07e143..78df5e16 100644
--- a/Source/TableDocument.m
+++ b/Source/TableDocument.m
@@ -103,6 +103,9 @@
 		taskProgressValueDisplayInterval = 1;
 		taskDrawTimer = nil;
 		taskFadeAnimator = nil;
+		taskCanBeCancelled = NO;
+		taskCancellationCallbackObject = nil;
+		taskCancellationCallbackSelector = NULL;
 
 		keyChainID = nil;
 	}
@@ -223,15 +226,17 @@
 		NSLog(@"Progress indicator layer could not be loaded; progress display will not function correctly.");
 	}
 
-	// Set up the progress indicator child window and later - add to main window, change indicator color and size
+	// Set up the progress indicator child window and layer - add to main window, change indicator color and size
+	[taskProgressIndicator setForeColor:[NSColor whiteColor]];
 	taskProgressWindow = [[NSWindow alloc] initWithContentRect:[taskProgressLayer bounds] styleMask:NSBorderlessWindowMask backing:NSBackingStoreBuffered defer:NO];
 	[taskProgressWindow setOpaque:NO];
-	[taskProgressWindow setIgnoresMouseEvents:YES];
 	[taskProgressWindow setBackgroundColor:[NSColor clearColor]];
 	[taskProgressWindow setAlphaValue:0.0];
-	[[taskProgressWindow contentView] addSubview:taskProgressLayer];
+	[taskProgressWindow orderWindow:NSWindowAbove relativeTo:[tableWindow windowNumber]];
 	[tableWindow addChildWindow:taskProgressWindow ordered:NSWindowAbove];
-	[taskProgressIndicator setForeColor:[NSColor whiteColor]];
+	[taskProgressWindow release];
+	[taskProgressWindow setContentView:taskProgressLayer];
+	[self centerTaskWindow];
 }
 
 /**
@@ -1258,7 +1263,7 @@
 		[historyControl setEnabled:NO];
 		databaseListIsSelectable = NO;
 		[[NSNotificationCenter defaultCenter] postNotificationName:SPDocumentTaskStartNotification object:self];
-
+		
 		// Schedule appearance of the task window in the near future
 		taskDrawTimer = [[NSTimer scheduledTimerWithTimeInterval:0.25 target:self selector:@selector(showTaskProgressWindow:) userInfo:nil repeats:NO] retain];
 	}
@@ -1333,6 +1338,9 @@
 	// Decrement the working level
 	_isWorkingLevel--;
 
+	// Ensure cancellation interface is reset
+	[self disableTaskCancellation];
+
 	// If all tasks have ended, re-enable the interface
 	if (!_isWorkingLevel) {
 
@@ -1358,6 +1366,57 @@
 }
 
 /**
+ * Allow a task to be cancelled, enabling the button with a supplied title
+ * and optionally supplying a callback object and function.
+ */
+- (void) enableTaskCancellationWithTitle:(NSString *)buttonTitle callbackObject:(id)callbackObject callbackFunction:(SEL)callbackFunction
+{
+
+	// If no task is active, return
+	if (!_isWorkingLevel) return;
+
+	if (callbackObject && callbackFunction) {
+		taskCancellationCallbackObject = callbackObject;
+		taskCancellationCallbackSelector = callbackFunction;
+	}
+	taskCanBeCancelled = YES;
+
+	[taskCancelButton setTitle:buttonTitle];
+	[taskCancelButton setEnabled:YES];
+	[taskCancelButton setHidden:NO];
+}
+
+/**
+ * Disable task cancellation.  Called automatically at the end of a task.
+ */
+- (void) disableTaskCancellation
+{
+
+	// If no task is active, return
+	if (!_isWorkingLevel) return;
+	
+	taskCanBeCancelled = NO;
+	taskCancellationCallbackObject = nil;
+	taskCancellationCallbackSelector = NULL;
+	[taskCancelButton setHidden:YES];
+}
+
+/**
+ * Action sent by the cancel button when it's active.
+ */
+- (IBAction) cancelTask:(id)sender
+{
+	if (!taskCanBeCancelled) return;
+
+	[taskCancelButton setEnabled:NO];
+	[mySQLConnection cancelCurrentQuery];
+
+	if (taskCancellationCallbackObject && taskCancellationCallbackSelector) {
+		[taskCancellationCallbackObject performSelector:taskCancellationCallbackSelector];
+	}
+}
+
+/**
  * Returns whether the document is busy performing a task - allows UI or actions
  * to be restricted as appropriate.
  */
@@ -3553,7 +3612,6 @@
 	if (spfSession) [spfSession release];
 	if (spfDocData) [spfDocData release];
 	if (keyChainID) [keyChainID release];
-	if (taskProgressWindow) [taskProgressWindow release];
 	
 	[super dealloc];
 }
diff --git a/Source/TableSource.m b/Source/TableSource.m
index 932d9cdd..6228e6f1 100644
--- a/Source/TableSource.m
+++ b/Source/TableSource.m
@@ -1033,8 +1033,20 @@ returns a dictionary containing enum/set field names as key and possible values
 }
 
 - (id)tableView:(NSTableView *)aTableView objectValueForTableColumn:(NSTableColumn *)aTableColumn row:(int)rowIndex
-{	
-	return [(aTableView == tableSourceView) ? [tableFields objectAtIndex:rowIndex] : [indexes objectAtIndex:rowIndex] objectForKey:[aTableColumn identifier]];
+{
+	NSDictionary *theRow;
+
+	if (aTableView == tableSourceView) {
+		
+		// Return a placeholder if the table is reloading
+		if (rowIndex >= [tableFields count]) return @"...";
+
+		theRow = [tableFields objectAtIndex:rowIndex];
+	} else {
+		theRow = [indexes objectAtIndex:rowIndex];
+	}
+
+	return [theRow objectForKey:[aTableColumn identifier]];
 }
 
 - (void)tableView:(NSTableView *)aTableView setObjectValue:(id)anObject forTableColumn:(NSTableColumn *)aTableColumn row:(int)rowIndex