개인적으로 개발을 하는데 DB Schema는 보존하면서 DB에 있는 모든 내용은 지우는 작업을 해야 할 필요가 생겼습니다. 테스트용 Dummy data를 발생시켜야 하는데 중복되는 dummy가 발생되면 안 되기 때문이었죠. 단순히 생각하면 모든 테이블을 지우고 새로 만들어버리면 되지만, 별로 그러고 싶지 않았습니다. 저는 SQLAlchemy와 Alembic의 조합으로 개발중이었는데, 테이블을 지우고 새로 만드는 것을 하려면 다음과 같은 작업을 해야할 필요가 있습니다.

  1. metadata에서 .drop_all() 하여 모든 테이블을 없앤다
  2. Alembic의 내부 버전 관리를 None으로 내린다.
  3. Alembic을 통해서 None에서부터 head까지 올린다.
  4. dummy data를 넣는다.

여기서 3번 과정을 그냥 .create_all()하면 모든 과정이 생략되고 최신의 DB가 생성됩니다. 문제는 너무나도 최신의 상태가 되어버려서, Alembic의 head와 일치하지 않는 상태가 되어버립니다.

발생 가능한 상황은 이런 식입니다.

  1. alembic의 head까지 이미 upgrade가 완료되있으나 아직 생성되지 않은 revision에 포함되어야 할 정보로 인해 Schema 불일치
  2. alembic이 아직 head까지 upgrade가 되있지도 않은데 head, 혹은 그 이상의 정보로 Schema가 생성되서 불일치

하지만 이 문제를 해결하려면 Alembic을 상당히 많이 뜯어봐야할 것 같은 예감이 들었습니다. 단순히 Database 내의 자료를 모두 지우기만 하면 되는 문제인데 너무 멀리 나가는 느낌이 들었습니다. 저는 이미 MySQL에 truncate table이라는 편리한 기능이 있음을 알고 있으니 그냥 이걸 쓰고 싶었습니다.

SQLAlchemy에서 RAW DDL Query 날리기

SQLAlchemy에서는 transaction을 시작하고 명령을 내릴 수 있는 방법이 두 가지 있는데, 하나는 engine 레벨이고 또 하나는 session 레벨입니다. 제가 하려는 것은 Truncate, 즉 DDL에 해당하는 Query이므로 engine레벨에서 작업하기로 했습니다.

engine = get_database_engine() # 여러분의 engine을 가져오세요!
with engine.begin() as conn:
    conn.execute('TRUNCATE TABLE users;') # 아직 MySQL만 되요!

이런 식으로 하면 DDL을 실행할 수 있습니다.

테이블 목록 받아오기

하지만 테이블을 모두 타이핑 해서 적어야 한다면 매우 귀찮은 작업일 것입니다. 이미 SQLAlchemy는 알고 있을테니 내부 변수에서 캐내보기로 했습니다. MetaData에겐 sorted_tables라는 속성과 tables라는 속성이 있습니다. 전자는 이미 foreign key constraint등을 고려하여 정렬된 table 순서입니다. 가장 의존되는 table이 맨 처음 나오게 되어 있습니다. 따라서 sorted_tables를 뒤집어서 접근하면 foreign key constraint를 최소한으로 건드리고도 작업이 가능합니다.

engine = get_database_engine() # 여러분의 engine을 가져오세요!
metadata = Base.metadata # 여러분의 Base를 가져오세요!
with engine.begin() as conn:
    for table in reversed(metadata.sorted_tables):
        conn.execute('TRUNCATE TABLE {};'.format(table.name)) # 아직 MySQL만 되요!

하지만, 이렇게 해도 foreign key constraint에서 완전히 해방되진 못합니다.

Foreign Key Constrant를 잠시 끄기

다행히도 제가 주로 쓰는 3개의 RDBMS는 모두 constraint 문제를 잠시 덮어둘 수 있습니다.

MySQL1/MariaDB2

SET FOREIGN_KEY_CHECKS=0;

이렇게 하면 외래키 검사를 하지 않습니다. 되돌리는 방법은 0 대신 1을 넣으면 됩니다. 기본적으로 세션 단위입니다.

