Hugh's Blog

Python 连接池简单实现

使用队列实现一个简单的连接池。

from queue import Queue

class Client(object):
    def query(self, sql:str):
        return ['ok']

    def disconnect(self):
        pass

class MyClient(object):
    def __init__(self):
        self.pool = None
        self._client = Client()

    def query(self, sql:str):
        try:
            data = self._client.query(sql)
        except:
            data = []

        if self.pool is not None:
            self.pool.free(self)

        return data

    def disconnect(self):
        self._client.disconnect()

class MyClientPool(object):
    def __init__(self, max_conns=5, max_timeout=5):
        self.max_conns = max_conns
        self.max_timeout = max_timeout

        self.free_conns = Queue(self.max_conns)
        for i in range(self.max_conns):
            self.free(self._create_conn())

    def _create_conn(self) -> MyClient:
        return MyClient()

    def free(self, conn:MyClient):
        if self.free_conns.full():
            conn.disconnect()
            return
        self.free_conns.put_nowait(conn)

    def get(self, timeout=None) -> MyClient:
        if timeout is None:
            timeout = self.max_timeout

        if self.free_conns.empty():
            conn = self._create_conn()
        else:
            conn = self.free_conns.get(timeout=timeout)

        conn.pool = self
        return conn

    def current_free(self):
        return self.free_conns.qsize()

pool = MyClientPool()

if __name__ == '__main__':
    print(pool.get().query('select 1'))
    print('free connections: ', pool.current_free())