From caf62f0d35bbf690851aadfbf390c559bf8a1377 Mon Sep 17 00:00:00 2001
From: yuemori <yuemori@aiming-inc.com>
Date: Mon, 1 Aug 2016 03:00:25 +0900
Subject: [PATCH] Implement distkey support to instance method

---
 lib/activerecord/shard_for/model.rb       | 20 ++++++++++++++---
 spec/activerecord/shard_for/model_spec.rb | 26 +++++++++++++++++++++++
 2 files changed, 43 insertions(+), 3 deletions(-)

diff --git a/lib/activerecord/shard_for/model.rb b/lib/activerecord/shard_for/model.rb
index f12be2a..bc1fc80 100644
--- a/lib/activerecord/shard_for/model.rb
+++ b/lib/activerecord/shard_for/model.rb
@@ -50,11 +50,11 @@ def shard_for(key)
         # @raise [ActiveRecord::ShardFor::MissingDistkeyAttribute]
         def put!(attributes)
           raise '`distkey` is not defined. Use `def_distkey`.' unless distkey
-
           @before_put_callback.call(attributes) if defined?(@before_put_callback) && @before_put_callback
-          key = attributes[distkey]
 
-          raise ActiveRecord::ShardFor::MissingDistkeyAttribute unless key || attributes[distkey.to_s]
+          key = fetch_distkey_from_attributes(attributes)
+
+          raise ActiveRecord::ShardFor::MissingDistkeyAttribute unless key
 
           shard_for(key).create!(attributes)
         end
@@ -139,6 +139,20 @@ def replicates_with(mapping)
         def switch(role_name, &block)
           replication_mapping.switch(self, role_name, &block)
         end
+
+        private
+
+        # @param [Hash] attributes
+        # @return [Object or nil] distkey
+        def fetch_distkey_from_attributes(attributes)
+          key = attributes[distkey] || attributes[distkey.to_s]
+          return key if key
+
+          instance = all_shards.first.new(attributes)
+          return unless instance.respond_to?(distkey)
+
+          instance.send(distkey)
+        end
       end
     end
   end
diff --git a/spec/activerecord/shard_for/model_spec.rb b/spec/activerecord/shard_for/model_spec.rb
index c1ae0e6..452075d 100644
--- a/spec/activerecord/shard_for/model_spec.rb
+++ b/spec/activerecord/shard_for/model_spec.rb
@@ -21,6 +21,22 @@ def self.generate_name
     end
   end
 
+  let!(:another_model) do
+    Class.new(ActiveRecord::Base) do
+      def self.name
+        'User'
+      end
+
+      include ActiveRecord::ShardFor::Model
+      use_cluster :character, :distkey
+      def_distkey :shard_no
+
+      def shard_no
+        1
+      end
+    end
+  end
+
   let(:user_attributes) { { name: 'Alice', email: 'alice@example.com' } }
 
   describe '.put!' do
@@ -39,6 +55,16 @@ def self.generate_name
           .to raise_error(ActiveRecord::ShardFor::MissingDistkeyAttribute)
       end
     end
+
+    context 'distkey be specify instance method' do
+      before { another_model.put!(user_attributes) }
+
+      it 'creates new record into shard 1' do
+        record = another_model.using(1).find_by(email: user_attributes[:email])
+        expect(record).to be_a(another_model)
+        expect(record).not_to be_nil
+      end
+    end
   end
 
   describe '.before_put' do