- In the SPMySQL.framework, separate framework-triggered connections and disconnections from external actions, and use that separation to perform safer disconnects

 - When closing a database document, add a new notification, and use that to resolve retain cycles affecting connection processes
 - Improve connection controller disconnection when the document is closed, fixing crashes, by building on those two features (addresses Issue #1396)
 - Use some of the new functionality to improve SSH and MySQL connection cancellation, making both cancelable in the interface and making both respond much more quickly

15 files changed, 229 insertions, 168 deletions

diff --git a/Frameworks/SPMySQLFramework/Source/SPMySQL Private APIs.h b/Frameworks/SPMySQLFramework/Source/SPMySQL Private APIs.h
index 64fa0bc3..c3c57e3f 100644
--- a/Frameworks/SPMySQLFramework/Source/SPMySQL Private APIs.h
+++ b/Frameworks/SPMySQLFramework/Source/SPMySQL Private APIs.h
@@ -41,10 +41,12 @@
 
 @interface SPMySQLConnection (PrivateAPI)
 
+- (BOOL)_connect;
 - (MYSQL *)_makeRawMySQLConnectionWithEncoding:(NSString *)encodingName isMasterConnection:(BOOL)isMaster;
 - (BOOL)_reconnectAllowingRetries:(BOOL)canRetry;
 - (BOOL)_reconnectAfterBackgroundConnectionLoss;
 - (BOOL)_waitForNetworkConnectionWithTimeout:(double)timeoutSeconds;
+- (void)_disconnect;
 - (void)_updateConnectionVariables;
 - (void)_restoreConnectionVariables;
 - (BOOL)_checkConnectionIfNecessary;
diff --git a/Frameworks/SPMySQLFramework/Source/SPMySQLConnection Categories/Max Packet Size.m b/Frameworks/SPMySQLFramework/Source/SPMySQLConnection Categories/Max Packet Size.m
index e0bfef52..867b7fba 100644
--- a/Frameworks/SPMySQLFramework/Source/SPMySQLConnection Categories/Max Packet Size.m
+++ b/Frameworks/SPMySQLFramework/Source/SPMySQLConnection Categories/Max Packet Size.m
@@ -150,7 +150,7 @@
 		if (newSize != NSNotFound) {
 
 			// Successfully increased the global size - reconnect to use it, and return success
-			[self reconnect];
+			[self _reconnectAllowingRetries:YES];
 			return YES;
 		}
 	}
diff --git a/Frameworks/SPMySQLFramework/Source/SPMySQLConnection Categories/Querying & Preparation.m b/Frameworks/SPMySQLFramework/Source/SPMySQLConnection Categories/Querying & Preparation.m
index bafa1e6e..22e35648 100644
--- a/Frameworks/SPMySQLFramework/Source/SPMySQLConnection Categories/Querying & Preparation.m
+++ b/Frameworks/SPMySQLFramework/Source/SPMySQLConnection Categories/Querying & Preparation.m
@@ -222,6 +222,11 @@
 	lastQueryWasCancelled = NO;
 	lastQueryWasCancelledUsingReconnect = NO;
 
+	// If a disconnect was requested, cancel the action
+	if (userTriggeredDisconnect) {
+		return nil;
+	}
+
 	// Check the connection state - if no connection is available, log an
 	// error and return.
 	if (state == SPMySQLDisconnected || state == SPMySQLConnecting) {
@@ -368,6 +373,7 @@
 
 	// Unlock the connection if appropriate - if not a streaming result type.
 	if (![theResult isKindOfClass:[SPMySQLStreamingResult class]]) {
+		[self _tryLockConnection];
 		[self _unlockConnection];
 
 		// Also perform restore if appropriate
@@ -514,7 +520,6 @@
 	// so set up a new connection to run the KILL command.
 	MYSQL *killerConnection = [self _makeRawMySQLConnectionWithEncoding:@"utf8" isMasterConnection:NO];
 
-
 	// If the new connection was successfully set up, use it to run a KILL command.
 	if (killerConnection) {
 		NSStringEncoding aStringEncoding = [SPMySQLConnection stringEncodingForMySQLCharset:mysql_character_set_name(killerConnection)];
@@ -553,7 +558,7 @@
 		} else {
 			NSLog(@"SPMySQL Framework: query cancellation failed due to cancellation query error (status %d)", killQueryStatus);
 		}
-	} else {
+	} else if (!userTriggeredDisconnect) {
 		NSLog(@"SPMySQL Framework: query cancellation failed because connection failed");
 	}
 
@@ -565,13 +570,13 @@
 		return;
 	}
 
