fix get connection cursor

Author: snower

tests/test_with_with.py

@@ -14,7 +14,7 @@ def test1(self):
             with connection.cursor() as cursor:
                 yield cursor.execute(sql)
                 datas = cursor.fetchall()
-                print (datas)
+                self.assertTrue(bool(datas))
 
 
 class TestAsyncCursor(BaseTestCase):
@@ -34,8 +34,9 @@ class TestAsyncCursor(BaseTestCase):
     def test1(self):
         sql = "select 1 as test"
         with (yield self.pool.Connection()) as connection:
-            with connection.cursor() as cursor:
-                yield cursor.execute(sql)
-                result = yield cursor.fetchone()
-                self.assertTrue('test' in result)
-                self.assertEqual(result['test'], 1)
+            cursor = connection.cursor()
+            yield cursor.execute(sql)
+            result = yield cursor.fetchone()
+            yield cursor.close()
+            self.assertTrue('test' in result)
+            self.assertEqual(result['test'], 1)

tormysql/client.py

@@ -70,6 +70,9 @@ def select_db(self, db):
         return async_call_method(self._connection.select_db, db)
 
     def cursor(self, cursor_cls=None):
+        if cursor_cls is None:
+            cursor_cls = self._connection.cursorclass
+            
         cursor = self._connection.cursor(
             cursor_cls.__delegate_class__ if cursor_cls and issubclass(cursor_cls, Cursor) else cursor_cls
         )

tormysql/cursor.py

@@ -34,6 +34,12 @@ def close(self):
         self._cursor = None
         return future
 
+    def nextset(self):
+        return async_call_method(self._cursor.nextset)
+
+    def mogrify(self, query, args=None):
+        return self._cursor.mogrify(query, args)
+
     def execute(self, query, args=None):
         return async_call_method(self._cursor.execute, query, args)
 
@@ -83,6 +89,19 @@ class DictCursor(Cursor):
 class SSCursor(Cursor):
     __delegate_class__ = OriginSSCursor
 
+    def close(self):
+        if self._cursor is None:
+            self._cursor.close()
+            future = Future()
+            future.set_result(None)
+        else:
+            future = async_call_method(self._cursor.close)
+        self._cursor = None
+        return future
+
+    def read_next(self):
+        return async_call_method(self._cursor.read_next)
+
     def fetchone(self):
         return async_call_method(self._cursor.fetchone)
 

tormysql/util.py

@@ -8,18 +8,17 @@
 
 def async_call_method(fun, *args, **kwargs):
     future = Future()
-    io_loop = IOLoop.current()
 
     def finish():
         try:
             result = fun(*args, **kwargs)
             if future._callbacks:
-                io_loop.add_callback(future.set_result, result)
+                IOLoop.current().add_callback(future.set_result, result)
             else:
                 future.set_result(result)
         except:
             if future._callbacks:
-                io_loop.add_callback(future.set_exc_info, sys.exc_info())
+                IOLoop.current().add_callback(future.set_exc_info, sys.exc_info())
             else:
                 future.set_exc_info(sys.exc_info())