implement timeout and cancelable_queries

2 files changed, 49 insertions, 4 deletions

diff --git a/it_dbi.class b/it_dbi.class
index 6d798a6..86c5e6d 100644
--- a/it_dbi.class
+++ b/it_dbi.class
@@ -43,6 +43,7 @@ class it_dbi
 		'throttle_writes' => 0, # sleep for 'throttle_writes' multiplied by the execution time after every write
 		'unbuffered' => false, # use MYSQLI_USE_RESULT (WARNING this is not at all equivalent to normal it_dbi WARNING)
 		'ignored_warnings' => "", # regex of additional mysql warnings numbers to ignore
+		'cancelable_queries' => false,
 	);
 
 	var $_key;         # Key of currently loaded record or null (public readonly)
@@ -1126,9 +1127,28 @@ function _tables($p) {
 	return (array)$result;
 }
 
+function __query($query, $p)
+{
+	if ($p['timeout'] || $this->_p['cancelable_queries'])
+	{
+		$starttime = microtime(true);
+		mysqli_query($this->_link, $query, ($p['unbuffered'] ? MYSQLI_USE_RESULT : MYSQLI_STORE_RESULT) | MYSQLI_ASYNC);
+		do {
+			$read = $error = $reject = [$this->_link];
+			if ($p['timeout'] && (microtime(true) - $starttime) > $p['timeout'])
+				return false;
+		} while (!mysqli_poll($read, $error, $reject, 1, 0));
+		return mysqli_reap_async_query($this->_link);
+	}
+	else
+	{
+		return mysqli_query($this->_link, $query, $p['unbuffered'] ? MYSQLI_USE_RESULT : MYSQLI_STORE_RESULT);
+	}
+}
+
 function _query($query, $p)
 {
-	if (!($result = mysqli_query($this->_link, $query, $p['unbuffered'] ? MYSQLI_USE_RESULT : MYSQLI_STORE_RESULT)) && $p['safety'])
+	if (!($result = $this->__query($query, $p)) && $p['safety'])
 	{
 		$errno = mysqli_errno($this->_link);
 		if (($p['safety'] < 2) && ($errno == 1062)) # Duplicate entry
@@ -1138,7 +1158,7 @@ function _query($query, $p)
 		{
 			it::log('sqllog', "it_dbi(): reconnecting mysqli_connect {$p['server']}, {$p['db']}");
 			$this->_connect(array('reconnect' => true));
-			$result = mysqli_query($this->_link, $query, $p['unbuffered'] ? MYSQLI_USE_RESULT : MYSQLI_STORE_RESULT);
+			$result = $this->__query($query, $p);
 		}
 	}
 
diff --git a/it_dbi_postgres.class b/it_dbi_postgres.class
index 73c76f2..ac9c4ae 100644
--- a/it_dbi_postgres.class
+++ b/it_dbi_postgres.class
@@ -109,6 +109,31 @@ function _connect_db($p) {
 	return [$result, $result ? '' : 'Could not connect'];
 }
 
+function __query($query, $p)
+{
+	if ($p['timeout'] || $this->_p['cancelable_queries'])
+	{
+		$starttime = microtime(true);
+		pg_send_query($this->_link, $query);
+		while (pg_connection_busy($this->_link)) {
+			if ($p['timeout'] && (microtime(true) - $starttime) > $p['timeout'])
+				return false;
+			$read = $error = [pg_socket($this->_link)];
+			socket_select($read, [], $error, 1, 0);
+		}
+
+		while ($newresult = @pg_get_result($this->_link))
+			$result = $newresult;
+		if ($result && pg_result_error($result))
+			$result = false;
+		return $result;
+	}
+	else
+	{
+		return @pg_query($this->_link, $query);
+	}
+}
+
 function _query($query, $p)
 {
 	if ($this->_p['keyfield'] && it::match('^INSERT ', $query))
@@ -117,7 +142,7 @@ function _query($query, $p)
 		$query .= ' RETURNING ' . $this->_escape_name($this->_p['keyfield']);
 	}
 
-	if (!($result = @pg_query($this->_link, $query)) && $p['safety'])
+	if (!($result = $this->__query($query, $p)) && $p['safety'])
 	{
 		if (($p['safety'] < 2) && it::match('duplicate key value', pg_last_error($this->_link))) # Duplicate entry
 			return null;
@@ -127,7 +152,7 @@ function _query($query, $p)
 		{
 			it::log('sqllog', "it_dbi(): reconnecting mysqli_connect {$p['server']}, {$p['db']}");
 			$this->_connect(array('reconnect' => true));
-			$result = mysqli_query($this->_link, $query, $p['unbuffered'] ? MYSQLI_USE_RESULT : MYSQLI_STORE_RESULT);
+			$result = $this->__query($query, $p);
 		}
 		*/
 	}