activerecord-sharding 0.1.0

data/lib/active_record/sharding/cluster_config.rb

@@ -0,0 +1,32 @@
+module ActiveRecord
+  module Sharding
+    class ClusterConfig
+      attr_reader :name
+
+      def initialize(name)
+        @name = name
+        @connection_registry = []
+      end
+
+      def register_connection(connection_name)
+        @connection_registry << connection_name
+      end
+
+      def fetch(modulo_key)
+        @connection_registry[modulo_key]
+      end
+
+      def registerd_connection_count
+        @connection_registry.count
+      end
+
+      def validate_config!
+        fail "Nothing registerd connections." if registerd_connection_count == 0
+      end
+
+      def connections
+        @connection_registry
+      end
+    end
+  end
+end

data/lib/active_record/sharding/config.rb

@@ -0,0 +1,34 @@
+module ActiveRecord
+  module Sharding
+    class Config
+      attr_reader :cluster_configs, :sequencer_configs
+
+      def initialize
+        @cluster_configs = {}
+        @sequencer_configs = {}
+      end
+
+      def define_cluster(cluster_name, &block)
+        cluster_config = ClusterConfig.new(cluster_name)
+        cluster_config.instance_eval(&block)
+        cluster_config.validate_config!
+        @cluster_configs[cluster_name] = cluster_config
+      end
+
+      def fetch_cluster_config(cluster_name)
+        @cluster_configs.fetch cluster_name
+      end
+
+      def define_sequencer(sequencer_name, &block)
+        sequencer_config = SequencerConfig.new sequencer_name
+        sequencer_config.instance_eval(&block)
+        sequencer_config.validate_config!
+        @sequencer_configs[sequencer_name] = sequencer_config
+      end
+
+      def fetch_sequencer_config(sequencer_name)
+        @sequencer_configs.fetch sequencer_name
+      end
+    end
+  end
+end

data/lib/active_record/sharding/database_tasks.rb

