iTranslated by AI
Managing Connections with contextmanager (mysql-connector-python)
Introduction
This is a cross-post from Qiita.
When accessing a database, it is necessary to call the close process for both the cursor and the database connection at the point when you finish using them. (If you don't, there is a possibility that connections will remain open, which can cause bugs such as the system failing to operate correctly once the maximum number of simultaneous connections is reached.)
In Python, using the with statement allows the close process to be called automatically. In other words, the language provides features for so-called resource management as a standard feature. This serves the same role as the try-with-resources statement in Java.
However, depending on the library, objects such as cursors and connections may not support the with statement.
In such cases, it can be a bit painful to implement the close process manually every time to ensure it is called at the appropriate timing.
To resolve such issues, a decorator is provided as a standard feature to wrap function return values into something that can be used with the with statement. That is contextmanager.
Practical Example with mysql-connector-python
When I checked while writing this article, I found that in the latest version, both cursor and connection have __enter__ and __exit__ implemented, making them compatible with the with statement.
It wasn't available in version 2.2, so it seems support was added at some point.
connection and cursor
First, let's create a class to manage connections and cursors.
import mysql.connector
class MySQLClient:
def __init__(self, host, port, user, password, database):
self.host = host
self.port = port
self.user = user
self.password = password
self.database = database
def get_connection(self):
return mysql.connector.connect(
host=self.host,
port=self.port,
user=self.user,
password=self.password,
database=self.database
)
@contextlib.contextmanager
def get_cursor(self):
with self.get_connection() as connection, connection.cursor(dictionary=True) as cursor:
yield cursor
Normally, with contextmanager, you would use try-finally and call the close process in the finally block. However, if what you are using already supports the with statement, you can write it as shown above.
(Essentially, this usage visually consolidates the management of multiple resources into one.)
The implementation above creates a connection pool by default.
With this approach, "once the cursor is finished, the connection is also closed (returned to the connection pool)." It is used like this:
client = MySQLClient(...)
with client.get_cursor() as cursor:
query = 'select * from user'
cursor.execute(query)
for row in cursor:
print(row)
Separation of SELECT statement results and processing
While connection management has been delegated to the with statement, there is still a slight problem in practice. In this state, "fetching the result of the SELECT statement" and "processing based on that result (business logic)" cannot be separated.
Usually, a web application displays the results of a SELECT statement on a screen, or a batch process creates some kind of file based on them. Since this part is a form of business logic, we want to be able to separate these classes from the infrastructure part, which is database access.
To do this, we should return the "cursor after the SELECT statement has been executed." Since the results of a SELECT statement are mostly retrieved via an iterator, the database connection must be maintained until all results are exhausted. This also fits the use case for contextmanager.
Create a class that simply executes a SQL query and returns its cursor, as shown below.
class ExecuteService:
def __init__(self, host, port, user, password, database):
self.client = MySQLClient(host, port, user, password, database)
@contextlib.contextmanager
def execute(query, params = None):
with self.client.get_cursor() as cursor:
cursor.execute(query, params)
yield cursor
By doing this, for example, when writing SQL results to a file, you can delegate the file-related information to a separate class. It might look something like this:
class BusinessLogic:
def __init__(self, *config):
self.execute_service = ExecuteService(*config)
def run(self):
query = 'select * from user where status = %(status)s'
params = {'status': 'active'}
output_file = Path('output.csv')
with execute_service.execute(query, params) as cursor, open(output_file, 'w') as file:
writer = csv.Dictwriter(file)
for row_dict in cursor:
writer.writerow(row_dict)
When writing to a file, information such as character encoding, line endings, filename, and path is required.
If you want to consolidate and separate these as well, it's a good idea to create a further separate class.
If there's a lot of logic involved in converting the SELECT results into the content to be written, it's also helpful to extract a class like a dataclass for that purpose.
Summary
For things that require some form of termination process, such as connections, using contextmanager to enable the with statement results in cleaner code. This also makes it easier to separate connection management from business logic.
Furthermore, regarding database access, if a library complies with PEP 249, the techniques in this article can be applied almost as they are.
Discussion