- Added SetTimeout (settimeout from python) to Connection object whic…

…h allows the default connect/read/write timeout to be set in whole seconds (as integer)

Connection.cpp

@@ -77,6 +77,7 @@ Connection::Connection (AMConnectionCAPI *_capi)
 {
 	PRINTMARK();
 
+	m_timeout = -1;
 	m_state = NONE;
 	m_sockfd = -1;
 	m_sockInst = NULL;
@@ -349,7 +350,7 @@ bool Connection::recvPacket()
 
 		time_t tsStart = time (0);
 
-		if (!m_capi.wouldBlock(m_sockInst, m_sockfd, AMC_READ, 10))
+		if (!m_capi.wouldBlock(m_sockInst, m_sockfd, AMC_READ, m_timeout == -1 ? 10 : m_timeout))
 		{
 			return false;
 		}
@@ -378,7 +379,7 @@ bool Connection::sendPacket()
 			break;
 		}
 
-		if (!m_capi.wouldBlock(m_sockInst, m_sockfd, AMC_WRITE, 10))
+		if (!m_capi.wouldBlock(m_sockInst, m_sockfd, AMC_WRITE, m_timeout == -1 ? 10 : m_timeout))
 		{
 			setError("Socket write timed out", 0);
 			return false;
@@ -447,6 +448,15 @@ bool Connection::connect(const char *_host, int _port, const char *_username, co
 		return false;
 	}
 
+	if (m_timeout != -1)
+	{
+		if (!setTimeout (m_timeout))
+		{
+			setError("setTimeout API failed", 0);
+			return false;
+		}
+	}
+
 	PRINTMARK();
 	m_sockfd = m_capi.getSocketFD(m_sockInst);
 
@@ -736,3 +746,18 @@ int Connection::getTxBufferSize()
 {
 	return (int) m_writer.getSize();
 }
+
+bool Connection::setTimeout(int timeout)
+{
+	m_timeout = timeout;
+
+	if (m_sockInst)
+	{
+		if (!m_capi.setTimeout(m_sockInst, timeout))
+		{
+			return false;
+		}
+	}
+
+	return true;
+}

Connection.h

@@ -105,6 +105,7 @@ class Connection
 
 	std::string m_errorMessage;
 	int m_errno;
+	int m_timeout;
 
 	AMConnectionCAPI m_capi;
 
@@ -123,6 +124,7 @@ class Connection
 	int getTxBufferSize();
 	bool isConnected(void);
 	bool close(void);
+	bool setTimeout(int timeout);
 
 protected:
 	void changeState(State _newState, const char *message);

amysql.c

@@ -121,6 +121,7 @@ void API_closeSocket(void *sock);
 void API_deleteSocket(void *sock);
 int API_wouldBlock(void *sock, int fd, int ops, int timeout);
 int API_connectSocket(void *sock, const char *host, int port);
+int API_setTimeout(void *sock, int timeoutSec);
 
 void *API_createResult(int columns)
 {
@@ -651,6 +652,7 @@ AMConnectionCAPI capi = {
 	API_closeSocket,
 	API_wouldBlock,
 	API_connectSocket,
+	API_setTimeout,
 	API_createResult,
 	API_resultSetField,
 	API_resultRowBegin,
@@ -673,6 +675,24 @@ int Connection_init(Connection *self, PyObject *arg)
 	return 0;
 }
 
+PyObject *Connection_setTimeout(Connection *self, PyObject *args)
+{
+	int timeout;
+
+	if (!PyArg_ParseTuple (args, "i", &timeout))
+	{
+		return NULL;
+	}
+
+	if (!AMConnection_SetTimeout(self->conn, timeout))
+	{
+		return NULL;
+	}
+
+	Py_RETURN_NONE;
+}	
+
+
 PyObject *Connection_isConnected(Connection *self, PyObject *args)
 {
 	if (AMConnection_IsConnected(self->conn))
@@ -1128,6 +1148,7 @@ static PyMethodDef Connection_methods[] = {
 	{"query", (PyCFunction)				Connection_query,				METH_VARARGS, "Performs a query. Arguments: query, arguments to escape"},
 	{"close", (PyCFunction)	Connection_close,	METH_NOARGS, "Closes connection"},
 	{"is_connected", (PyCFunction) Connection_isConnected, METH_NOARGS, "Check connection status"},
+	{"settimeout", (PyCFunction) Connection_setTimeout, METH_VARARGS, "Sets connection timeout in seconds"},
 	{NULL}
 };
 static PyMemberDef Connection_members[] = {

amysql.h

@@ -85,6 +85,7 @@ typedef struct __AMConnectionCAPI
 	void (*closeSocket)(void *instance);
 	int (*wouldBlock)(void *instance, int fd, int ops, int timeout);
 	int (*connectSocket)(void *sock, const char *host, int port);
+	int (*setTimeout)(void *sock, int timeout);
 
 	void *(*createResult)(int columns);
 	void (*resultSetField)(void *result, int ifield, AMTypeInfo *ti, void *name, size_t cbName);
@@ -93,6 +94,8 @@ typedef struct __AMConnectionCAPI
 	void (*resultRowEnd)(void *result);
 	void (*destroyResult)(void *result);
 	void *(*resultOK)(UINT64 affected, UINT64 insertId, int serverStatus, const char *message, size_t len);
+
+
 } AMConnectionCAPI;
 
 
@@ -115,5 +118,6 @@ int AMConnection_GetTxBufferSize (AMConnection conn);
 int AMConnection_GetRxBufferSize (AMConnection conn);
 int AMConnection_IsConnected (AMConnection conn);
 int AMConnection_Close (AMConnection conn);
+int AMConnection_SetTimeout(AMConnection conn, int timeout);
 
 #endif

capi.cpp

@@ -84,6 +84,11 @@ EXPORT_ATTR int AMConnection_Connect (AMConnection conn, const char *_host, int
 	return ((Connection *)conn)->connect(_host, _port, _username, _password, _database, _autoCommit, (MYSQL_CHARSETS) _charset) ? 1 : 0;
 }
 
+EXPORT_ATTR int AMConnection_SetTimeout(AMConnection conn, int timeout)
+{
+	return ((Connection *)conn)->setTimeout(timeout) ? 1 : 0;
+}
+
 EXPORT_ATTR int AMConnection_GetLastError (AMConnection conn, const char **_ppOutMessage, int *_outErrno)
 {
 	return ((Connection *)conn)->getLastError(_ppOutMessage, _outErrno) ? 1 : 0;
@@ -108,3 +113,4 @@ EXPORT_ATTR int AMConnection_Close (AMConnection conn)
 {
 	return ((Connection *)conn)->close() ? 1 : 0;
 }
+

io_cpython.c

@@ -142,6 +142,34 @@ void *API_createSocket(int family, int type, int proto)
 	return sockobj;
 }
 
+int API_setTimeout(void *sock, int timeoutSec)
+{
+	PyObject *intobj;
+	PyObject *retobj;
+	PyObject *methodObj;
+
+	PRINTMARK();
+	intobj = PyFloat_FromDouble( (double) timeoutSec);
+
+	methodObj = PyString_FromString("settimeout");
+	PRINTMARK();
+	retobj = PyObject_CallMethodObjArgs ((PyObject *) sock, methodObj, intobj, NULL);
+	Py_DECREF(intobj);
+	Py_DECREF(methodObj);
+	PRINTMARK();
+
+	if (retobj == NULL)
+	{
+		PyErr_Clear();
+		return 0;
+	}
+
+	Py_DECREF(retobj);
+	return 1;
+
+}   
+
+
 int API_getSocketFD(void *sock)
 {
 	int ret;

io_gevent.c

@@ -60,6 +60,7 @@ SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 #include <amysql.h>
 #include <Python.h>
 #include <stdio.h>
+#include <time.h>
 
 #ifndef TRUE
 #define TRUE 1
@@ -141,6 +142,33 @@ void *API_createSocket(int family, int type, int proto)
 	return sockobj;
 }
 
+int API_setTimeout(void *sock, int timeoutSec)
+{
+	PyObject *intobj;
+	PyObject *retobj;
+	PyObject *methodObj;
+
+	PRINTMARK();
+	intobj = PyFloat_FromDouble( (double) timeoutSec);
+
+	methodObj = PyString_FromString("settimeout");
+	PRINTMARK();
+	retobj = PyObject_CallMethodObjArgs ((PyObject *) sock, methodObj, intobj, NULL);
+	Py_DECREF(intobj);
+	Py_DECREF(methodObj);
+	PRINTMARK();
+
+	if (retobj == NULL)
+	{
+		PyErr_Clear();
+		return 0;
+	}
+
+	Py_DECREF(retobj);
+	return 1;
+
+}   
+
 int API_getSocketFD(void *sock)
 {
 	int ret;
@@ -203,8 +231,9 @@ int API_connectSocket(void *sock, const char *host, int port)
 	PyTuple_SET_ITEM(addrTuple, 1, PyInt_FromLong(port));
 
 	connectStr = PyString_FromString("connect");
+
 	res = PyObject_CallMethodObjArgs( (PyObject *) sock, connectStr, addrTuple, NULL);
-
+    
 	Py_DECREF(connectStr);
 	Py_DECREF(addrTuple);
 

tests.py

@@ -107,16 +107,26 @@ def testDoubleConnect(self):
         cnn = amysql.Connection()
         cnn.connect(DB_HOST, DB_PORT, DB_USER, DB_PASSWD, DB_DB)
 
-    def testConnectTwice(self):
+    def testConnectFails(self):
         cnn = amysql.Connection()
-        cnn.connect (DB_HOST, 3306, DB_USER, DB_PASSWD, DB_DB)
+
         try:
-            cnn.connect (DB_HOST, 3306, DB_USER, DB_PASSWD, DB_DB)
+            cnn.connect (DB_HOST, 31337, DB_USER, DB_PASSWD, DB_DB)
             assert False, "Expected exception"
         except(RuntimeError):
             pass
         pass
 
+    def testConnectDNSFails(self):
+        cnn = amysql.Connection()
+
+        try:
+            cnn.connect ("thisplaceisnowere", 31337, DB_USER, DB_PASSWD, DB_DB)
+            assert False, "Expected exception"
+        except(RuntimeError):
+            pass
+        pass
+
     def testParallelQuery(self):
 
         def query(s):
@@ -132,28 +142,38 @@ def query(s):
         gevent.joinall([ch1, ch2, ch3])
 
         end = time.time()
-        self.assertAlmostEqual(3.0, end - start, places = 1)
-
-    def testConnectFails(self):
-        cnn = amysql.Connection()
+        self.assertAlmostEqual(3.0, end - start, places = 2)
 
+    def testConnectTimeout(self):
+        cnn = amysql.Connection()
+        cnn.settimeout(1)
+
+        start = time.clock()
         try:
-            cnn.connect (DB_HOST, 31337, DB_USER, DB_PASSWD, DB_DB)
-            assert False, "Expected exception"
+            cnn.connect (DB_HOST, 31481, DB_USER, DB_PASSWD, DB_DB)
+
         except(RuntimeError):
-            pass
-        pass
+            elapsed = time.clock() - start
 
-    def testConnectDNSFails(self):
+            if (elapsed > 2):
+                assert False, "Timeout isn't working"
+            return
+
+        assert False, "Expected expection"
+
+
+
+    def testConnectTwice(self):
         cnn = amysql.Connection()
-
+        cnn.connect (DB_HOST, 3306, DB_USER, DB_PASSWD, DB_DB)
         try:
-            cnn.connect ("thisplaceisnowere", 31337, DB_USER, DB_PASSWD, DB_DB)
+            cnn.connect (DB_HOST, 3306, DB_USER, DB_PASSWD, DB_DB)
             assert False, "Expected exception"
         except(RuntimeError):
             pass
         pass
-
+
+
 
     def testConnectClosed(self):
         cnn = amysql.Connection()
@@ -562,7 +582,8 @@ def testCharsets(self):
             self.assertEquals(result, expected)
 if __name__ == '__main__':
     unittest.main()
-
+
+"""            
 if __name__ == '__main__':
     from guppy import hpy
     hp = hpy()
@@ -571,5 +592,4 @@ def testCharsets(self):
         unittest.main()
         heap = hp.heapu()
         print heap
-
-
+"""