-	if (state == SPMySQLDisconnecting) return;
+	if (state == SPMySQLDisconnecting || state == SPMySQLDisconnected) return;
 
 	// Reset the connection with a reconnect.  Unlock the connection beforehand,
 	// to allow the reconnect, but lock it again afterwards to restore the expected
 	// state (query execution process should unlock as appropriate).
 	[self _unlockConnection];
-	[self reconnect];
+	[self _reconnectAllowingRetries:YES];
 	[self _lockConnection];
 
 	// Reset tracking bools to cover encompassed queries
diff --git a/Frameworks/SPMySQLFramework/Source/SPMySQLConnection.m b/Frameworks/SPMySQLFramework/Source/SPMySQLConnection.m
index 1470d1bf..e12063a2 100644
--- a/Frameworks/SPMySQLFramework/Source/SPMySQLConnection.m
+++ b/Frameworks/SPMySQLFramework/Source/SPMySQLConnection.m
@@ -250,66 +250,8 @@ const char *SPMySQLSSLPermissibleCiphers = "DHE-RSA-AES256-SHA:AES256-SHA:DHE-RS
  */
 - (BOOL)connect
 {
-
-	// If a connection is already active in some form, throw an exception
-	if (state != SPMySQLDisconnected && state != SPMySQLConnectionLostInBackground) {
-		[NSException raise:NSInternalInconsistencyException format:@"Attempted to connect a connection that is not disconnected (%d).", state];
-		return NO;
-	}
-	state = SPMySQLConnecting;
-
-	// Lock the connection for safety
-	[self _lockConnection];
-
-	// Attempt the connection
-	mySQLConnection = [self _makeRawMySQLConnectionWithEncoding:encoding isMasterConnection:YES];
-
-	// If the connection failed, reset state and return
-	if (!mySQLConnection) {
-		[self _unlockConnection];
-		state = SPMySQLDisconnected;
-		return NO;
-	}
-
-	// Successfully connected - record connected state and reset tracking variables
-	state = SPMySQLConnected;
 	userTriggeredDisconnect = NO;
-	initialConnectTime = mach_absolute_time();
-	mysqlConnectionThreadId = mySQLConnection->thread_id;
-	lastConnectionUsedTime = 0;
-
-	// Update SSL state
-	connectedWithSSL = NO;
-	if (useSSL) connectedWithSSL = (mysql_get_ssl_cipher(mySQLConnection))?YES:NO;
-	if (useSSL && !connectedWithSSL) {
-		if ([delegate respondsToSelector:@selector(connectionFellBackToNonSSL:)]) {
-			[delegate connectionFellBackToNonSSL:self];
-		}
-	}
-
-	// Reset keepalive variables
-	lastKeepAliveTime = 0;
-	keepAlivePingFailures = 0;
-
-	// Clear the connection error record
-	[self _updateLastErrorID:NSNotFound];
-	[self _updateLastErrorMessage:nil];
-
-	// Unlock the connection
-	[self _unlockConnection];
-
-	// Update connection variables to be in sync with the server state.  As this performs
-	// a query, ensure the connection is still up afterwards (!)
-	[self _updateConnectionVariables];
-	if (state != SPMySQLConnected) return NO;
-
-	// Now connection is established and verified, reset the counter
-	reconnectionRetryAttempts = 0;
-
-	// Update the maximum query size
-	[self _updateMaxQuerySize];
-
-	return YES;
+	return [self _connect];
 }
 
 /**
@@ -321,6 +263,7 @@ const char *SPMySQLSSLPermissibleCiphers = "DHE-RSA-AES256-SHA:AES256-SHA:DHE-RS
  */
 - (BOOL)reconnect
 {
+	userTriggeredDisconnect = NO;
 	return [self _reconnectAllowingRetries:YES];
 }
 
