3 files changed, 50 insertions, 5 deletions
diff --git a/Frameworks/MCPKit/MCPFoundationKit/MCPConnection.m b/Frameworks/MCPKit/MCPFoundationKit/MCPConnection.m
index 4413e1c0..2463041a 100644
--- a/Frameworks/MCPKit/MCPFoundationKit/MCPConnection.m
+++ b/Frameworks/MCPKit/MCPFoundationKit/MCPConnection.m
@@ -1365,6 +1365,8 @@ static void forcePingTimeout(int signalNumber)
 		// On success, capture the results
 		if (0 == queryResultCode) {
 			
+			queryAffectedRows = mysql_affected_rows(mConnection);
+
 			if (mysql_field_count(mConnection) != 0) {
 				
 				// For normal result sets, fetch the results and unlock the connection
@@ -1391,10 +1393,8 @@ static void forcePingTimeout(int signalNumber)
 			
 			queryErrorMessage = [[NSString alloc] initWithString:@""];
 			queryErrorId = 0;
-			if (streamResultType == MCP_NO_STREAMING) {
+			if (streamResultType == MCP_NO_STREAMING && queryAffectedRows == -1) {
 				queryAffectedRows = mysql_affected_rows(mConnection);
-			} else {
-				queryAffectedRows = 0;
 			}
 			
 		// On failure, set the error messages and IDs
diff --git a/Frameworks/MCPKit/MCPFoundationKit/MCPStreamingResult.h b/Frameworks/MCPKit/MCPFoundationKit/MCPStreamingResult.h
index 146f81e7..38e59d1d 100644
--- a/Frameworks/MCPKit/MCPFoundationKit/MCPStreamingResult.h
+++ b/Frameworks/MCPKit/MCPFoundationKit/MCPStreamingResult.h
@@ -58,5 +58,6 @@ typedef struct SP_MYSQL_ROWS {
 
 // Results fetching
 - (NSArray *)fetchNextRowAsArray;
+- (void) cancelResultLoad;
 
 @end
 \ No newline at end of file
diff --git a/Frameworks/MCPKit/MCPFoundationKit/MCPStreamingResult.m b/Frameworks/MCPKit/MCPFoundationKit/MCPStreamingResult.m
index 3f3e2073..c63a9f0b 100644
--- a/Frameworks/MCPKit/MCPFoundationKit/MCPStreamingResult.m
+++ b/Frameworks/MCPKit/MCPFoundationKit/MCPStreamingResult.m
@@ -285,13 +285,57 @@
 	return returnArray;
 }
 
+/*
+ * Ensure the result set is fully processed and freed without any processing
+ */
+- (void) cancelResultLoad
+{
+	MYSQL_ROW theRow;
+
+	// Loop through all the rows and ensure the rows are fetched.
+	// If fully streaming, loop through the rows directly
+	if (fullyStreaming) {
+		while (1) {
+			theRow = mysql_fetch_row(mResult);
+
+			// If no data was returned, we're at the end of the result set - return.
+			if (theRow == NULL) return;
+		}
+
+	// If in cached-streaming/fast download mode, loop until all data is fetched and freed
+	} else {
+
+		while (1) {
+		
+			// Check to see whether we need to wait for the data to be availabe
+			// - if so, wait 1ms before checking again
+			while (!dataDownloaded && processedRowCount == downloadedRowCount) usleep(1000);
+
+			// If all rows have been processed, we're at the end of the result set - return
+			// once all memory has been freed
+			if (processedRowCount == downloadedRowCount) {
+				while (!dataFreed) usleep(1000);
+				[parentConnection unlockConnection];
+				connectionUnlocked = YES;
+				return;
+			}
+			processedRowCount++;
+		}
+	}
+}
+
 #pragma mark -
 #pragma mark Overrides for safety
 
+/**
+ * If numOfRows is used before the data is fully downloaded, -1 will be returned;
+ * otherwise the number of rows is returned.
+ */
 - (my_ulonglong)numOfRows
 {
-	NSLog(@"numOfRows cannot be used with streaming results");
-	return 0;
+	if (!dataDownloaded) return -1;
+
+	return downloadedRowCount;
 }
 
 - (void)dataSeek:(my_ulonglong) row