PostgreSQL3

SET CONSTRAINTS ALL DEFERRED;

이렇게 하면 현 transaction에서 모든 검사를 하지 않습니다.

SQLite4

PRAGMA foreign_keys = OFF;

이렇게 하면 각 connection마다 검사를 하지 않습니다. 주의할 점은, sqlite는 기본값이 OFF입니다.

Truncate table

DBMS별로 constraint 종료법이 이정도로 다른 것으로 볼 때, truncate 또한 다를 것이 명백합니다.

MySQL5/MariaDB6

TRUNCATE TABLE tb_name;

PostgreSQL7

기본적으로 MySQL과 용법이 같습니다만 옵션을 더 줄 수 있습니다.

TRUNCATE TABLE tb_name RESTART IDENTITY CASCADE;

RESTART IDENTITY로 인해 sequence의 번호가 처음부터 시작됩니다. CASCADE로 인해 foreign key에 걸려도 무시합니다.

SQLite8

사실 가상 기상천외한 동작방식은 SQLite로부터 찾을 수 있는데, 다른 DDL이 필요 없습니다.

DELETE FROM tb_name;

SQLite는 WHERE절을 생략하면 자동으로 truncate에 최적화된 동작을 수행한다고 합니다.

조합하기

조합해보면 다음과 같습니다.

일단 모든 경우를 if문으로 표기해봤습니다.

engine = get_database_engine() # 여러분의 engine을 가져오세요!
metadata = Base.metadata # 여러분의 Base를 가져오세요!
engine_name = engine.name
with engine.begin() as conn:
    if engine_name == 'mysql':
        conn.execute('SET FOREIGN_KEY_CHECKS=0;')
    elif engine_name == 'postgresql':
        conn.execute('SET CONSTRAINTS ALL DEFERRED;')
    elif engine_name == 'sqlite':
        conn.execute('PRAGMA foreign_keys = OFF;')

    for table in reversed(metadata.sorted_tables):
        if engine_name == 'mysql':
            conn.execute('TRUNCATE TABLE {};'.format(table.name))
        elif engine_name == 'postgresql':
            conn.execute('TRUNCATE TABLE {} RESTART IDENTITY CASCADE;'.format(table.name))
        elif engine_name == 'sqlite':
            conn.execute('DELETE FROM {};'.format(table.name))

    if engine_name == 'mysql':
        conn.execute('SET FOREIGN_KEY_CHECKS=1;')
    elif engine_name == 'postgresql':
        conn.execute('SET CONSTRAINTS ALL IMMEDIATE;')
    elif engine_name == 'sqlite':
        conn.execute('PRAGMA foreign_keys = ON;')

하지만 소스가 너무 길어지는 느낌이 있으므로 조금 더 정리해보자면 다음과 같아집니다.

engine = get_database_engine() # 여러분의 engine을 가져오세요!
metadata = Base.metadata # 여러분의 Base를 가져오세요!
engine_name = engine.name
foreign_key_turn_off = {
    'mysql': 'SET FOREIGN_KEY_CHECKS=0;',
    'postgresql': 'SET CONSTRAINTS ALL DEFERRED;',
    'sqlite': 'PRAGMA foreign_keys = OFF;',
}
foreign_key_turn_on = {
    'mysql': 'SET FOREIGN_KEY_CHECKS=1;',
    'postgresql': 'SET CONSTRAINTS ALL IMMEDIATE;',
    'sqlite': 'PRAGMA foreign_keys = ON;',
}
truncate_query = {
    'mysql': 'TRUNCATE TABLE {};',
    'postgresql': 'TRUNCATE TABLE {} RESTART IDENTITY CASCADE;',
    'sqlite': 'DELETE FROM {};',
}

with engine.begin() as conn:
    conn.execute(foreign_key_turn_off[engine.name])

    for table in reversed(metadata.sorted_tables):
        conn.execute(truncate_query[engine.name].format(table.name))

    conn.execute(foreign_key_turn_on[engine.name])