@@ -329,49 +272,8 @@ const char *SPMySQLSSLPermissibleCiphers = "DHE-RSA-AES256-SHA:AES256-SHA:DHE-RS
  */
 - (void)disconnect
 {
-
-	// If state is connection lost, set state directly to disconnected.
-	if (state == SPMySQLConnectionLostInBackground) {
-		state = SPMySQLDisconnected;
-	}
-
-	// Only continue if a connection is active
-	if (state != SPMySQLConnected && state != SPMySQLConnecting) {
-		return;
-	}
-
-	state = SPMySQLDisconnecting;
-
-	// If a query is active, cancel it
-	[self cancelCurrentQuery];
-
-	// Allow any pings or cancelled queries  to complete, inside a time limit of ten seconds
-	uint64_t disconnectStartTime_t = mach_absolute_time();
-	while (![self _tryLockConnection]) {
-		usleep(100000);
-		if (_elapsedSecondsSinceAbsoluteTime(disconnectStartTime_t) > 10) break;
-	}
-	[self _unlockConnection];
-	[self _cancelKeepAlives];
-
-	// Close the underlying MySQL connection if it still appears to be active, and not reading
-	// or writing.  While this may result in a leak of the MySQL object, it prevents crashes
-	// due to attempts to close a blocked/stuck connection.
-	if (!mySQLConnection->net.reading_or_writing && mySQLConnection->net.vio && mySQLConnection->net.buff) {
-		mysql_close(mySQLConnection);
-	}
-	mySQLConnection = NULL;
-
-	// If using a connection proxy, disconnect that too
-	if (proxy) {
-		[proxy performSelectorOnMainThread:@selector(disconnect) withObject:nil waitUntilDone:YES];
-	}
-
-	// Clear host-specific information
-	if (serverVersionString) [serverVersionString release], serverVersionString = nil;
-	if (database) [database release], database = nil;
-
-	state = SPMySQLDisconnected;
+	userTriggeredDisconnect = YES;
+	[self _disconnect];
 }
 
 #pragma mark -
@@ -433,7 +335,7 @@ const char *SPMySQLSSLPermissibleCiphers = "DHE-RSA-AES256-SHA:AES256-SHA:DHE-RS
 	// attempt to reconnect once, and if that fails will ask the user how to proceed - whether
 	// to keep reconnecting, or whether to disconnect.
 	if (!connectionVerified) {
-		connectionVerified = [self reconnect];
+		connectionVerified = [self _reconnectAllowingRetries:YES];
 	}
 
 	// Update the connection tracking use variable if the connection was confirmed,