@@ -0,0 +1,234 @@
+module ActiveRecord
+  module Sharding
+    module DatabaseTasks
+      extend self
+
+      def info
+        puts "All clusters registered to activerecord-sharding"
+        puts
+        clusters.each do |cluster|
+          puts "= Cluster: #{cluster.name} ="
+          cluster.connections.each do |name|
+            puts "- #{name}"
+          end
+          puts
+        end
+        puts_sequencers
+      end
+
+      def puts_sequencers
+        return unless sequencers
+
+        puts "All sequencers registered to activerecord-sharding"
+        puts
+        sequencers.each do |sequencer|
+          puts "= Sequencer: #{sequencer.name} ="
+          puts "- Connection:#{sequencer.connection_name} Table:#{sequencer.table_name}"
+          puts
+        end
+      end
+
+      def ar5?
+        ActiveRecord::VERSION::MAJOR == 5
+      end
+
+      def ar4?
+        ActiveRecord::VERSION::MAJOR == 4
+      end
+
+      def ar42?
+        ar4? && ActiveRecord::VERSION::MINOR == 2
+      end
+
+      def ar41?
+        ar4? && ActiveRecord::VERSION::MINOR == 1
+      end
+
+      def ar417_above?
+        ar41? && ActiveRecord::VERSION::TINY > 7
+      end
+
+      def clusters
+        ActiveRecord::Sharding.config.cluster_configs.values
+      end
+
+      def cluster_names
+        ActiveRecord::Sharding.config.cluster_configs.keys
+      end
+
+      def sequencer_names
+        ActiveRecord::Sharding.config.sequencer_configs.keys
+      end
+
+      def fetch_cluster_config(cluster_name)
+        ActiveRecord::Sharding.config.fetch_cluster_config cluster_name
+      end
+
+      def sequencers
+        ActiveRecord::Sharding.config.sequencer_configs.values
+      end
+
+      def fetch_sequencer_config(sequencer_name)
+        ActiveRecord::Sharding.config.fetch_sequencer_config sequencer_name
+      end
+
+      def to_rake_task(task_name)
+        Rake::Task[task_name]
+      end
+
+      module TasksForMultipleClusters
+        def invoke_task_for_all_clusters(task_name)
+          cluster_names.each do |cluster_name|
+            invoke_task task_name, cluster_name
+          end
+        end
+
+        def invoke_task(name, cluster_name)
+          task_name = "active_record:sharding:#{name}"
+          to_rake_task(task_name).invoke cluster_name.to_s
+          to_rake_task(task_name).reenable
+        end
+
+        def invoke_task_for_all_sequencers(task_name)
+          sequencer_names.each do |sequencer_name|
+            invoke_task_for_sequencer task_name, sequencer_name
+          end
+        end
+
+        def invoke_task_for_sequencer(name, sequencer_name)
+          task_name = "active_record:sharding:sequencer:#{name}"
+          to_rake_task(task_name).invoke sequencer_name.to_s
+          to_rake_task(task_name).reenable
+        end
+      end
+      extend TasksForMultipleClusters
+
+      module TaskOrganizerForSingleClusterTask
+        def create_all_databases(args)
+          exec_task_for_all_databases "create", args
+        end
+
+        def drop_all_databases(args)
+          exec_task_for_all_databases "drop", args
+        end
+
+        def load_schema_all_databases(args)
+          exec_task_for_all_databases "load_schema", args
+        end
+
+        private
+
+        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
+
+        def cluster_name_or_error(name, args)
+          unless cluster_name = args[:cluster_name]
+            $stderr.puts <<-MSG
+Missing cluster_name. Find cluster_name via `rake active_record:sharding:info` then call `rake "active_record:sharding:#{name}[$cluster_name]"`.
+          MSG
+            exit
+          end
+          cluster_name
+        end
+
+        def cluster_or_error(cluster_name)
+          fetch_cluster_config cluster_name.to_sym
+        rescue KeyError
+          $stderr.puts %(cluster name "#{cluster_name}" not found.)
+          exit
+        end
+      end
+      extend TaskOrganizerForSingleClusterTask
+
+      module TasksForSingleConnection
+        def create(connection_name)
+          configuration = ActiveRecord::Base.configurations[connection_name]
+          ActiveRecord::Tasks::DatabaseTasks.create(configuration)
+          ActiveRecord::Base.establish_connection(configuration)
+        end
+
+        def drop(connection_name)
+          configuration = ActiveRecord::Base.configurations[connection_name]
+          ActiveRecord::Tasks::DatabaseTasks.drop configuration
+        end
+
+        def execute(connection_name, sql)
+          configuration = ActiveRecord::Base.configurations[connection_name]
+          ActiveRecord::Base.establish_connection(configuration).connection.execute sql
+        end
+
+        def load_schema(connection_name)
+          configuration = ActiveRecord::Base.configurations[connection_name]
+
+          case
+          when ar5?
+            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
+            fail "This version of ActiveRecord is not supported: v#{ActiveRecord::VERSION::STRING}"
+          end
+        end
+      end
+      extend TasksForSingleConnection
+
+      module TasksForSingleSequencerTask
+        def create_sequencer_database(args)
+          exec_task_for_sequencer_database "create", args
+        end
+
+        def drop_sequencer_database(args)
+          exec_task_for_sequencer_database "drop", args
+        end
+
+        def create_table_sequencer_database(args)
+          sequencer = sequencer_or_error "create_table", args
+          create_table_sql = "CREATE TABLE #{sequencer.table_name} (id BIGINT unsigned NOT NULL DEFAULT 0) ENGINE=MyISAM"
+          execute sequencer.connection_name.to_s, create_table_sql
+        end
+
+        def insert_initial_record_sequencer_database(args)
+          sequencer = sequencer_or_error "insert_initial_record", args
+          insert_initial_record_sql = "INSERT INTO #{sequencer.table_name} VALUES (0)"
+          execute sequencer.connection_name.to_s, insert_initial_record_sql
+        end
+
+        private
+
+        def exec_task_for_sequencer_database(task_name, args)
+          sequencer = sequencer_or_error task_name, args
+          __send__ task_name, sequencer.connection_name.to_s
+        end
+
+        def sequencer_or_error(task_name, args)
+          sequencer_name = sequencer_name_or_error task_name, args
+          fetch_sequencer_config sequencer_name.to_sym
+        rescue KeyError
+          $stderr.puts %(sequencer name "#{sequencer_name}" not found.)
+          exit
+        end
+
+        def sequencer_name_or_error(task_name, args)
+          unless sequencer_name = args[:sequencer_name]
+            # rubocop:disable Metrics/LineLength
+            $stderr.puts <<-MSG
+Missing sequencer_name. Find sequencer_name via `rake active_record:sharding:info` then call `rake "active_record:sharding:sequencer#{task_name}[$sequencer_name]"`.
+          MSG
+            exit
+            # rubocop:enable Metrics/LineLength
+          end
+          sequencer_name
+        end
+      end
+      extend TasksForSingleSequencerTask
+    end # module DatabaseTasks
+  end
+end

data/lib/active_record/sharding/errors.rb

@@ -0,0 +1,9 @@
+module ActiveRecord
+  module Sharding
+    class Error < ::StandardError
+    end
+
+    class MissingShardingKeyAttribute < Error
+    end
+  end
+end

