Bài viết này dựa trên ý tưởng của bài “Creating PostgreSQL Test Data with SQL, PL/pgSQL, and Python” [1]. Tuy nhiên, tôi cụ thể hoá từng bước và thêm một số câu hỏi cần thảo luận về việc tạo mock data cho PostgreSQL

Trong quá trình phát triển phần mềm, việc tạo mock data giúp việc test phần mềm dễ hơn. Khi gặp sự cố, có thể tạo các dữ liệu tương tự ở PROD và test ở DEV hoặc STAGE. Ta cũng có thể dùng mock data để test performance của ứng dụng, giả sử lượng data do ứng dụng sinh ra trong 5-10 năm tới.

Trên thị trường, có tool hỗ trợ việc này, cả trả phí lẫn miễn phí. Với các tool miễn phí, thường chỉ giới hạn 1 lượng mock data nhất định, khoảng nhỏ hơn 500 rows. Điều này không đủ để thực hiện load test.

Tôi cũng từng viết các ứng dụng tạo mock data bằng Java. Tuy nhiên, việc này tốt rất nhiều thời gian và đôi lúc khá khó khăn. Bài viết ở [1] gợi ý tôi 1 hướng mới trong việc tạo mock khi sử dụng chính các tool cung cấp bởi database. Với stored procedured, hàm generate_series và hàm random, ta có thể tạo ra nhiều dạng dữ liệu khác nhau.

Bài này tôi sẽ tập trung theo hướng sử dụng PL/pgSQL hơn là SQL.

Tạo schema


CREATE SCHEMA music;

SET search_path = music, $user, public;

/*
 * table schema creation 
 */

CREATE TABLE artists (
    artist_id serial PRIMARY KEY,
    name TEXT NOT NULL
);

CREATE TABLE albums (
    album_id serial PRIMARY KEY,
    artist_id INTEGER REFERENCES artists(artist_id) ON DELETE CASCADE,
    title TEXT NOT NULL,
    released DATE NOT NULL
);

CREATE TABLE genres (
    genre_id serial PRIMARY KEY,
    name TEXT NOT NULL
);

CREATE TABLE album_genres (
    album_id INTEGER REFERENCES albums(album_id) ON DELETE CASCADE,
    genre_id INTEGER REFERENCES genres(genre_id) ON DELETE CASCADE,
    PRIMARY KEY (album_id, genre_id)
);

Tạo lorem text

Đối với các field dạng text như tên, địa chỉ, title, … random text rất khó nhìn và trông không thực. Ta sẽ tạo 1 bảng chứa các word trong tiếng anh, sau đó, tạo 1 hàm để tạo ngẫu nhiên 1 câu chứa n word. Như vậy, sẽ dễ nhìn hơn rất nhiều

Tạo bảng words và thử chạy random từng word

-- Trong hướng dẫn tạo TEMPORARY table, nhưng tôi tạo luôn permanent table, vì lười phải tạo lại mỗi khi kết thúc session  
CREATE TABLE words (word TEXT);

-- when using docker OR remote server, the file need to be in docker or that remote server. 
-- So, user can use import function of dbeaver to import data with local file
-- Can download top 100k english word at: https://gist.githubusercontent.com/h3xx/1976236/raw/bbabb412261386673eff521dddbe1dc815373b1d/wiki-100k.txt
COPY words (word) FROM '/usr/share/dict/words' WHERE word NOT LIKE '%''%';

select count(*) from words ;

-- Randomly order the table and pick the first result
SELECT * FROM words ORDER BY random() LIMIT 1;

Tạo hàm tạo câu với n word ngẫu nhiên

/*
 * Create a function to generate random sentence in plpsql style 
 */
CREATE OR REPLACE FUNCTION generate_random_title(num_words int default 1) 
RETURNS text AS $$
DECLARE
  words text;
BEGIN
  SELECT initcap(array_to_string(array(
    SELECT * FROM words ORDER BY random() LIMIT num_words
  ), ' ')) INTO words;
  RETURN words;
END;
$$ LANGUAGE plpgsql;

-- test function
SELECT generate_random_title(ceil(random() * 4 + _g * 0)::int)
FROM generate_series(1, 4) AS _g;

Tạo hàm cho mỗi bảng

Trước tiên, xoá data trên toàn bộ bảng theo thứ tự

TRUNCATE albums, artists, genres, album_genres;

Hàm tạo mock data cho từng bảng