@@ -506,6 +408,84 @@ const char *SPMySQLSSLPermissibleCiphers = "DHE-RSA-AES256-SHA:AES256-SHA:DHE-RS
 @implementation SPMySQLConnection (PrivateAPI)
 
 /**
+ * Handle a connection using previously set parameters, returning success or failure.
+ */
+- (BOOL)_connect
+{
+
+	// If a connection is already active in some form, throw an exception
+	if (state != SPMySQLDisconnected && state != SPMySQLConnectionLostInBackground) {
+		[NSException raise:NSInternalInconsistencyException format:@"Attempted to connect a connection that is not disconnected (%d).", state];
+		return NO;
+	}
+	state = SPMySQLConnecting;
+
+	if (userTriggeredDisconnect) {
+		return NO;
+	}
+
+	// Lock the connection for safety
+	[self _lockConnection];
+
+	// Attempt the connection
+	mySQLConnection = [self _makeRawMySQLConnectionWithEncoding:encoding isMasterConnection:YES];
+
+	// If the connection failed, reset state and return
+	if (!mySQLConnection) {
+		[self _unlockConnection];
+		state = SPMySQLDisconnected;
+		return NO;
+	}
+
+	// If the connection was cancelled, clean up and don't continue
+	if (userTriggeredDisconnect) {
+		mysql_close(mySQLConnection);
+		[self _unlockConnection];
+		mySQLConnection = NULL;
+		return NO;
+	}
+
+	// Successfully connected - record connected state and reset tracking variables
+	state = SPMySQLConnected;
+	initialConnectTime = mach_absolute_time();
+	mysqlConnectionThreadId = mySQLConnection->thread_id;
+	lastConnectionUsedTime = 0;
+
+	// Update SSL state
+	connectedWithSSL = NO;
+	if (useSSL) connectedWithSSL = (mysql_get_ssl_cipher(mySQLConnection))?YES:NO;
+	if (useSSL && !connectedWithSSL) {
+		if ([delegate respondsToSelector:@selector(connectionFellBackToNonSSL:)]) {
+			[delegate connectionFellBackToNonSSL:self];
+		}
+	}
+
+	// Reset keepalive variables
+	lastKeepAliveTime = 0;
+	keepAlivePingFailures = 0;
+
+	// Clear the connection error record
+	[self _updateLastErrorID:NSNotFound];
+	[self _updateLastErrorMessage:nil];
+
+	// Unlock the connection
+	[self _unlockConnection];
+
+	// Update connection variables to be in sync with the server state.  As this performs
+	// a query, ensure the connection is still up afterwards (!)
+	[self _updateConnectionVariables];
+	if (state != SPMySQLConnected) return NO;
+
+	// Now connection is established and verified, reset the counter
+	reconnectionRetryAttempts = 0;
+
+	// Update the maximum query size
+	[self _updateMaxQuerySize];
+
+	return YES;
+}
+
+/**
  * Make a connection using the class connection settings, returning a MySQL
  * connection object on success.
  */
@@ -740,7 +720,7 @@ const char *SPMySQLSSLPermissibleCiphers = "DHE-RSA-AES256-SHA:AES256-SHA:DHE-RS
 
 	// If not using a proxy, or if the proxy successfully connected, trigger a connection
 	if (!proxy || [proxy state] == SPMySQLProxyConnected) {
-		[self connect];
+		[self _connect];
 	}
 
 	// If the reconnection succeeded, restore the connection state as appropriate
@@ -856,6 +836,56 @@ const char *SPMySQLSSLPermissibleCiphers = "DHE-RSA-AES256-SHA:AES256-SHA:DHE-RS
 }
 
 /**
+ * Perform a disconnect of any active connections, cleaning up state to match.
+ */
+- (void)_disconnect
+{
+
+	// If state is connection lost, set state directly to disconnected.
+	if (state == SPMySQLConnectionLostInBackground) {
+		state = SPMySQLDisconnected;
+	}
+
+	// Only continue if a connection is active
+	if (state != SPMySQLConnected && state != SPMySQLConnecting) {
+		return;
+	}
+
+	state = SPMySQLDisconnecting;
+
+	// If a query is active, cancel it
+	[self cancelCurrentQuery];
+
+	// Allow any pings or cancelled queries  to complete, inside a time limit of ten seconds
+	uint64_t disconnectStartTime_t = mach_absolute_time();
+	while (![self _tryLockConnection]) {
+		usleep(100000);
+		if (_elapsedSecondsSinceAbsoluteTime(disconnectStartTime_t) > 10) break;
+	}
+	[self _unlockConnection];
+	[self _cancelKeepAlives];
+
+	// Close the underlying MySQL connection if it still appears to be active, and not reading
+	// or writing.  While this may result in a leak of the MySQL object, it prevents crashes
+	// due to attempts to close a blocked/stuck connection.
+	if (mySQLConnection && !mySQLConnection->net.reading_or_writing && mySQLConnection->net.vio && mySQLConnection->net.buff) {
+		mysql_close(mySQLConnection);
+	}
+	mySQLConnection = NULL;
+
+	// If using a connection proxy, disconnect that too
+	if (proxy) {
+		[proxy performSelectorOnMainThread:@selector(disconnect) withObject:nil waitUntilDone:YES];
+	}
+
+	// Clear host-specific information
+	if (serverVersionString) [serverVersionString release], serverVersionString = nil;
+	if (database) [database release], database = nil;
+
+	state = SPMySQLDisconnected;
+}
+
+/**
  * Update connection variables from the server, collecting state and ensuring
  * settings like encoding are in sync.
  */
