I would make a few suggestions. The last two, listed separately, may be of marginal value to you and do increase the complexity of the code considerably. I have included a sample implementation that incorporates these two ideas mostly for my future reference and use. Ironically, I could now use somebody to review this for me!
- I would prepend to your class's attribute names a '_' signifying that they are "private."
- If a client uses the context manager method
connection to obtain a connection, there is a default timeout argument value. However, the same is not true if the client wants to explicitly call get_connection and release_connection. But if these two methods are not meant to be called by the client, then rename these methods so that they have a leading '_'. Otherwise, you can have defaults for both styles of acquisition with that default value specified only once for consistency:
def get_connection(self, timeout=None):
if timeout is None:
timeout = 10
... # rest of code omitted
@contextmanager
def connection(self, timeout=None):
conn = self.get_connection(timeout)
... # rest of code omitted
- You might wish to implement a
close method that closes all the connections. This method should only be called when it is known that all the connections have been returned to the pool for good.
- I have updated the signature on the
__init__ method so that the client can add additional parameters the call to sqlite3.connect.
- If you timeout during acquisition, the
RuntimeError message should be changed to "Timeout: No available connection in the pool."
- A connection pool is just a special case of a more general "resource" pool where the resource in question is a reusable connection. Consider creating an abstract base class,
ResourcePool, that can be used to allocate any type of resource. This class would have an abstract method, allocate_resource, that is overridden in a subclass (e.g. your ConectionPool class) and creates an instance of a specific type of resource. Thus, ResourcePool class would have method names such as allocate_resource instead of create_connection, get_resource instead of get_connection, etc. Your ConnectionPool subclass can define, for example, a get_connection method that simply delegates to the base class get_resource method so that the client can use connection-specific method names.
If you want to be slightly more efficient at the cost of additional complexity, then:
- You might wish to lazily create connections. Store as attributes the maximum number of connections to create, the number of connections created and the number of connections currently in the pool, i.e. the
Queue instance, that are immediately available for allocation. As long as there are unused connections in the queue, you can simply get the next available one. But if the queue is empty, then if the number of connections that have been created is less than the maximum number of connections permitted, you simply create another connection, update the count of created connections and return the new connection. This logic is a critical section that needs to be serialized under control of a threading.Lock instance. Lazy allocation will save you something if the maximum number of concurrent connections in use at one time is considerably less than your pool size and if creating a connection uses a lot of resources. This is probably not your situation.
- If a connection is returned to the pool with uncommitted updates, it would be wise to perform a rollback on the connection.
- You might also wish to use a
collections.deque instance instead of a queue.Queue, since it would be a bit faster. But your code also becomes a bit more complicated.
Update
Your implementation, which used a queue.Queue instance to implement the pool (as well as my previous sample implementation) assumes that the connections would be obtained by threads running concurrently for if it were only a single thread running, what would be the point of that thread obtaining multiple connections to the same database? You would then really only need a pool size of 1. But in this case, you would still be better off just implementing a singleton connection and forget about using a pool.
So the problem now become this: A sqlite3 connection can only be used on the thread that created it (at least that is the case for my Python 3.8.5). So such a pool is really not suitable for a multithreading environment.
My conclusion is that such a pool needs to be using instead asyncio and thus your original implementation needs to be using a asyncio.Queue instance. So my new sample implementation would be:
Sample asyncio Implementation Using a deque and Lazy Connection Creation
import asyncio
from abc import ABC, abstractmethod
from collections import deque
from contextlib import asynccontextmanager
class ResourcePool(ABC):
def __init__(self, max_resources):
if max_resources < 1:
raise ValueError(f'Invalid max_resources argument: {max_resources}')
self._max_resources = max_resources
self._request_lock = asyncio.Lock()
self._resource_returned = asyncio.Condition()
self._pool = deque()
self._n_resources_created = 0
self._is_hashable_resource = self.is_hashable_resource()
self._in_use = set()
async def close(self):
"""Close all resources."""
async with self._request_lock:
while self._pool:
res = self._pool.popleft()
await self.close_resource(res)
while self._in_use:
res = self._in_use.pop()
await self.close_resource(res)
async def close_resource(self, res):
"""Subclasses should override this method if the resource does
not implement a close method."""
await res.close()
def is_hashable_resource(self):
"""Override this function and return False if the resource cannot be
added to a set. Otherwise, we can do some additional error checking and
ensure that all resources are closed when method close is called."""
return True
@abstractmethod
async def create_resource(self):
pass
async def get_resource(self, timeout=None):
if timeout is None:
timeout = 10
while True:
async with self._request_lock:
if self._pool:
res = self._pool.popleft()
if self._is_hashable_resource:
self._in_use.add(res)
return res
# Can we create another resource?
if self._n_resources_created < self._max_resources:
self._n_resources_created += 1
res = await self.create_resource()
if self._is_hashable_resource:
self._in_use.add(res)
return res
# We must wait for a resource to be returned
async def wait_for_resource():
async with self._resource_returned:
await self._resource_returned.wait()
try:
await asyncio.wait_for(wait_for_resource(), timeout=timeout)
# The pool now has at least one resource available and
# we will succeed on next iteration.
except asyncio.TimeoutError:
raise RuntimeError("Timeout: No available resource in the pool.")
async def release_resource(self, res):
if self._is_hashable_resource:
if res not in self._in_use:
raise Exception('Releasing unknown object')
# Could raise exception if two threads are releasing the same resource:
self._in_use.remove(res)
self._pool.append(res)
# If someone is waiting for a resource:
async with self._resource_returned:
self._resource_returned.notify()
@asynccontextmanager
async def resource(self, timeout=None):
res = await self.get_resource(timeout)
try:
yield res
finally:
await self.release_resource(res)
import aiosqlite3
class ConnectionPool(ResourcePool):
def __init__(self, max_connections, database):
super().__init__(max_connections)
self._database = database
async def create_resource(self):
return await aiosqlite3.connect(self._database, asyncio.get_running_loop())
def get_connection(self, timeout=None):
return self.get_resource(timeout)
async def release_resource(self, connection):
await connection.rollback() # clean up
return await super().release_resource(connection)
def release_connection(self, connection):
return self.release_resource(connection)
def connection(self, timeout=None):
return self.resource(timeout)
if __name__ == "__main__":
async def worker(pool, n):
try:
for i in range(10):
async with pool.connection() as conn:
cursor = await conn.cursor()
await cursor.execute('SELECT COUNT(*) FROM braintree_transaction')
row = await cursor.fetchone()
await cursor.close()
print(n, i, row[0])
await asyncio.sleep(.0001)
except Exception as e:
print(e)
connection_pool = ConnectionPool(2, database='fncpl.db')
tasks = [worker(connection_pool, n) for n in range(3)]
loop = asyncio.get_event_loop()
loop.run_until_complete(asyncio.gather(*tasks))
loop.run_until_complete(connection_pool.close())
loop.close()
A sample implementation that is based on threads (not suitable for a sqlite3 pool) would be:
Sample Multi-threaded Implementation
from abc import ABC, abstractmethod
from collections import deque
from threading import Lock, Condition
from contextlib import contextmanager
class ResourcePool(ABC):
def __init__(self, max_resources):
if max_resources < 1:
raise ValueError(f'Invalid max_resources argument: {max_resources}')
self._max_resources = max_resources
self._request_lock = Lock()
self._resource_returned = Condition()
self._pool = deque()
self._n_resources_created = 0
self._is_hashable_resource = self.is_hashable_resource()
self._in_use = set()
def close(self):
"""Close all resources."""
with self._request_lock:
while self._pool:
self.close_resource(self._pool.popleft())
while self._in_use:
self.close_resource(self._in_use.pop())
def close_resource(self, res):
"""Subclasses should override this method if the resource does
not implement a close method."""
res.close()
def __del__(self):
"""Close all resources when the pool is garbage
collected. Note that due to how close is implemented, it may be
called multiple times."""
self.close()
def is_hashable_resource(self):
"""Override this function and return False if the resource cannot be
added to a set. Otherwise, we can do some additional error checking and
ensure that all resources are closed when method close is called."""
return True
@abstractmethod
def create_resource(self):
pass
def get_resource(self, timeout=None):
if timeout is None:
timeout = 10
with self._request_lock:
while True:
if self._pool:
res = self._pool.popleft()
if self._is_hashable_resource:
self._in_use.add(res)
return res
# Can we create another resource?
if self._n_resources_created < self._max_resources:
self._n_resources_created += 1
res = self.create_resource()
if self._is_hashable_resource:
self._in_use.add(res)
return res
# We must wait for a resource to be returned
with self._resource_returned:
if not self._resource_returned.wait(timeout=timeout):
raise RuntimeError("Timeout: No available resource in the pool.")
# The pool now has at least one resource available and
# we will succeed on next iteration.
def release_resource(self, res):
if self._is_hashable_resource:
if res not in self._in_use:
raise Exception('Releasing unknown object')
# Could raise exception if two threads are releasing the same resource:
self._in_use.remove(res)
self._pool.append(res)
# Notify the thread waiting for a resource, if any:
with self._resource_returned:
self._resource_returned.notify()
@contextmanager
def resource(self, timeout=None):
res = self.get_resource(timeout)
try:
yield res
finally:
self.release_resource(res)
import mysql.connector
class ConnectionPool(ResourcePool):
def __init__(self, max_connections, database, user, password, time_zone='America/New_York'):
super().__init__(max_connections)
self._database = database
self._user = user
self._password = password
self._time_zone = time_zone
def create_resource(self):
return mysql.connector.connect(database=self._database,
user=self._user,
password=self._password,
charset='utf8mb4',
use_unicode=True,
init_command=f"set session time_zone='{self._time_zone}'"
)
def get_connection(self, timeout=None):
return self.get_resource(timeout)
def release_resource(self, connection):
connection.rollback() # clean up
return super().release_resource(connection)
def release_connection(self, connection):
return self.release_resource(connection)
def connection(self, timeout=None):
return self.resource(timeout)
if __name__ == '__main__':
from multiprocessing.pool import ThreadPool
def worker(pool, n):
try:
for i in range(10):
with pool.connection() as conn:
cursor = conn.cursor()
cursor.execute('SELECT COUNT(*) FROM transaction')
cnt = cursor.fetchone()[0]
cursor.close()
print(n, i, cnt)
except Exception as e:
print(e)
connection_pool = ConnectionPool(2, database='test_db', user='***', password='***')
thread_pool = ThreadPool(3)
for n in range(3):
thread_pool.apply_async(worker, args=(connection_pool, n))
thread_pool.close()
thread_pool.join()
Note that in both of these implementations if the resource being created can be hashed, then we can create a set to which we add resources that have been given out to clients. This allows us to provide an additional check to ensure that a client cannot return to the pool a resource that was never taken from the pool. It also allows the close method to close all resources including those that have not been returned back to the pool yet. There is no reason why this logic cannot be incorporated into the asyncio version.
