Loading a representative sample of data into a Mastodon database for migration tests

sample_load.sql

DO $$
DECLARE ratio float;
DECLARE statuses_count integer;
DECLARE conversations_count integer;
DECLARE notifications_count integer;
DECLARE media_attachments_count integer;
DECLARE oauth_access_tokens_count integer;
DECLARE stream_entries_count integer;
DECLARE favourites_count integer;
DECLARE mentions_count integer;
DECLARE statuses_tags_count integer;
DECLARE preview_cards_statuses_count integer;
DECLARE accounts_count integer;
DECLARE preview_cards_count integer;
DECLARE users_count integer;
DECLARE subscriptions_count integer;
DECLARE blocks_count integer;
DECLARE session_activations_count integer;
DECLARE follow_requests_count integer;
DECLARE follows_count integer;
DECLARE tags_count integer;
DECLARE oauth_access_grants_count integer;
DECLARE web_push_subscriptions_count integer;
DECLARE oauth_applications_count integer;
DECLARE mutes_count integer;
DECLARE settings_count integer;
DECLARE web_settings_count integer;
BEGIN
  SELECT 0.1 INTO ratio;
  SELECT ceil(33597681 * ratio) INTO statuses_count;
  SELECT ceil(18729250 * ratio) INTO conversations_count;
  SELECT ceil(12705005 * ratio) INTO notifications_count;
  SELECT ceil(2254240 * ratio) INTO media_attachments_count;
  SELECT ceil(1911506 * ratio) INTO oauth_access_tokens_count;
  SELECT ceil(597439 * ratio) INTO stream_entries_count;
  SELECT ceil(507040 * ratio) INTO favourites_count;
  SELECT ceil(405300 * ratio) INTO mentions_count;
  SELECT ceil(377983 * ratio) INTO statuses_tags_count;
  SELECT ceil(320344 * ratio) INTO preview_cards_statuses_count;
  SELECT ceil(287677 * ratio) INTO accounts_count;
  SELECT ceil(193562 * ratio) INTO preview_cards_count;
  SELECT ceil(148902 * ratio) INTO users_count;
  SELECT ceil(148865 * ratio) INTO subscriptions_count;
  SELECT ceil(138923 * ratio) INTO blocks_count;
  SELECT ceil(72603 * ratio) INTO session_activations_count;
  SELECT ceil(19857 * ratio) INTO follow_requests_count;
  SELECT ceil(18868 * ratio) INTO follows_count;
  SELECT ceil(16052 * ratio) INTO tags_count;
  SELECT ceil(11427 * ratio) INTO oauth_access_grants_count;
  SELECT ceil(9101 * ratio) INTO web_push_subscriptions_count;
  SELECT ceil(7666 * ratio) INTO oauth_applications_count;
  SELECT ceil(6247 * ratio) INTO mutes_count;
  SELECT ceil(4074 * ratio) INTO settings_count;
  SELECT ceil(1068 * ratio) INTO web_settings_count;

  -- We need accounts, users, and statuses for a bunch of foreign keys
  INSERT INTO accounts (id, username, created_at, updated_at) SELECT
    generate_series(2,accounts_count),
    generate_series(2,accounts_count),
    NOW(), NOW();
  INSERT INTO users (id, email, account_id, created_at, updated_at) SELECT
    generate_series(2,users_count),
    generate_series(2,users_count),
    generate_series(2,users_count),
    NOW(), NOW();
  INSERT INTO statuses (id, text, account_id, application_id, in_reply_to_account_id, created_at, updated_at) SELECT
    generate_series(1,statuses_count),
    'Test Status Text',
    generate_series(0,statuses_count-1) % accounts_count + 1,
    generate_series(1,statuses_count),
    generate_series(0,statuses_count-1) % accounts_count + 1,
    NOW(), NOW();

  -- Skiping conversations as there are no columns to migrate
  INSERT INTO notifications (id, activity_id, account_id, from_account_id, created_at, updated_at) SELECT
    generate_series(1,notifications_count),
    generate_series(1,notifications_count),
    generate_series(0,notifications_count-1) % accounts_count + 1,
    generate_series(0,notifications_count-1) % accounts_count + 1,
    NOW(), NOW();
  INSERT INTO media_attachments (id, status_id, account_id, created_at, updated_at) SELECT
    generate_series(1,media_attachments_count),
    generate_series(1,media_attachments_count),
    generate_series(0,media_attachments_count-1) % accounts_count + 1,
    NOW(), NOW();
  INSERT INTO oauth_applications (id, name, uid, secret, redirect_uri, owner_id, created_at, updated_at) SELECT
    generate_series(2,oauth_applications_count),
    'OAuth Name',
    generate_series(2,oauth_applications_count),
    'secret',
    'redirect_uri',
    generate_series(2,oauth_applications_count),
    NOW(), NOW();
  INSERT INTO oauth_access_tokens (id, token, refresh_token, resource_owner_id, application_id, created_at) SELECT
    generate_series(1,oauth_access_tokens_count),
    generate_series(1,oauth_access_tokens_count),
    generate_series(1,oauth_access_tokens_count),
    generate_series(0,oauth_access_tokens_count-1) % users_count + 1,
    generate_series(0,oauth_access_tokens_count-1) % oauth_applications_count + 1,
    NOW();
  INSERT INTO stream_entries (id, account_id, created_at, updated_at) SELECT
    generate_series(1,stream_entries_count),
    generate_series(0,stream_entries_count-1) % accounts_count + 1,
    NOW(), NOW();
  INSERT INTO favourites (id, account_id, status_id, created_at, updated_at) SELECT
    generate_series(1,favourites_count),
    generate_series(0,favourites_count-1) % accounts_count + 1,
    generate_series(1,favourites_count),
    NOW(), NOW();
  INSERT INTO mentions (id, account_id, status_id, created_at, updated_at) SELECT
    generate_series(1,mentions_count),
    generate_series(0,mentions_count-1) % accounts_count + 1,
    generate_series(1,mentions_count),
    NOW(), NOW();
  -- Skiping statuses_tags as there are no columns to migrate
  -- Skiping preview_cards_statuses as there are no columns to migrate
  -- Skiping preview_cards as there are no columns to migrate
  INSERT INTO subscriptions (id, callback_url, account_id, created_at, updated_at) SELECT
    generate_series(1,subscriptions_count),
    'callback',
    generate_series(1,subscriptions_count),
    NOW(), NOW();
  INSERT INTO blocks (id, account_id, target_account_id, created_at, updated_at) SELECT
    generate_series(2,blocks_count+1),
    generate_series(2,blocks_count+1),
    generate_series(2,blocks_count+1),
    NOW(), NOW();
  INSERT INTO session_activations (id, session_id, access_token_id, user_id, web_push_subscription_id, created_at, updated_at) SELECT
    generate_series(1,session_activations_count),
    generate_series(1,session_activations_count),
    generate_series(1,session_activations_count),
    generate_series(1,session_activations_count),
    generate_series(1,session_activations_count),
    NOW(), NOW();
  INSERT INTO follow_requests (id, account_id, target_account_id, created_at, updated_at) SELECT
    generate_series(1,follow_requests_count),
    generate_series(1,follow_requests_count),
    generate_series(1,follow_requests_count),
    NOW(), NOW();
  INSERT INTO follows (id, account_id, target_account_id, created_at, updated_at) SELECT
    generate_series(1,follows_count),
    generate_series(1,follows_count),
    generate_series(1,follows_count),
    NOW(), NOW();
  INSERT INTO tags (id, name, created_at, updated_at) SELECT
    generate_series(1,tags_count),
    generate_series(1,tags_count),
    NOW(), NOW();
  INSERT INTO oauth_access_grants (id, token, expires_in, redirect_uri, application_id, resource_owner_id, created_at) SELECT
    generate_series(1,oauth_access_grants_count),
    generate_series(1,oauth_access_grants_count),
    360000,
    'redirect_uri',
    generate_series(0,oauth_access_grants_count-1) % oauth_applications_count + 1,
    generate_series(1,oauth_access_grants_count),
    NOW();
  -- Skipping web_push_subscriptions as there are no columns to migrate
  INSERT INTO mutes (id, account_id, target_account_id, created_at, updated_at) SELECT
    generate_series(1,mutes_count),
    generate_series(1,mutes_count),
    generate_series(1,mutes_count),
    NOW(), NOW();
  INSERT INTO settings (id, thing_id, var, value, created_at, updated_at) SELECT
    generate_series(1,settings_count),
    generate_series(1,settings_count),
    'var',
    'value',
    NOW(), NOW();
  INSERT INTO web_settings (id, user_id, created_at, updated_at) SELECT
    generate_series(1,web_settings_count),
    generate_series(1,web_settings_count),
    NOW(), NOW();

  SELECT SETVAL('account_domain_blocks_id_seq', COALESCE(MAX(id), 1) ) FROM account_domain_blocks;
  SELECT SETVAL('accounts_id_seq', COALESCE(MAX(id), 1) ) FROM accounts;
  SELECT SETVAL('blocks_id_seq', COALESCE(MAX(id), 1) ) FROM blocks;
  SELECT SETVAL('conversation_mutes_id_seq', COALESCE(MAX(id), 1) ) FROM conversation_mutes;
  SELECT SETVAL('conversations_id_seq', COALESCE(MAX(id), 1) ) FROM conversations;
  SELECT SETVAL('custom_emojis_id_seq', COALESCE(MAX(id), 1) ) FROM custom_emojis;
  SELECT SETVAL('deprecated_preview_cards_id_seq', COALESCE(MAX(id), 1) ) FROM deprecated_preview_cards;
  SELECT SETVAL('domain_blocks_id_seq', COALESCE(MAX(id), 1) ) FROM domain_blocks;
  SELECT SETVAL('favourites_id_seq', COALESCE(MAX(id), 1) ) FROM favourites;
  SELECT SETVAL('follow_requests_id_seq', COALESCE(MAX(id), 1) ) FROM follow_requests;
  SELECT SETVAL('follows_id_seq', COALESCE(MAX(id), 1) ) FROM follows;
  SELECT SETVAL('imports_id_seq', COALESCE(MAX(id), 1) ) FROM imports;
  SELECT SETVAL('media_attachments_id_seq', COALESCE(MAX(id), 1) ) FROM media_attachments;
  SELECT SETVAL('mentions_id_seq', COALESCE(MAX(id), 1) ) FROM mentions;
  SELECT SETVAL('mutes_id_seq', COALESCE(MAX(id), 1) ) FROM mutes;
  SELECT SETVAL('notifications_id_seq', COALESCE(MAX(id), 1) ) FROM notifications;
  SELECT SETVAL('oauth_access_grants_id_seq', COALESCE(MAX(id), 1) ) FROM oauth_access_grants;
  SELECT SETVAL('oauth_access_tokens_id_seq', COALESCE(MAX(id), 1) ) FROM oauth_access_tokens;
  SELECT SETVAL('oauth_applications_id_seq', COALESCE(MAX(id), 1) ) FROM oauth_applications;
  SELECT SETVAL('preview_cards_id_seq', COALESCE(MAX(id), 1) ) FROM preview_cards;
  SELECT SETVAL('reports_id_seq', COALESCE(MAX(id), 1) ) FROM reports;
  SELECT SETVAL('session_activations_id_seq', COALESCE(MAX(id), 1) ) FROM session_activations;
  SELECT SETVAL('settings_id_seq', COALESCE(MAX(id), 1) ) FROM settings;
  SELECT SETVAL('site_uploads_id_seq', COALESCE(MAX(id), 1) ) FROM site_uploads;
  SELECT SETVAL('status_pins_id_seq', COALESCE(MAX(id), 1) ) FROM status_pins;
  SELECT SETVAL('statuses_id_seq', COALESCE(MAX(id), 1) ) FROM statuses;
  SELECT SETVAL('stream_entries_id_seq', COALESCE(MAX(id), 1) ) FROM stream_entries;
  SELECT SETVAL('subscriptions_id_seq', COALESCE(MAX(id), 1) ) FROM subscriptions;
  SELECT SETVAL('tags_id_seq', COALESCE(MAX(id), 1) ) FROM tags;
  SELECT SETVAL('users_id_seq', COALESCE(MAX(id), 1) ) FROM users;
  SELECT SETVAL('web_push_subscriptions_id_seq', COALESCE(MAX(id), 1) ) FROM web_push_subscriptions;
  SELECT SETVAL('web_settings_id_seq', COALESCE(MAX(id), 1) ) FROM web_settings;
END $$;