data/lib/active_record/sharding/model.rb

@@ -0,0 +1,59 @@
+require "active_support/concern"
+
+module ActiveRecord
+  module Sharding
+    module Model
+      extend ActiveSupport::Concern
+
+      included do
+        class_attribute :cluster_router, instance_writer: false
+        class_attribute :shard_repository, instance_writer: false
+        class_attribute :sharding_key, instance_writer: false
+      end
+
+      module ClassMethods
+        def use_sharding(name, algorithm = :modulo)
+          config = ActiveRecord::Sharding.config.fetch_cluster_config name
+          if algorithm == :modulo
+            self.cluster_router = ActiveRecord::Sharding::ModuloRouter.new config
+          end
+          self.shard_repository = ActiveRecord::Sharding::ShardRepository.new config, self
+          self.abstract_class = true
+        end
+
+        def define_sharding_key(column)
+          self.sharding_key = column.to_sym
+        end
+
+        def before_put(&block)
+          @before_put_callback = block
+        end
+
+        def put!(attributes)
+          fail "`sharding_key` is not defined. Use `define_sharding_key`." unless sharding_key
+
+          @before_put_callback.call(attributes) if @before_put_callback
+
+          if key = attributes[sharding_key] || attributes[sharding_key.to_s]
+            shard_for(key).create!(attributes)
+          else
+            fail ActiveRecord::Sharding::MissingShardingKeyAttribute
+          end
+        end
+
+        def shard_for(key)
+          connection_name = cluster_router.route key
+          shard_repository.fetch connection_name
+        end
+
+        def all_shards
+          shard_repository.all
+        end
+
+        def define_parent_methods(&block)
+          instance_eval(&block)
+        end
+      end
+    end
+  end
+end

data/lib/active_record/sharding/modulo_router.rb

@@ -0,0 +1,14 @@
+module ActiveRecord
+  module Sharding
+    class ModuloRouter
+      def initialize(cluster_config)
+        @cluster_config = cluster_config
+      end
+
+      def route(id)
+        modulo_key = id % @cluster_config.registerd_connection_count
+        @cluster_config.fetch modulo_key
+      end
+    end
+  end
+end

data/lib/active_record/sharding/railtie.rb

@@ -0,0 +1,9 @@
+module ActiveRecord
+  module Sharding
+    class Railtie < ::Rails::Railtie
+      rake_tasks do
+        load File.expand_path("../../../tasks/activerecord-sharding.rake", __FILE__)
+      end
+    end
+  end
+end

data/lib/active_record/sharding/sequencer.rb

@@ -0,0 +1,40 @@
+require "active_support/concern"
+
+module ActiveRecord
+  module Sharding
+    module Sequencer
+      extend ActiveSupport::Concern
+
+      included do
+        class_attribute :sequencer_repository, instance_writer: false
+        class_attribute :sequencer_name, instance_writer: false
+        class_attribute :sequencer_config, instance_writer: false
+      end
+
+      module ClassMethods
+        def use_sequencer(name)
+          self.sequencer_name = name
+          self.sequencer_config = ActiveRecord::Sharding.config.fetch_sequencer_config name
+          self.sequencer_repository = ActiveRecord::Sharding::SequencerRepository.new sequencer_config, self
+          self.abstract_class = true
+        end
+
+        def current_sequence_id
+          execute_sql "id"
+        end
+
+        def next_sequence_id
+          execute_sql "id +1"
+        end
+
+        def execute_sql(last_insert_id_args)
+          connection = sequencer_repository.fetch(sequencer_name).connection
+          connection.execute "UPDATE `#{sequencer_config.table_name}` SET id = LAST_INSERT_ID(#{last_insert_id_args})"
+          res = connection.execute "SELECT LAST_INSERT_ID()"
+          new_id = res.first.first.to_i
+          new_id
+        end
+      end
+    end
+  end
+end

data/lib/active_record/sharding/sequencer_config.rb

@@ -0,0 +1,26 @@
+module ActiveRecord
+  module Sharding
+    class SequencerConfig
+      attr_reader :name, :table_name, :connection_name
+
+      def initialize(name)
+        @name = name
+        @table_name = nil
+        @connection_name = nil
+      end
+
+      def register_connection(connection_name)
+        @connection_name = connection_name
+      end
+
+      def register_table_name(table_name)
+        @table_name = table_name
+      end
+
+      def validate_config!
+        fail "Nothing connection. Please call register_connection" if @connection_name.blank?
+        fail "Nothing table_name. Please call register_table_name" if @table_name.blank?
+      end
+    end
+  end
+end