Add pools module

Author: methane

MANIFEST.in

@@ -1,3 +1,3 @@
 include README.rst LICENSE CHANGELOG
 include runtests.py tox.ini
-include example.py
+include example.py example_pool

README.rst

@@ -12,9 +12,17 @@ This package contains a fork of PyMySQL supporting Tornado.
 Example
 -------
 
+Simple
+~~~~~~~
+
 .. include:: example.py
     :code: python
 
+pool
+~~~~~
+
+.. include:: example_pool.py
+    :code: python
 
 Requirements
 -------------

example_pool.py

@@ -0,0 +1,34 @@
+#!/usr/bin/env python
+from __future__ import print_function
+import random
+
+from tornado import ioloop, gen
+from tornado_mysql import pools
+
+
+pools.DEBUG = True
+
+
+POOL = pools.Pool(
+    dict(host='127.0.0.1', port=3306, user='root', passwd='', db='mysql'),
+    max_idle_connections=1,
+    max_recycle_sec=3)
+
+
+@gen.coroutine
+def worker(n):
+    for i in range(10):
+        t = 1
+        print(n, "sleeping", t, "seconds")
+        cur = yield POOL.execute("SELECT SLEEP(%s)", (t,))
+        print(n, cur.fetchall())
+
+
+@gen.coroutine
+def main():
+    workers = [worker(i) for i in range(10)]
+    yield workers
+
+
+ioloop.IOLoop.current().run_sync(main)
+print(POOL._opened_conns)

tornado_mysql/connections.py

@@ -738,6 +738,8 @@ def connect(self):
             yield self._get_server_information()
             yield self._request_authentication()
 
+            self.connected_time = self.io_loop.time()
+
             if self.sql_mode is not None:
                 yield self.query("SET sql_mode=%s" % (self.sql_mode,))
 

tornado_mysql/cursors.py

@@ -89,7 +89,7 @@ def nextset(self):
             raise gen.Return()
         if not current_result.has_next:
             raise gen.Return()
-        conn.next_result()
+        yield conn.next_result()
         self._do_get_result()
         raise gen.Return(True)
 

tornado_mysql/pools.py

@@ -0,0 +1,158 @@
+"""Connection pool"""
+from __future__ import absolute_import, division, print_function
+
+from collections import deque
+import sys
+import warnings
+
+from tornado.ioloop import IOLoop
+from tornado.gen import coroutine, Return
+from tornado.concurrent import Future
+from tornado_mysql import connect
+
+
+DEBUG = False
+
+
+def _debug(*msg):
+    if DEBUG:
+        print(*msg)
+
+
+class Pool(object):
+    """Connection pool like Golang's database/sql.DB.
+
+    This connection pool is based on autocommit mode.
+    You can execute query without knowing connection.
+
+    When transaction is necessary, you can checkout transaction object.
+    """
+
+    def __init__(self,
+                 connect_kwargs,
+                 max_idle_connections=1,
+                 max_recycle_sec=3600,
+                 io_loop=None,
+                 ):
+        """
+        :param dict connect_kwargs: kwargs for tornado_mysql.connect()
+        :param int max_idle_connections: Max number of keeping connections.
+        :param int max_recycle_sec: How long connections are recycled.
+        """
+        connect_kwargs['autocommit'] = True
+        self.io_loop = io_loop or IOLoop.current()
+        self.connect_kwargs = connect_kwargs
+        self.max_idle_connections = max_idle_connections
+        self.max_recycle_sec = max_recycle_sec
+
+        self._opened_conns = 0
+        self._free_conn = deque()
+
+    def _get_conn(self):
+        now = self.io_loop.time()
+        while self._free_conn:
+            conn = self._free_conn.popleft()
+            if now - conn.connected_time > self.max_recycle_sec:
+                self._close_async(conn)
+                continue
+            _debug("Reusing connection from pool (opened=%d)" % (self._opened_conns,))
+            fut = Future()
+            fut.set_result(conn)
+            return fut
+
+        self._opened_conns += 1
+        _debug("Creating new connection (opened=%d)" % (self._opened_conns,))
+        return connect(**self.connect_kwargs)
+
+    def _put_conn(self, conn):
+        if (len(self._free_conn) < self.max_idle_connections and
+                self.io_loop.time() - conn.connected_time < self.max_recycle_sec):
+            self._free_conn.append(conn)
+        else:
+            self._close_async(conn)
+
+    def _close_async(self, conn):
+        self.io_loop.add_future(conn.close_async(), callback=lambda f: None)
+        self._opened_conns -= 1
+
+    def _close_conn(self, conn):
+        conn.close()
+        self._opened_conns -= 1
+
+    @coroutine
+    def execute(self, query, params):
+        """Execute query in pool.
+
+        Returns future yielding closed cursor.
+        You can get rows, lastrowid, etc from the cursor.
+
+        :return: Future of cursor
+        :rtype: Future
+        """
+        conn = yield self._get_conn()
+        try:
+            cur = conn.cursor()
+            yield cur.execute(query, params)
+            yield cur.close()
+            self._put_conn(conn)
+        except:
+            self._opened_conns -= 1
+            conn.close()
+            raise
+        raise Return(cur)
+
+    @coroutine
+    def begin(self):
+        """Start transaction
+
+        Wait to get connection and returns `Transaction` object.
+
+        :return: Future[Transaction]
+        :rtype: Future
+        """
+        conn = yield self._get_conn()
+        trx = Transaction(self, conn)
+        raise gen.Return(trx)
+
+
+class Transaction(object):
+    """Represents transaction in pool"""
+    def __init__(self, pool, conn):
+        self._pool = pool
+        self._conn = conn
+
+    def _ensure_conn(self):
+        if self._conn is None:
+            raise Exception("Transaction is closed already")
+
+    def _close(self):
+        self._pool._put_conn(self._conn)
+        self._pool = self._conn = None
+
+    @coroutine
+    def execute(self, query, args):
+        """
+        :return: Future[Cursor]
+        :rtype: Future
+        """
+        self._ensure_conn()
+        cur = self._conn.cursor()
+        yield cur.execute(query, args)
+        raise Return(cur)
+
+    @coroutine
+    def commit(self):
+        self._ensure_conn()
+        yield self._conn.commit()
+        self._close()
+
+    @coroutine
+    def rollback(self):
+        self._ensure_conn()
+        yield self._conn.rollback()
+        self._close()
+
+    def __del__(self):
+        if self._pool is not None:
+            warnings.warn("Transaction has not committed or rollbacked.")
+            self._pool._close_conn(self._conn)