yuemori/activerecord-shard_for

module ActiveRecord
  module ShardFor
    module DatabaseTasks
      module_function

      # @return [Boolean]
      def ar6?
        ActiveRecord::VERSION::MAJOR == 6
      end

      # @return [Boolean]
      def ar5?
        ActiveRecord::VERSION::MAJOR == 5
      end

      # @return [Boolean]
      def ar4?
        ActiveRecord::VERSION::MAJOR == 4
      end

      # @return [Boolean]
      def ar42?
        ar4? && ActiveRecord::VERSION::MINOR == 2
      end

      # @return [Boolean]
      def ar41?
        ar4? && ActiveRecord::VERSION::MINOR == 1
      end

      # @return [Boolean]
      def ar417_above?
        ar41? && ActiveRecord::VERSION::TINY > 7
      end

      # Show information of database sharding config.
      def info
        puts 'All clusters registered to ActiveRecord::ShardFor'
        puts
        clusters.each do |cluster|
          puts "= Cluster: #{cluster.name} ="
          cluster.connections.each do |name|
            puts "- #{name}"
          end
          puts
        end
      end

      # @private
      # @param [String] task_name
      # @return [Rake::Task]
      def to_rake_task(task_name)
        Rake::Task[task_name]
      end

      # @private
      # @return [Array<Symbol>]
      def cluster_names
        ActiveRecord::ShardFor.config.cluster_configs.keys
      end

      # @private
      # @return [Array<ActiveRecord::ShardFor::ClusterConfig>]
      def clusters
        ActiveRecord::ShardFor.config.cluster_configs.values
      end

      # @private
      # @return [ActiveRecord::ShardFor::ClusterConfig]
      # @raise [KeyError]
      def fetch_cluster_config(cluster_name)
        ActiveRecord::ShardFor.config.fetch_cluster_config(cluster_name)
      end

      # For mock-ablity
      # @private
      def exit_with_error
        exit 1
      end

      module TasksForMultipleClusters
        # @param [String] task_name
        def invoke_task_for_all_clusters(task_name)
          cluster_names.each do |cluster_name|
            invoke_task(task_name, cluster_name)
          end
        end

        # @private
        # @param [String] name
        # @param [Symbol] cluster_name
        def invoke_task(name, cluster_name)
          task_name = "activerecord:shard_for:#{name}"
          to_rake_task(task_name).invoke(cluster_name.to_s)
          to_rake_task(task_name).reenable
        end
      end
      extend TasksForMultipleClusters

      # Organize cluster config and handle error for invalid args, call single
      # cluster task with each single connection config.
      module TaskOrganizerForSingleClusterTask
        # @param [Hash{Symbol => String}] args
        def create_all_databases(args)
          exec_task_for_all_databases('create', args)
        end

        # @param [Hash{Symbol => String}] args
        def drop_all_databases(args)
          exec_task_for_all_databases('drop', args)
        end

        # @param [Hash{Symbol => String}] args
        def load_schema_all_databases(args)
          exec_task_for_all_databases('load_schema', args)
        end

        private

        # @param [String] task_name
        # @param [Hash{Symbol => String}] args
        def exec_task_for_all_databases(task_name, args)
          cluster_name = cluster_name_or_error(task_name, args)
          cluster = cluster_or_error(cluster_name)
          cluster.connections.each do |connection_name|
            __send__(task_name, connection_name.to_s)
          end
        end

        # @param [String] name A task name
        # @param [Hash{Symbol => String}] args
        # @return [String]
        def cluster_name_or_error(name, args)
          cluster_name = args[:cluster_name]
          return cluster_name if cluster_name

          $stderr.puts <<-MSG
  Missing cluster_name. Find cluster_name via `rake activerecord:shard_for:info` then call `rake "activerecord:shard_for:#{name}[$cluster_name]"`.
          MSG
          exit_with_error
        end

        # @param [String] cluster_name
        # @return [ActiveRecord::ShardFor::ClusterConfig]
        def cluster_or_error(cluster_name)
          fetch_cluster_config(cluster_name.to_sym)
        rescue KeyError
          $stderr.puts %(!cluster name "#{cluster_name}" not found.!)
          exit_with_error
        end
      end
      extend TaskOrganizerForSingleClusterTask

      # Create, drop, load_schema for single connection config.
      module TasksForSingleConnection
        # @param [String] connection_name
        def create(connection_name)
          configuration = ActiveRecord::Base.configurations[connection_name]
          ActiveRecord::Tasks::DatabaseTasks.create(configuration)
          # Re-configure using configuration with database
          ActiveRecord::Base.establish_connection(configuration)
        end

        # @param [String] connection_name
        def drop(connection_name)
          configuration = ActiveRecord::Base.configurations[connection_name]
          ActiveRecord::Tasks::DatabaseTasks.drop(configuration)
        end

        # @param [String] connection_name
        def load_schema(connection_name)
          configuration = ActiveRecord::Base.configurations[connection_name]

          case
          when ar5? || ar6?
            ActiveRecord::Tasks::DatabaseTasks.load_schema(configuration, :ruby)
          when ar42? || ar417_above?
            ActiveRecord::Tasks::DatabaseTasks.load_schema_for(configuration, :ruby)
          when ar41?
            ActiveRecord::Base.establish_connection(configuration)
            ActiveRecord::Tasks::DatabaseTasks.load_schema(:ruby)
          else
            raise "This version of ActiveRecord is not supported: v#{ActiveRecord::VERSION::STRING}"
          end
        end
      end
      extend TasksForSingleConnection
    end
  end
end