🐍

contextmanagerを使ったconnectionの管理 (mysql-connector-python)

2021/12/28に公開

はじめに

Qiitaとダブルポストです。

DBアクセスする場合、cursorとDBへのconnectionを、使用が終わったタイミングで、それぞれのclose処理を呼ぶ必要があります。(やらないと、コネクションが残ったままになる可能性があり、同時接続数の上限になったタイミングで、正しい動作をしなくなるなど、なんらかのバグの要因になるためです)

Pythonでは、with構文を使うことで、自動でclose処理が呼ばれます。いわゆる、リソース管理のための機能を、言語の標準機能として提供しているということです。Javaでいうところの、try-with-resources構文と同じ役割です。

https://docs.python.org/3.9/reference/compound_stmts.html#the-with-statement

ただ、ライブラリによっては、cursorやconnectionなどのオブジェクトが、with構文に対応していないこともあります。
この場合、自力でclose処理を、適正なタイミングで呼び出すように、毎回実装するのは少々つらいものがあります。

そういった問題を解消するために、関数の返り値をwith構文で使えるものにwrapするためのdecolatorで、標準機能で用意されています。それがcontextmanagerです。

https://docs.python.org/3.9/library/contextlib.html?highlight=contextmanager#contextlib.contextmanager

mysql-connector-pythonでの実践例

この記事を書いてるときに確認したら、最新のバージョンでは、cursor,connectionともに、enter, __exit__が実装されて、with構文が使えるようになっていました。

バージョンが2.2系のときはなかったので、どこかで対応したようです。

connectionとcursor

まず、connectionとcursorを管理するためのクラスをつくります。

mysql_client.py

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
           

普通、contextmanagerでは、try-finallyを使って、finallyでclose処理を呼び出すのですが、使っているものが、with構文を使えるようになっていれば、今回のような書き方ができます。
(実態は、複数のリソース管理しているのを、見た目上はひとつにまとめた、というのが今回の使い方になります)

上記の書き方で、デフォルトでconnection poolがつくられます。

https://dev.mysql.com/doc/connector-python/en/connector-python-connectargs.html

この書き方で、「cursorを使い終わったら、connectionも切断(connection poolに返す)」になります。
使い方としてはこんな感じです。

client = MySQLClient(...)
with client.get_cursor() as cursor:
     query = 'select * from user'
     cursor.execute(query)
     for row in cursor:
         print(row) 

select文の結果と、加工処理の分離

コネクション管理はwith構文におしつけられましたが、実践ではまだすこし問題があります。
このままでは、「select文の結果の取得」と「結果をもとに加工処理(ビジネスロジック)」が分離できないです。

だいたいは、select文の結果をもとに、webアプリなら、それを画面に表示しますし、バッチなら、何かしらのファイルを作成します。この部分は、ある種のビジネスロジックなので、DBアクセスというインフラ部分とは、クラスを分けられるようにしたいです。

そのためには、「select文を実行後のcursor」を返却させればよいです。iteratorでselect文の結果を取得するのが大半なので、結果を取り尽くすまでは、DB connectionを維持する必要があります。
これもcontextmanagerのユースケースにあてはまります。

下記のようにして、SQLを実行してそのcursorを返すだけのクラスを作成します。

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 
               

このようにすると、例えば、SQLの結果をファイルに書き出す場合、ファイルに関する情報を、別のクラスに持たせられます。例えば、下記のような感じです。

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) 
        

ファイルを書き出す場合は、文字コード、改行コード、ファイル名、パスなどの情報が必要です。
これも一箇所にまとめて分離したければ、更にクラスを切り出すとよいです。
select文の結果 -> 書き出す内容 に変換するための、dataclassのようなクラスも、変換の内容がおおければ、切り出すとよいです。

まとめ

コネクションなど、なんらかの終了処理が必要なものは、contextmanagerを使って、with構文を使えるようにすると、すっきりしたコードになります。それによって、コネクション管理とビジネスロジックの分離もしやすくなります。

更に、DBアクセスについていえば、PEP 249を満たしているライブラリなら、この記事のテクニックがほぼそのまま使えます。

Discussion