diff --git a/Source/SPConnectionController.h b/Source/SPConnectionController.h
index 6e439953..80cfd132 100644
--- a/Source/SPConnectionController.h
+++ b/Source/SPConnectionController.h
@@ -157,7 +157,6 @@
 	BOOL isEditingItemName;
     BOOL reverseFavoritesSort;
 	BOOL initComplete;
-	BOOL mySQLConnectionCancelled;
 	BOOL favoriteNameFieldWasAutogenerated;
 	
 #ifndef SP_REFACTOR	/* ivars */
@@ -211,7 +210,7 @@
 
 // Connection processes
 - (IBAction)initiateConnection:(id)sender;
-- (IBAction)cancelMySQLConnection:(id)sender;
+- (IBAction)cancelConnection:(id)sender;
 
 #ifndef SP_REFACTOR
 // Interface interaction
diff --git a/Source/SPConnectionController.m b/Source/SPConnectionController.m
index 409b89c0..23fe0911 100644
--- a/Source/SPConnectionController.m
+++ b/Source/SPConnectionController.m
@@ -93,6 +93,8 @@ static NSString *SPExportFavoritesFilename = @"SequelProFavorites.plist";
 
 - (void)_startEditingConnection;
 
+- (void)_documentWillClose:(NSNotification *)notification;
+
 static NSComparisonResult _compareFavoritesUsingKey(id favorite1, id favorite2, void *key);
 #endif
 
@@ -291,21 +293,41 @@ static NSComparisonResult _compareFavoritesUsingKey(id favorite1, id favorite2,
 }
 
 /**
- * Cancels (or rather marks) the current connection is to be cancelled once established.
- *
- * Note, that once called this method does not mark the connection attempt to be immediately cancelled as
- * there is no reliable way to actually cancel connection attempts via the MySQL client libs. Once the
- * connection is established it will be immediately killed.
+ * Cancels the current connection - both SSH and MySQL.
  */
-- (IBAction)cancelMySQLConnection:(id)sender
+- (IBAction)cancelConnection:(id)sender
 {
 #ifndef SP_REFACTOR
 	[connectButton setEnabled:NO];
 	
 	[progressIndicatorText setStringValue:NSLocalizedString(@"Cancelling...", @"cancelling task status message")];
 	[progressIndicatorText display];
-	
-	mySQLConnectionCancelled = YES;
+
+	if (mySQLConnection) {
+		[NSThread detachNewThreadSelector:@selector(disconnect) toTarget:mySQLConnection withObject:nil];
+	}
+#endif
+
+	cancellingConnection = YES;
+
+	// Cancel the MySQL connection - handing it off to a background thread - if one is present
+	if (mySQLConnection) {
+		[mySQLConnection setDelegate:nil];
+		[NSThread detachNewThreadSelector:@selector(disconnect) toTarget:mySQLConnection withObject:nil];
+		[mySQLConnection autorelease];
+		mySQLConnection = nil;
+	}
+
+	// Cancel the SSH tunnel if present
+	if (sshTunnel) {
+		[sshTunnel disconnect];
+		[sshTunnel release];
+		sshTunnel = nil;
+	}
+
+#ifndef SP_REFACTOR
+	// Restore the connection interface
+	[self _restoreConnectionInterface];
 #endif
 }
 
@@ -1558,9 +1580,7 @@ static NSComparisonResult _compareFavoritesUsingKey(id favorite1, id favorite2,
 	// Re-enable favorites table view
 	[favoritesOutlineView setEnabled:YES];
 	[(NSView *)favoritesOutlineView display];
-	
-	mySQLConnectionCancelled = NO;
-	
+
 	// Revert the connect button back to its original selector
 	[connectButton setAction:@selector(initiateConnection:)];
 }
