SQLAlchemyのModelクラスをMySQLから自動生成する

2016/12/23

SQLAlchemyのMySQLモデルを書くのが面倒で、きっと何かツールがあるのでしょうが、探せなかったので作りました。

JOINとかには対応してないので用途によっては生成後書き足す事も必要です。
この例ではAmazon RDSのMySQLを使用しています。
config.pyはそれぞれの環境にあわせて作成してください。
全ての型を網羅していません。自分が取り急ぎ必要な型だけサポートしています。

GitHubのコードはこちらです。

テーブル情報の取得

まず、pymysqlでデータベースに接続して show tables で全テーブル情報の配列を取得します。
書き出し先のテキストファイルも用意します。

conn = pymysql.connect(
    host,
    user=user,
    passwd=password,
    db=db_name,
    charset='utf8',
    cursorclass=pymysql.cursors.DictCursor
)

cur = conn.cursor()
sql = 'show tables'
cur.execute(sql)
tables = cur.fetchall()

write_file = open(db_name + '_models.txt', 'w')

conn = pymysql.connect(

host,

user=user,

passwd=password,

db=db_name,

charset='utf8',

cursorclass=pymysql.cursors.DictCursor

)

cur = conn.cursor()

sql = 'show tables'

cur.execute(sql)

tables = cur.fetchall()

write_file = open(db_name + '_models.txt', 'w')

show tablesの結果

テーブル名がTables_in_[db_name]というカラムで出力されます

Tables_in_db_name
table_name_1
table_name_2
…

カラム情報の取得

テーブル情報配列をループしてtable[‘Tables_in_’ + db_name]でテーブル名を取得します。
そして、show columns from table_name でテーブルの情報を取得します。

table_name = table['Tables_in_' + db_name]
sql = 'show columns from ' + table_name
cur.execute(sql)
columns = cur.fetchall()

table_name = table['Tables_in_' + db_name]

sql = 'show columns from ' + table_name

cur.execute(sql)

columns = cur.fetchall()

show columns from ‘ + table_nameの結果

カラム名、型、Null許可か、プライマリキーか、デフォルト値といった情報が出力されます。
なので、この情報をもとにモデルクラスを書き出します。

Field	Type	Null	Key	Default	Extra
id	int(10)	NO	PRI
name	varchar(20)	YES		NULL

型と長さとunsigned

Type列の出力値からSQLAlchemyの文法にもとづいて生成します。

def get_type(type_str):
    sta = type_str.find('(')
    end = type_str.find(')')
    ts = type_str[0:3]
    length = -1
    unsigned = ''
    ret_type = ''

    if sta > 0:
        length = type_str[sta+1:end]
    if type_str.find('unsigned') > -1:
        unsigned = 'unsigned=True'

    if type_str.find('decimal') > -1:
        ret_type = 'DECIMAL()'
    elif type_str.find('bigint') > -1:
        ret_type = 'BIGINT(' + unsigned + ')'
    elif type_str.find('varchar') > -1:
        ret_type = 'VARCHAR(' + str(length) + ')'
    elif type_str.find('datetime') > -1:
        ret_type = 'DATETIME()'
    elif type_str.find('char') > -1:
        ret_type = 'CHAR(' + str(length) + ')'
    elif type_str.find('text') > -1:
        ret_type = 'TEXT()'
    elif type_str.find('tinyint') > -1:
        ret_type = 'TINYINT(' + unsigned + ')'
    elif type_str.find('time') > -1:
        ret_type = 'TIME()'
    elif type_str.find('int') > -1:
        ret_type = 'INTEGER(' + unsigned + ')'
    return ret_type

def get_type(type_str):

sta = type_str.find('(')

end = type_str.find(')')

ts = type_str[0:3]

length = -1

unsigned = ''

ret_type = ''

if sta > 0:

length = type_str[sta+1:end]

if type_str.find('unsigned') > -1:

unsigned = 'unsigned=True'

if type_str.find('decimal') > -1:

ret_type = 'DECIMAL()'

elif type_str.find('bigint') > -1:

ret_type = 'BIGINT(' + unsigned + ')'

elif type_str.find('varchar') > -1:

ret_type = 'VARCHAR(' + str(length) + ')'

elif type_str.find('datetime') > -1:

ret_type = 'DATETIME()'

elif type_str.find('char') > -1:

ret_type = 'CHAR(' + str(length) + ')'

elif type_str.find('text') > -1:

ret_type = 'TEXT()'

elif type_str.find('tinyint') > -1:

ret_type = 'TINYINT(' + unsigned + ')'

elif type_str.find('time') > -1:

ret_type = 'TIME()'

elif type_str.find('int') > -1:

ret_type = 'INTEGER(' + unsigned + ')'

return ret_type

モデルクラスの書き出し

テーブル名を書き出した後、カラムを1行づつインデントを付けて書き出します。

write_file.write('\n\n')
write_file.write('class ' + table_name.upper() + '(Base):\n')
write_file.write("    __tablename__ = '" + table_name.upper() + "'\n\n")

for column in columns:
    field_str = column['Field']
    type_str = column['Type']
    null_str = column['Null']
    key_str = column['Key']
    pri_str = ''
    null_able = ''
    if key_str == 'PRI':
        pri_str = ' primary_key=True, '
    if null_str == 'NO':
        null_able = 'nullable=False)'
    else:
        null_able = 'nullable=True)'

    line_str = '    '
    line_str += field_str + ' = Column('
    line_str += get_type(type_str) + ', '
    line_str += pri_str
    line_str += null_able

    write_file.write(line_str + '\n')

write_file.write('\n\n')

write_file.write('class ' + table_name.upper() + '(Base):\n')

write_file.write(" __tablename__ = '" + table_name.upper() + "'\n\n")

for column in columns:

field_str = column['Field']

type_str = column['Type']

null_str = column['Null']

key_str = column['Key']

pri_str = ''

null_able = ''

if key_str == 'PRI':

pri_str = ' primary_key=True, '

if null_str == 'NO':

null_able = 'nullable=False)'

else:

null_able = 'nullable=True)'

line_str = ' '

line_str += field_str + ' = Column('

line_str += get_type(type_str) + ', '

line_str += pri_str

line_str += null_able

write_file.write(line_str + '\n')

出来た！

class TABLE_NAME(Base):
    __tablename__ = 'TABLE_NAME'

    ID = Column(INTEGER(unsigned=True),  primary_key=True, nullable=False)
    NAME = Column(VARCHAR(20), nullable=True)