データベースごとにCREATE文を書き分けるのが面倒なのでSQLAlchemyでテーブル定義を書きたい話

こんにちは。
ECシステム開発チームのいまづです。
みなさん、sqlcつかっていますか？

なんでsqlc使っていますかなのかと言いますと、ちょっと試しに作ってみたいものがあってsqlcとechoあたりで簡単なAPIサーバーを書いてみようとしていたんですね。
で、sqlcを使い始めるにあたって、スキーマ定義のSQL文を書く（MySQLの例）のですね。
ところがSQLiteでまずやってみた後、MySQLに乗せ換えてみようとするとCREATE TABLE文とかはデータベースの種類によって変わってきますよね。

めんどくさいですね。

sqlcはマイグレーションツールに対応しているので、それを使えよ、というのはそうなのですが、SQLAlchemyで書くのに慣れているので今回はSQLAlchemyを使ってCREATE文を出力することにしました。

Pythonを多用していることもあり、書いておけば後でPythonで使いたくなったときにそのまま使えますしね。

SQLAlchemyにはDDLを出力する用に create_mock_engine関数が用意されています。これを使います。

ほとんど書いてある例そのままですが、このままでは動かなくないですか？

SQLAlchemyのドキュメントから引用

from sqlalchemy import create_mock_engine

def dump(sql, *multiparams, **params):
    print(sql.compile(dialect=engine.dialect))

engine = create_mock_engine('postgresql+psycopg2://', dump)
metadata.create_all(engine, checkfirst=False)

create_mock_engine に指定している dump 関数の中で create_mock_engine が作成するはずの engine が指定されている？

よくわからなかったので、こんな感じにしました。
テーブル定義は Creating and Dropping Database Tables の例を使っています。

"""SQLAlchemyでテーブル定義を書いて、dialectに応じたCREATE文を出力する。
"""
from typing import Dict, Any

from sqlalchemy import (
    create_mock_engine,
    MetaData,
    ForeignKey,
    Table,
    Column,
    Integer,
    String,
    URL,
)


# SQLAlchemyのドキュメントにあるテーブル定義の例
# https://docs.sqlalchemy.org/en/20/core/metadata.html#creating-and-dropping-database-tables
metadata_obj = MetaData()

user = Table(
    "user",
    metadata_obj,
    Column("user_id", Integer, primary_key=True),
    Column("user_name", String(16), nullable=False),
    Column("email_address", String(60), key="email"),
    Column("nickname", String(50), nullable=False),
)

user_prefs = Table(
    "user_prefs",
    metadata_obj,
    Column("pref_id", Integer, primary_key=True),
    Column("user_id", Integer, ForeignKey("user.user_id"), nullable=False),
    Column("pref_name", String(40), nullable=False),
    Column("pref_value", String(100)),
)


class Executor:
    """create_mock_engine で指定する executor でdialectを指定可能にするためのクラス
    """

    def __init__(self, dialect_cls):
        self.dialect_cls = dialect_cls

    def dump(self, sql, *multiparams, **params):
        """create_mock_engine で指定する executor
        """
        print('{};\n'.format(
            str(
                sql.compile(dialect=self.dialect_cls())
            ).strip()
        ))


def main(dialect_name: str, **kwargs: Dict[str, Any]):
    url = URL.create(
        drivername=dialect_name,  # SQLを出力するだけなら、`dialect+driver` ではなくてもよい。
    )
    executor = Executor(url.get_dialect())
    engine = create_mock_engine(url, executor=executor.dump)
    metadata_obj.create_all(engine)


if __name__ == '__main__':
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument(
        '--dialect-name',
        default='sqlite',
        help='dialect name. e.g. sqlite, mysql, postgresql'
    )

    args = parser.parse_args()
    main(**vars(args))

Usageはこうです。

usage: schema.py [-h] [--dialect-name DIALECT_NAME]

options:
  -h, --help            show this help message and exit
  --dialect-name DIALECT_NAME
                        dialect name. e.g. sqlite, mysql, postgresql

これをオプションを指定せずに実行するとSQLite用のSQLが出力されます。

$ rye run python schema.py
CREATE TABLE user (
        user_id INTEGER NOT NULL,
        user_name VARCHAR(16) NOT NULL,
        email_address VARCHAR(60),
        nickname VARCHAR(50) NOT NULL,
        PRIMARY KEY (user_id)
);

CREATE TABLE user_prefs (
        pref_id INTEGER NOT NULL,
        user_id INTEGER NOT NULL,
        pref_name VARCHAR(40) NOT NULL,
        pref_value VARCHAR(100),
        PRIMARY KEY (pref_id),
        FOREIGN KEY(user_id) REFERENCES user (user_id)
);

MySQLを指定した場合。

 rye run python schema.py --dialect-name mysql
CREATE TABLE user (
        user_id INTEGER NOT NULL AUTO_INCREMENT,
        user_name VARCHAR(16) NOT NULL,
        email_address VARCHAR(60),
        nickname VARCHAR(50) NOT NULL,
        PRIMARY KEY (user_id)
);

CREATE TABLE user_prefs (
        pref_id INTEGER NOT NULL AUTO_INCREMENT,
        user_id INTEGER NOT NULL,
        pref_name VARCHAR(40) NOT NULL,
        pref_value VARCHAR(100),
        PRIMARY KEY (pref_id),
        FOREIGN KEY(user_id) REFERENCES user (user_id)
);

PostgreSQLの場合。

 rye run python schema.py --dialect-name postgresql
CREATE TABLE "user" (
        user_id SERIAL NOT NULL,
        user_name VARCHAR(16) NOT NULL,
        email_address VARCHAR(60),
        nickname VARCHAR(50) NOT NULL,
        PRIMARY KEY (user_id)
);

CREATE TABLE user_prefs (
        pref_id SERIAL NOT NULL,
        user_id INTEGER NOT NULL,
        pref_name VARCHAR(40) NOT NULL,
        pref_value VARCHAR(100),
        PRIMARY KEY (pref_id),
        FOREIGN KEY(user_id) REFERENCES "user" (user_id)
);

これで自分でデータベースごとに書き分けなくて良さそうです。

INSERT文などのSQLが欲しい場合も、create_mock_engine で作ったインスタンス engine に対してSQL文の実行コードを書いてやれば同様にSQL文が出力されます。

    conn = engine.connect()
    conn.execute(
        insert(user).
        values(user_name="user1", email="user1@example.com", nickname="user1")
    )

これを追加した場合に出力されるSQLの例。

INSERT INTO user (user_name, email_address, nickname) VALUES (?, ?, ?);

変数をバインドした結果が欲しい場合は、compileする際に compile_kwargsでliteral_bindsを指定します。

class Executor:
    """create_mock_engine で指定する executor でdialectを指定可能にするためのクラス
    """

    def __init__(self, dialect_cls):
        self.dialect_cls = dialect_cls

    def dump(self, sql, *multiparams, **params):
        """create_mock_engine で指定する executor
        """
        print('{};\n'.format(
            str(
                sql.compile(
                    dialect=self.dialect_cls(),
                    compile_kwargs={
                        "literal_binds": True,
                    },
                )
            ).strip()
        ))

出力されるINSERT文。

INSERT INTO user (user_name, email_address, nickname) VALUES ('user1', 'user1@example.com', 'user1');

うまくつくれば、sqlcで必要なクエリ定義にも使えるかもしれませんね。
では！

システムエンジニア募集中です！

興味のある方はぜひカジュアル面談へ！お待ちしています。