【PostgreSQL】シーケンス採番付きのテストデータを複製してみた

はじめに

こんにちは、新卒2年目のyk_itgです。
業務の中で内容が同じ大量のテストデータが必要なテストがあったのですが、手作業で作成するとたくさんの手順を踏まなくてはならなかったり、入力を間違えたりして大変なので、なんとかSQLでできないか考えてみました。

はじめに
例えばこんなデータ
複製する
関数にしてみる
- 実行結果
終わりに

例えばこんなデータ

例えばこのような感じのデータがあったとします。

CREATE TABLE chohyo (
    chohyo_id INTEGER PRIMARY KEY, 
    status INTEGER, 
    user_name TEXT, 
    title TEXT, 
    hoge_flag INTEGER, 
    total_price INTEGER
);
CREATE SEQUENCE chohyo_sequence;

CREATE TABLE chohyo_meisai (
    chohyo_id INTEGER, 
    meisai_number INTEGER,
    hoge_name TEXT, 
    price INTEGER, 
    PRIMARY KEY(chohyo_id, meisai_number), 
    FOREIGN KEY(chohyo_id) references chohyo (chohyo_id) on delete cascade
);

chohyoとchohyo_meisaiは1：多の関係で、chohyoを作成する毎にシーケンスによってidに値が割り振られていくイメージです。

postgres=# SELECT * FROM chohyo WHERE chohyo_id = 1;
-[ RECORD 1 ]------------
chohyo_id   | 1
status      | 0
user_name   | taro
title       | テスト帳票1
hoge_flag   | 0
total_price | 1000

postgres=# SELECT * FROM chohyo_meisai WHERE chohyo_id = 1;
-[ RECORD 1 ]-+---------
chohyo_id     | 1
meisai_number | 0
hoge_name     | コーヒー
price         | 500
-[ RECORD 2 ]-+---------
chohyo_id     | 1
meisai_number | 1
hoge_name     | パン
price         | 500

今回はこのデータを複製する方法を考えていきます。

複製する

SQLでテスト用のダミーデータを作成する方法としては作成したいデータの値を入れたINSERT文を作る、COPYを使う方法等ありますが、それぞれカラムを追加したときに対応する必要があったり、ファイルを操作する必要があったりで面倒です。
なるべくデータベース内で完結してほしいので、以下の手順で複製していきます。

コピーしたいデータをtmpテーブルにコピーする
変更したい部分のみを書き換える
編集したデータを元のテーブルに挿入する

1. コピーしたいデータをtmpテーブルにコピーする

CREATE TABLE AS文*1を使って、コピーするデータをtmpテーブルにコピーします。
TEMPORARYを指定するとセッションが切れたときにtmpテーブルは削除されるので、不要なデータが残る心配が無くなります。
下のSQLを実行する前に\set id 1やpsql -v id=1などでコピーしたいデータを指定しておいてください。

CREATE TEMPORARY TABLE tmp_chohyo 
AS SELECT * FROM chohyo WHERE chohyo_id = :id;
CREATE TEMPORARY TABLE tmp_chohyo_meisai 
AS SELECT * FROM chohyo_meisai WHERE chohyo_id = :id;

2. 変更したい部分のみを書き換える

tmpテーブルにデータをコピーできたので、次は加工します。
下のSQLではchohyoのstatus以外はそのままにして、chohyo_idをシーケンスによって採番しています。

UPDATE tmp_chohyo SET chohyo_id = nextval('chohyo_sequence'), status = 1;
UPDATE tmp_chohyo_meisai SET chohyo_id = (SELECT chohyo_id FROM tmp_chohyo);

3. 編集したデータを元のテーブルに挿入する

あとは加工したデータを元のテーブルに挿入するだけです。

INSERT INTO chohyo SELECT * FROM tmp_chohyo;
INSERT INTO chohyo_meisai SELECT * FROM tmp_chohyo_meisai;

実行結果

postgres=# SELECT COUNT(*) FROM chohyo WHERE user_name = 'taro';
-[ RECORD 1 ]
count | 2

関数にしてみる

上のSQLではデータごとに1回実行しなくてはならず面倒なので、PL/pgSQL *2の関数にしてみます。
長くなりそうなので説明は省きますが、このような感じになります。

CREATE OR REPLACE FUNCTION copy_chohyo(chohyoId integer, count integer) RETURNS VOID AS $$
DECLARE
    recordCount integer;
    nextId integer;
    resultChohyoCount integer;
    resultMeisaiCount integer;
BEGIN
    SELECT COUNT(*) FROM chohyo WHERE chohyo_id = chohyoId INTO recordCount;
    RAISE NOTICE '件数 = %', recordCount;
    IF recordCount <> 1 THEN
        RETURN;
    END IF;

    -- テーブル構造をコピー
    CREATE TEMPORARY TABLE tmp_chohyo AS SELECT * FROM chohyo WHERE false;
    CREATE TEMPORARY TABLE tmp_chohyo_meisai AS SELECT * FROM chohyo_meisai WHERE false;

    -- コピーする数だけtmpにコピー&加工
    FOR i IN 1..count LOOP
        INSERT INTO tmp_chohyo SELECT * FROM chohyo WHERE chohyo_id = chohyoId;
        INSERT INTO tmp_chohyo_meisai SELECT * FROM chohyo_meisai WHERE chohyo_id = chohyoId;

        -- あらかじめ次の値を保存
        SELECT nextval('chohyo_sequence') AS next_id INTO nextId;

        UPDATE tmp_chohyo SET chohyo_id = nextId WHERE chohyo_id = chohyoId;
        UPDATE tmp_chohyo_meisai SET chohyo_id = nextId WHERE chohyo_id = chohyoId;
    END LOOP;

    -- コピーデータを挿入
    INSERT INTO chohyo SELECT * FROM tmp_chohyo;
    GET DIAGNOSTICS resultChohyoCount = ROW_COUNT;
    RAISE NOTICE 'chohyo の挿入件数 = %', resultChohyoCount;
    INSERT INTO chohyo_meisai SELECT * FROM tmp_chohyo_meisai;
    GET DIAGNOSTICS resultMeisaiCount = ROW_COUNT;
    RAISE NOTICE 'chohyo_meisai の挿入件数 = %', resultMeisaiCount;

    -- 完了後処理
    DROP TABLE tmp_chohyo;
    DROP TABLE tmp_chohyo_meisai;

    RETURN;
END;
$$ LANGUAGE plpgsql;

実行結果

postgres=# SELECT copy_chohyo(1, 5);
NOTICE:  件数 = 1
NOTICE:  chohyo の挿入件数 = 5
NOTICE:  chohyo_meisai の挿入件数 = 10
-[ RECORD 1 ]-
copy_chohyo | 

postgres=# SELECT COUNT(*) FROM chohyo WHERE user_name = 'taro';
-[ RECORD 1 ]
count | 6

終わりに

今回はSQLのみでシーケンス採番付きのテストデータを複製する方法を考えてみましたが、いかがでしたでしょうか。テストでデータを複製する際の参考になれば幸いです。

エンジニア中途採用サイト
ラクスでは、エンジニア・デザイナーの中途採用を積極的に行っております！
ご興味ありましたら是非ご確認をお願いします。

https://career-recruit.rakus.co.jp/career_engineer/
カジュアル面談お申込みフォーム
どの職種に応募すれば良いかわからないという方は、カジュアル面談も随時行っております。
以下フォームよりお申込みください。
forms.gle
イベント情報
会社の雰囲気を知りたい方は、毎週開催しているイベントにご参加ください！ rakus.connpass.com