@@ -1767,6 +1787,18 @@ static NSComparisonResult _compareFavoritesUsingKey(id favorite1, id favorite2,
 
 #pragma mark -
 
+- (void)_documentWillClose:(NSNotification *)notification
+{
+	dbDocument = nil;
+	if (mySQLConnection) {
+		[mySQLConnection setDelegate:nil];
+		[NSThread detachNewThreadSelector:@selector(disconnect) toTarget:mySQLConnection withObject:nil];
+		[mySQLConnection autorelease];
+		mySQLConnection = nil;
+	}
+	if (sshTunnel) [sshTunnel setConnectionStateChangeSelector:nil delegate:nil], [sshTunnel disconnect], [sshTunnel release];
+}
+
 - (void)dealloc
 {
 	[[NSNotificationCenter defaultCenter] removeObserver:self];
@@ -1809,12 +1841,6 @@ static NSComparisonResult _compareFavoritesUsingKey(id favorite1, id favorite2,
 	for (id retainedObject in nibObjectsToRelease) [retainedObject release];
 	
 	[nibObjectsToRelease release];
-	
-	if (mySQLConnection) {
-		[mySQLConnection setDelegate:nil];
-		[mySQLConnection release];
-	}
-	if (sshTunnel) [sshTunnel setConnectionStateChangeSelector:nil delegate:nil], [sshTunnel disconnect], [sshTunnel release];
 
 	if (connectionKeychainID) [connectionKeychainID release];
 	if (connectionKeychainItemName) [connectionKeychainItemName release];
diff --git a/Source/SPConnectionControllerDataSource.m b/Source/SPConnectionControllerDataSource.m
index a318db85..17da21c7 100644
--- a/Source/SPConnectionControllerDataSource.m
+++ b/Source/SPConnectionControllerDataSource.m
@@ -39,6 +39,7 @@
 @interface SPConnectionController ()
 
 - (void)_reloadFavoritesViewData;
+- (void)_saveCurrentDetailsCreatingNewFavorite:(BOOL)createNewFavorite validateDetails:(BOOL)validateDetails;
 
 @end
 
diff --git a/Source/SPConnectionControllerInitializer.m b/Source/SPConnectionControllerInitializer.m
index b3f06691..2f4d725e 100644
--- a/Source/SPConnectionControllerInitializer.m
+++ b/Source/SPConnectionControllerInitializer.m
@@ -85,7 +85,6 @@ static NSString *SPConnectionViewNibName = @"ConnectionView";
 		sshTunnel = nil;
 		mySQLConnection = nil;
 		cancellingConnection = NO;
-		mySQLConnectionCancelled = NO;
 		favoriteNameFieldWasAutogenerated = NO;
 		
 		[self loadNib];
@@ -169,6 +168,11 @@ static NSString *SPConnectionViewNibName = @"ConnectionView";
  */
 - (void)registerForNotifications
 {
+	[[NSNotificationCenter defaultCenter] addObserver:self
+											 selector:@selector(_documentWillClose:) 
+												 name:SPDocumentWillCloseNotification 
+											   object:dbDocument];
+
 #ifndef SP_REFACTOR
 	[[NSNotificationCenter defaultCenter] addObserver:self 
 											 selector:@selector(scrollViewFrameChanged:) 
diff --git a/Source/SPConnectionHandler.h b/Source/SPConnectionHandler.h
index 765bae73..73e1f00b 100644
--- a/Source/SPConnectionHandler.h
+++ b/Source/SPConnectionHandler.h
@@ -48,7 +48,6 @@
 - (void)mySQLConnectionEstablished;
 - (void)sshTunnelCallback:(SPSSHTunnel *)theTunnel;
 
-- (void)cancelConnection;
 - (void)addConnectionToDocument;
 
 - (void)failConnectionWithTitle:(NSString *)theTitle errorMessage:(NSString *)theErrorMessage detail:(NSString *)errorDetail rawErrorText:(NSString *)rawErrorText;
diff --git a/Source/SPConnectionHandler.m b/Source/SPConnectionHandler.m
index 1a85bc50..3495eb95 100644
--- a/Source/SPConnectionHandler.m
+++ b/Source/SPConnectionHandler.m
@@ -69,9 +69,9 @@ static NSString *SPLocalhostAddress = @"127.0.0.1";
 		[progressIndicatorText setStringValue:NSLocalizedString(@"Connecting...", @"Generic connecting very short status message")];
 	}
 	[progressIndicatorText display];
-	
+
 	[connectButton setTitle:NSLocalizedString(@"Cancel", @"cancel button")];
-	[connectButton setAction:@selector(cancelMySQLConnection:)];
+	[connectButton setAction:@selector(cancelConnection:)];
 	[connectButton setEnabled:YES];
 	[connectButton display];
 #endif
@@ -250,6 +250,11 @@ static NSString *SPLocalhostAddress = @"127.0.0.1";
 	}
 	[progressIndicatorText display];
 	
+	[connectButton setTitle:NSLocalizedString(@"Cancel", @"cancel button")];
+	[connectButton setAction:@selector(cancelConnection:)];
+	[connectButton setEnabled:YES];
+	[connectButton display];
+
 	// Trim whitespace and newlines from the SSH host field before attempting to connect
 	[self setSshHost:[[self sshHost] stringByTrimmingCharactersInSet:[NSCharacterSet whitespaceAndNewlineCharacterSet]]];
 	
@@ -285,18 +290,10 @@ static NSString *SPLocalhostAddress = @"127.0.0.1";
 {	
 	isConnecting = NO;
 	
-	// If the user hit cancel during the connection attempt, or a test connection is
-	// occurring, kill the connection once established and reset the UI.
-	if (mySQLConnectionCancelled || isTestingConnection) {		
-		if ([mySQLConnection isConnected]) {
-			[mySQLConnection disconnect];
-			[mySQLConnection release], mySQLConnection = nil;
-		}
-		
-		// Kill the SSH connection if present
-		[self cancelConnection];
-		
-		[self _restoreConnectionInterface];
+	// If the user is only testing the connection, kill the connection
+	// once established and reset the UI.
+	if (isTestingConnection) {
+		[self cancelConnection:self];
 
 		if (isTestingConnection) {
 			[self _showConnectionTestResult:NSLocalizedString(@"Connection succeeded", @"Connection success very short status message")];
@@ -387,21 +384,6 @@ static NSString *SPLocalhostAddress = @"127.0.0.1";
 	}
 }
 
-/*
- * Cancel connection.
- */
-- (void)cancelConnection
-{
-	cancellingConnection = YES;
-
-	if (!sshTunnel) return;
-
-	[sshTunnel disconnect];
-	[sshTunnel release];
-	
-	sshTunnel = nil;
-}
-
 /**
  * Add the connection to the parent document and restore the
  * interface, allowing the application to run as normal.
@@ -463,9 +445,8 @@ static NSString *SPLocalhostAddress = @"127.0.0.1";
 		[[(NSObject *)delegate onMainThread] connectionControllerConnectAttemptFailed:self];
 	}
 	
-	// Only display the connection error message if there is a window visible and the connection attempt
-	// wasn't cancelled even though it failed.
-	if ([[dbDocument parentWindow] isVisible] && (!mySQLConnectionCancelled)) {
+	// Only display the connection error message if there is a window visible
+	if ([[dbDocument parentWindow] isVisible]) {
 		SPBeginAlertSheet(theTitle, NSLocalizedString(@"OK", @"OK button"), (errorDetail) ? NSLocalizedString(@"Show Detail", @"Show detail button") : nil, (isSSHTunnelBindError) ? NSLocalizedString(@"Use Standard Connection", @"use standard connection button") : nil, [dbDocument parentWindow], self, @selector(connectionFailureSheetDidEnd:returnCode:contextInfo:), @"connect", theErrorMessage);
 	}
 }
diff --git a/Source/SPConstants.h b/Source/SPConstants.h
index 3978db6f..03da11a6 100644
--- a/Source/SPConstants.h
+++ b/Source/SPConstants.h
@@ -394,6 +394,7 @@ extern NSString *SPExportFilenameFormat;
 extern NSString *SPContentFilters;
 extern NSString *SPDocumentTaskEndNotification;
 extern NSString *SPDocumentTaskStartNotification;
+extern NSString *SPDocumentWillCloseNotification;
 extern NSString *SPActivitiesUpdateNotification;
 extern NSString *SPFieldEditorSheetFont;
 extern NSString *SPLastSQLFileEncoding;
diff --git a/Source/SPConstants.m b/Source/SPConstants.m
index b72c0217..33f5e22a 100644
--- a/Source/SPConstants.m
+++ b/Source/SPConstants.m
@@ -199,6 +199,7 @@ NSString *SPExportFilenameFormat                 = @"SPExportFilenameFormat";
 NSString *SPContentFilters                       = @"ContentFilters";
 NSString *SPDocumentTaskEndNotification          = @"DocumentTaskEnded";
 NSString *SPDocumentTaskStartNotification        = @"DocumentTaskStarted";
+NSString *SPDocumentWillCloseNotification        = @"DocumentWillClose";
 NSString *SPActivitiesUpdateNotification         = @"ActivitiesUpdateNotification";
 NSString *SPFieldEditorSheetFont                 = @"FieldEditorSheetFont";
 NSString *SPLastSQLFileEncoding                  = @"lastSqlFileEncoding";
diff --git a/Source/SPDatabaseDocument.m b/Source/SPDatabaseDocument.m
index 0bfbb16f..fa46db2d 100644
--- a/Source/SPDatabaseDocument.m
+++ b/Source/SPDatabaseDocument.m
@@ -4065,8 +4065,11 @@ static NSString *SPRenameDatabaseAction = @"SPRenameDatabase";
 
 
 	[mySQLConnection setDelegate:nil];
-	if (_isConnected) [self closeConnection];
-	else [connectionController cancelConnection];
+	if (_isConnected) {
+		[self closeConnection];
+	} else {
+		[connectionController cancelConnection:self];
+	}
 #ifndef SP_REFACTOR
 	if ([[[SPQueryController sharedQueryController] window] isVisible]) [self toggleConsole:self];
 	[createTableSyntaxWindow orderOut:nil];
@@ -5737,7 +5740,9 @@ static NSString *SPRenameDatabaseAction = @"SPRenameDatabase";
 
 #endif
 
-	[databaseStructureRetrieval destroy];
+	// Tell listeners that this database document is being closed - fixes retain cycles and allows cleanup
+	[[NSNotificationCenter defaultCenter] postNotificationName:SPDocumentWillCloseNotification object:self];
+
 	[databaseStructureRetrieval release];
 
 	[allDatabases release];
diff --git a/Source/SPDatabaseStructure.h b/Source/SPDatabaseStructure.h
index f382944a..57249c93 100644
--- a/Source/SPDatabaseStructure.h
+++ b/Source/SPDatabaseStructure.h
@@ -53,7 +53,7 @@
 // Setup and teardown
 - (id)initWithDelegate:(SPDatabaseDocument *)theDelegate;
 - (void)setConnectionToClone:(SPMySQLConnection *)aConnection;
-- (void)destroy;
+- (void)destroy:(NSNotification *)notification;
 
 // Information
 - (SPMySQLConnection *)connection;
diff --git a/Source/SPDatabaseStructure.m b/Source/SPDatabaseStructure.m
index 16a5b205..2ceb2728 100644
--- a/Source/SPDatabaseStructure.m
+++ b/Source/SPDatabaseStructure.m
@@ -81,6 +81,11 @@
 		structure = [[NSMutableDictionary alloc] initWithCapacity:1];
 		allKeysofDbStructure = [[NSMutableArray alloc] initWithCapacity:20];
 
+		[[NSNotificationCenter defaultCenter] addObserver:self
+												 selector:@selector(destroy:) 
+													 name:SPDocumentWillCloseNotification 
+												   object:theDelegate];
+
 		// Set up the connection, thread management and data locks
 		pthread_mutex_init(&threadManagementLock, NULL);
 		pthread_mutex_init(&dataLock, NULL);
@@ -105,7 +110,7 @@
 /**
  * Ensure that processing is completed.
  */
-- (void)destroy
+- (void)destroy:(NSNotification *)notification
 {
 	delegate = nil;
 
@@ -127,7 +132,9 @@
 
 - (void)dealloc
 {
-	[self destroy];
+	[[NSNotificationCenter defaultCenter] removeObserver:self];
+
+	[self destroy:nil];
 	[structureRetrievalThreads release];
 
 	pthread_mutex_destroy(&threadManagementLock);