--- create genres 
CREATE OR REPLACE FUNCTION generate_genres(genre_options text[] default array['Hip Hop', 'Jazz', 'Rock', 'Electronic']) 
RETURNS void AS $$
BEGIN
  -- Convert each array option into a row and insert them into genres table.
  INSERT INTO genres (name) SELECT unnest(genre_options);
END;
$$ LANGUAGE plpgsql;

--------------

-- create artists
CREATE OR REPLACE FUNCTION generate_artists(size int default 10) 
RETURNS void AS $$
DECLARE
  artist_name text;
BEGIN
  FOR i IN 1..size LOOP
    SELECT generate_random_title(ceil(random() * 2)::int) INTO artist_name;
    -- About 50% of the time, add 'DJ ' to the front of the artist's name.
    IF random() > 0.5 THEN
      artist_name = 'DJ ' || artist_name;
    END IF;
    INSERT INTO artists (name)
    SELECT artist_name;
  END LOOP;
END;
$$ LANGUAGE plpgsql;

---------- generate album
CREATE OR REPLACE FUNCTION generate_albums(size int default 10) 
RETURNS void AS $$
DECLARE
  artist_name text;
BEGIN
  FOR i IN 1..size LOOP
    INSERT INTO albums (artist_id, title, released)
    VALUES (
      (SELECT artist_id FROM artists ORDER BY random() LIMIT 1),
      (SELECT generate_random_title(ceil(random() * 2)::int)),
      (now() - '5 years'::interval * random())::date
    );
  END LOOP;
END;
$$ LANGUAGE plpgsql;

--- generate album_genres 
CREATE OR REPLACE FUNCTION generate_albums_genres(size int default 10) 
RETURNS void AS $$
DECLARE
  dj_album RECORD;
BEGIN
  FOR i in 1..size LOOP
    INSERT INTO album_genres (album_id, genre_id)
    VALUES (
      (SELECT album_id FROM albums ORDER BY random() LIMIT 1),
      (SELECT genre_id FROM genres ORDER BY random() LIMIT 1)
    )
    -- If we insert a row that already exists, do nothing (don't raise an error)
    ON CONFLICT DO NOTHING;
  END LOOP;
 
  -- Ensure all albums by a 'DJ' artist belong to the Electronic genre.
  FOR dj_album IN
    SELECT albums.* FROM albums
    INNER JOIN artists USING (artist_id)
    WHERE artists.name LIKE 'DJ %'
  LOOP
    RAISE NOTICE 'Ensuring DJ album % belongs to Electronic genre!', quote_literal(dj_album.title);
    INSERT INTO album_genres (album_id, genre_id)
    SELECT dj_album.album_id, (SELECT genre_id FROM genres WHERE name = 'Electronic')
    -- If we insert a row that already exists, do nothing (don't raise an error)
    ON CONFLICT DO NOTHING;
  END LOOP;
END;
$$ LANGUAGE plpgsql;

Gọi hàm tạo test data

-- create genres 
select * from generate_genres();

select * from genres ;

-- create 10k artists cost 1m2s
select * from generate_artists(10000)

select count(*) from artists ;

-- create 50k alumns cost 5m51s
select generate_albums(50000);

select count(*) from albums ;

-- create 10k albums_genres cost 27s
select generate_albums_genres(10000);

select count(*) from album_genres ;

Thảo luận

  • Việc lấy ngẫu nhiên 1 dòng trong bảng chưa tối ưu, hiện tại ta thực hiện SELECT album_id FROM albums ORDER BY random() LIMIT 1. Việc này cần sort toàn bộ bảng theo random(). Nếu bảng lớn, việc này tốn nhiều thời gian và memory vì ORDER BY cần lấy tất cả rows, sau đó mới thực hiện sort. Để hạn chế việc này, ta có thể dùng câu lệnh thay thế SELECT artist_id FROM artists where artist_id > floor(random() * 10000) limit 1. Việc thực hiện riêng rẻ thì câu lệnh 2 nhanh hơn cầu lệnh gốc, nhưng khi áp dụng vào việc generate mock data, tôi vẫn chưa thấy sự khác biệt. Một nguyên nhân có thể nghĩ đến là dữ liệu hiện tại vẫn đang nhỏ. Nếu có từ 5m - 10m dòng, có lẽ sẽ thấy sự khác biệt ?

References

  1. Creating PostgreSQL Test Data with SQL, PL/pgSQL, and Python