gitlab-org--gitlab-foss/scripts/regenerate-schema

#!/usr/bin/env ruby

# frozen_string_literal: true

require 'net/http'
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

  def execute
    Dir.chdir(File.expand_path('..', __dir__)) do
      checkout_ref
      checkout_clean_schema
      hide_migrations
      reset_db
      unhide_migrations
      migrate
    ensure
      unhide_migrations
    end
  end

  private

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

    run %Q[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, download it from the server.
  def remote_checkout_clean_schema
    return false unless project_url

    uri = URI.join("#{project_url}/", 'raw/', "#{merge_base}/", FILENAME)

    download_schema(uri)
  end

  ##
  # Download the schema from the given +uri+.
  def download_schema(uri)
    puts "Downloading #{uri}..."

    Net::HTTP.start(uri.host, uri.port, use_ssl: true) do |http|
      request = Net::HTTP::Get.new(uri.request_uri)
      http.read_timeout = 500
      http.request(request) do |response|
        raise("Failed to download file: #{response.code} #{response.message}") if response.code.to_i != 200

        File.open(FILENAME, 'w') do |io|
          response.read_body do |chunk|
            io.write(chunk)
          end
        end
      end
    end

    true
  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 %Q[git checkout #{merge_base} -- #{FILENAME}]
    run %Q[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

  ##
  # Run rake task to reset the database.
  def reset_db
    run %q[bin/rails db:reset RAILS_ENV=test]
  end

  ##
  # Run rake task to run migrations.
  def migrate
    run %q[bin/rails db:migrate RAILS_ENV=test]
  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"
    ret = system(cmd)
    puts "\e[0m"
    raise("Command failed") unless ret
  end

  ##
  # Return the base commit between source and target branch.
  def merge_base
    @merge_base ||= `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'] || '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 project URL from CI environment variable.
  def project_url
    ENV['CI_PROJECT_URL']
  end

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

SchemaRegenerator.new.execute