regenerate-schema

#!/usr/bin/env ruby

# frozen_string_literal: true

ENV['RAILS_ENV'] = 'test'

require 'optparse'
require 'open3'
require 'fileutils'
require 'uri'

class SchemaRegenerator
  ##
  # Filename of the schema
  #
  # This file is being regenerated by this script.
  FILENAME = 'db/structure.sql'

  ##
  # Directories where migrations are stored
  #
  # The methods +hide_migrations+ and +unhide_migrations+ will rename
  # these to disable/enable migrations.
  MIGRATION_DIRS = %w[db/migrate db/post_migrate].freeze

  ##
  # Directory where we store schema versions
  #
  # The remove_schema_migration_files removes files added in this
  # directory when it runs.
  SCHEMA_MIGRATIONS_DIR = 'db/schema_migrations/'

  def initialize(options)
    @rollback_testing = options.delete(:rollback_testing)
    @init_schema_loading = options.delete(:init_schema_loading)
  end

  def execute
    Dir.chdir(File.expand_path('..', __dir__)) do
      # Note: `db:drop` must run prior to hiding migrations.
      #
      # Executing a Rails DB command e.g., `reset`, `drop`, etc. triggers running the initializers.
      # During the initialization, the default values for `application_settings` need to be set.
      # Depending on the presence of migrations, the default values are either faked or inserted.
      #
      # 1. If no migration is detected, all the necessary columns are in place from `db/structure.sql`.
      # The default values can be inserted into `application_settings` table.
      #
      # 2. If a migration is detected, at least one column may be missing from `db/structure.sql`
      # and needs to be added through the detected migration. In this case, the default values are faked.
      # If not, an error would be raised e.g., "NoMethodError: undefined method `some_setting`"
      #
      # See https://gitlab.com/gitlab-org/gitlab/-/merge_requests/135085#note_1628210334 for more info.
      #
      load_tasks
      drop_db
      checkout_ref
      checkout_clean_schema
      hide_migrations
      remove_schema_migration_files
      stop_spring
      setup_db
      unhide_migrations
      migrate
      dump_schema
      rollback if @rollback_testing
    ensure
      unhide_migrations
    end
  end

  private

  def load_tasks
    require_relative '../config/environment'
    Gitlab::Application.load_tasks
  end

  ##
  # Git checkout +CI_COMMIT_SHA+.
  #
  # When running from CI, checkout the clean commit,
  # not the merged result.
  def checkout_ref
    return unless ci?

    run %(git checkout #{source_ref})
    run %q(git clean -f -- db)
  end

  ##
  # Checkout the clean schema from the target branch
  def checkout_clean_schema
    remote_checkout_clean_schema || local_checkout_clean_schema
  end

  ##
  # Get clean schema from remote servers
  #
  # This script might run in CI, using a shallow clone, so to checkout
  # the file, fetch the target branch from the server.
  def remote_checkout_clean_schema
    return false unless project_url
    return false unless target_project_url

    run %(git remote add target_project #{target_project_url}.git)
    run %(git fetch target_project #{target_branch}:#{target_branch})

    local_checkout_clean_schema
  end

  ##
  # Git checkout the schema from target branch.
  #
  # Ask git to checkout the schema from the target branch and reset
  # the file to unstage the changes.
  def local_checkout_clean_schema
    run %(git checkout #{merge_base} -- #{FILENAME})
    run %(git reset -- #{FILENAME})
  end

  ##
  # Move migrations to where Rails will not find them.
  #
  # To reset the database to clean schema defined in +FILENAME+, move
  # the migrations to a path where Rails will not find them, otherwise
  # +db:reset+ would abort. Later when the migrations should be
  # applied, use +unhide_migrations+ to bring them back.
  def hide_migrations
    MIGRATION_DIRS.each do |dir|
      File.rename(dir, "#{dir}__")
    end
  end

  ##
  # Undo the effect of +hide_migrations+.
  #
  # Place back the migrations which might be moved by
  # +hide_migrations+.
  def unhide_migrations
    error = nil

    MIGRATION_DIRS.each do |dir|
      File.rename("#{dir}__", dir)
    rescue Errno::ENOENT
      nil
    rescue StandardError => e
      # Save error for later, but continue with other dirs first
      error = e
    end

    raise error if error
  end

  ##
  # Remove files added to db/schema_migrations
  #
  # In order to properly reset the database and re-run migrations
  # the schema migrations for new migrations must be removed.
  def remove_schema_migration_files
    (untracked_schema_migrations + committed_schema_migrations).each do |schema_migration|
      FileUtils.rm(schema_migration)
    end
  end

  ##
  # List of untracked schema migrations
  #
  # Get a list of schema migrations that are not tracked so we can remove them
  def untracked_schema_migrations
    git_command = "git ls-files --others --exclude-standard -- #{SCHEMA_MIGRATIONS_DIR}"
    run(git_command).chomp.split("\n")
  end

  ##
  # List of untracked schema migrations
  #
  # Get a list of schema migrations that have been committed since the last
  def committed_schema_migrations
    git_command = "git diff --name-only --diff-filter=A #{merge_base} -- #{SCHEMA_MIGRATIONS_DIR}"
    run(git_command).chomp.split("\n")
  end

  ##
  # Stop spring before modifying the database
  def stop_spring
    run %q(bin/spring stop)
  end

  ##
  # Run rake task to drop the database.
  def drop_db
    run_rake_task 'db:drop'
  end

  ##
  # Run rake task to setup the database.
  def setup_db
    run_rake_task(@init_schema_loading ? 'db:create' : 'db:setup')
  end

  ##
  # Run rake task to run migrations.
  def migrate
    run_rake_task 'db:migrate'
  end

  ##
  # Run rake task to dump schema.
  def dump_schema
    run_rake_task 'db:schema:dump'
  end

  ##
  # Run rake task to rollback migrations.
  def rollback
    (untracked_schema_migrations + committed_schema_migrations).sort.reverse_each do |filename|
      version = filename[/\d+\Z/]
      run %(bin/rails db:rollback:main db:rollback:ci RAILS_ENV=test VERSION=#{version})
    end
  end

  ##
  # Run the given +cmd+.
  #
  # The command is colored green, and the output of the command is
  # colored gray.
  # When the command failed an exception is raised.
  def run(cmd)
    puts "\e[32m$ #{cmd}\e[37m"
    stdout_str, stderr_str, status = Open3.capture3(cmd)
    puts "#{stdout_str}#{stderr_str}\e[0m"
    raise("Command failed: #{stderr_str}") unless status.success?

    stdout_str
  end

  def run_rake_task(*tasks, env: {})
    Array.wrap(tasks).each do |task|
      env.each { |k, v| ENV[k.to_s] = v.to_s }

      puts "\e[32m$ bin/rails #{task} RAILS_ENV=test #{env.map { |m| m.join('=') }.join(' ')}\e[37m"
      Rake::Task[task].invoke
    end
  end

  ##
  # Return the base commit between source and target branch.
  def merge_base
    @merge_base ||= run("git merge-base #{target_branch} #{source_ref}").chomp
  end

  ##
  # Return the name of the target branch
  #
  # Get source ref from CI environment variable, or read the +TARGET+
  # environment+ variable, or default to +HEAD+.
  def target_branch
    ENV['CI_MERGE_REQUEST_TARGET_BRANCH_NAME'] || ENV['TARGET'] || ENV['CI_DEFAULT_BRANCH'] || 'master'
  end

  ##
  # Return the source ref
  #
  # Get source ref from CI environment variable, or default to +HEAD+.
  def source_ref
    ENV['CI_COMMIT_SHA'] || 'HEAD'
  end

  ##
  # Return the source project URL from CI environment variable.
  def project_url
    ENV['CI_PROJECT_URL']
  end

  ##
  # Return the target project URL from CI environment variable.
  def target_project_url
    ENV['CI_MERGE_REQUEST_PROJECT_URL']
  end

  ##
  # Return whether the script is running from CI
  def ci?
    ENV['CI']
  end
end

if $PROGRAM_NAME == __FILE__
  options = {}

  OptionParser.new do |opts|
    opts.on("-r", "--rollback-testing", String, "Enable rollback testing") do
      options[:rollback_testing] = true
    end

    opts.on("-i", "--init-schema-loading", String,
      "Enable clean migrations starting from the beginning - loading init_structure.sql") do
      options[:init_schema_loading] = true
    end

    opts.on("-h", "--help", "Prints this help") do
      puts opts
      exit
    end
  end.parse!

  SchemaRegenerator.new(options).execute
end