MONGOID-5582 Correctly interpret shard keys in embedded documents

lib/mongoid/shardable.rb

@@ -53,12 +53,10 @@ def shard_key_fields
     # @return [ Hash ] The shard key selector.
     #
     # @api private
-    def shard_key_selector
-      selector = {}
-      shard_key_fields.each do |field|
-        selector[field.to_s] = send(field)
+    def shard_key_selector(prefer_persisted: false)
+      shard_key_fields.each_with_object({}) do |field, selector|
+        selector[field.to_s] = shard_key_field_value(field.to_s, prefer_persisted: prefer_persisted)
       end
-      selector
     end
 
     # Returns the selector that would match the existing version of this
@@ -72,11 +70,29 @@ def shard_key_selector
     #
     # @api private
     def shard_key_selector_in_db
-      selector = {}
-      shard_key_fields.each do |field|
-        selector[field.to_s] = new_record? ? send(field) : attribute_was(field)
+      shard_key_selector(prefer_persisted: true)
+    end
+
+    # Returns the value for the named shard key. If the field identifies
+    # an embedded document, the key will be parsed and recursively evaluated.
+    # If `prefer_persisted` is true, the value last persisted to the database
+    # will be returned, regardless of what the current value of the attribute
+    # may be.
+    #
+    # @param [String] field The name of the field to evaluate
+    # @param [ true|false ] prefer_persisted Whether or not to prefer the
+    #   persisted value over the current value.
+    #
+    # @return [ Object ] The value of the named field.
+    def shard_key_field_value(field, prefer_persisted:)
+      if field.include?(".")
+        relation, remaining = field.split(".", 2)
+        send(relation)&.shard_key_field_value(remaining, prefer_persisted: prefer_persisted)
+      elsif prefer_persisted && !new_record?
+        attribute_was(field)
+      else
+        send(field)
       end
-      selector
     end
 
     module ClassMethods

spec/mongoid/shardable_models.rb

@@ -59,3 +59,17 @@ class SmDriver
 class SmNotSharded
   include Mongoid::Document
 end
+
+class SmReviewAuthor
+  include Mongoid::Document
+  embedded_in :review, class_name: "SmReview", touch: false
+  field :name, type: String
+end
+
+class SmReview
+  include Mongoid::Document
+
+  embeds_one :author, class_name: "SmReviewAuthor"
+
+  shard_key "author.name" => 1
+end

spec/mongoid/shardable_spec.rb

@@ -27,11 +27,11 @@
     context 'when full syntax is used' do
       context 'with symbol value' do
         it 'sets shard key fields to symbol value' do
-          SmProducer.shard_key_fields.should == %i(age gender)
+          expect(SmProducer.shard_key_fields).to be == %i(age gender)
         end
 
         it 'sets shard config' do
-          SmProducer.shard_config.should == {
+          expect(SmProducer.shard_config).to be == {
             key: {age: 1, gender: 'hashed'},
             options: {
               unique: true,
@@ -41,163 +41,258 @@
         end
 
         it 'keeps hashed as string' do
-          SmProducer.shard_config[:key][:gender].should == 'hashed'
+          expect(SmProducer.shard_config[:key][:gender]).to be == 'hashed'
         end
       end
 
       context 'with string value' do
         it 'sets shard key fields to symbol value' do
-          SmActor.shard_key_fields.should == %i(age gender hello)
+          expect(SmActor.shard_key_fields).to be == %i(age gender hello)
         end
 
         it 'sets shard config' do
-          SmActor.shard_config.should == {
+          expect(SmActor.shard_config).to be == {
             key: {age: 1, gender: 'hashed', hello: 'hashed'},
             options: {},
           }
         end
 
         it 'sets hashed to string' do
-          SmActor.shard_config[:key][:gender].should == 'hashed'
+          expect(SmActor.shard_config[:key][:gender]).to be == 'hashed'
         end
       end
 
       context 'when passed association name' do
         it 'uses foreign key as shard key in shard config' do
-          SmDriver.shard_config.should == {
+          expect(SmDriver.shard_config).to be == {
             key: {age: 1, agency_id: 'hashed'},
             options: {},
           }
         end
 
         it 'uses foreign key as shard key in shard key fields' do
-          SmDriver.shard_key_fields.should == %i(age agency_id)
+          expect(SmDriver.shard_key_fields).to be == %i(age agency_id)
         end
       end
+
+      context 'when passed subdocument fields' do
+      end
     end
 
     context 'when shorthand syntax is used' do
       context 'with symbol value' do
         it 'sets shard key fields to symbol value' do
-          SmMovie.shard_key_fields.should == %i(year)
+          expect(SmMovie.shard_key_fields).to be == %i(year)
         end
       end
 
       context 'with string value' do
         it 'sets shard key fields to symbol value' do
-          SmTrailer.shard_key_fields.should == %i(year)
+          expect(SmTrailer.shard_key_fields).to be == %i(year)
         end
       end
 
       context 'when passed association name' do
         it 'uses foreign key as shard key in shard config' do
-          SmDirector.shard_config.should == {
+          expect(SmDirector.shard_config).to be == {
             key: {agency_id: 1},
             options: {},
           }
         end
 
         it 'uses foreign key as shard key in shard key fields' do
-          SmDirector.shard_key_fields.should == %i(agency_id)
+          expect(SmDirector.shard_key_fields).to be == %i(agency_id)
         end
       end
     end
   end
 
   describe '#shard_key_selector' do
     subject { instance.shard_key_selector }
-    let(:klass) { Band }
-    let(:value) { 'a-brand-name' }
+
+    context 'when key is an immediate attribute' do
+      let(:klass) { Band }
+      let(:value) { 'a-brand-name' }
+
+      before { klass.shard_key(:name) }
 
-    before { klass.shard_key(:name) }
+      context 'when record is new' do
+        let(:instance) { klass.new(name: value) }
 
-    context 'when record is new' do
-      let(:instance) { klass.new(name: value) }
+        it { is_expected.to eq({ 'name' => value }) }
 
-      it { is_expected.to eq({ 'name' => value }) }
+        context 'changing shard key value' do
+          let(:new_value) { 'a-new-value' }
 
-      context 'changing shard key value' do
-        let(:new_value) { 'a-new-value' }
+          before do
+            instance.name = new_value
+          end
 
-        before do
-          instance.name = new_value
+          it { is_expected.to eq({ 'name' => new_value }) }
         end
+      end
 
-        it { is_expected.to eq({ 'name' => new_value }) }
+      context 'when record is persisted' do
+        let(:instance) { klass.create!(name: value) }
+
+        it { is_expected.to eq({ 'name' => value }) }
+
+        context 'changing shard key value' do
+          let(:new_value) { 'a-new-value' }
+
+          before do
+            instance.name = new_value
+          end
+
+          it { is_expected.to eq({ 'name' => new_value }) }
+        end
       end
     end
 
-    context 'when record is persisted' do
-      let(:instance) { klass.create!(name: value) }
+    context 'when key is an embedded attribute' do
+      let(:klass) { SmReview }
+      let(:value) { 'Arthur Conan Doyle' }
+      let(:key)   { 'author.name' }
 
-      it { is_expected.to eq({ 'name' => value }) }
+      context 'when record is new' do
+        let(:instance) { klass.new(author: { name: value }) }
 
-      context 'changing shard key value' do
-        let(:new_value) { 'a-new-value' }
+        it { is_expected.to eq({ key => value }) }
 
-        before do
-          instance.name = new_value
+        context 'changing shard key value' do
+          let(:new_value) { 'Jules Verne' }
+
+          before do
+            instance.author.name = new_value
+          end
+
+          it { is_expected.to eq({ key => new_value }) }
         end
+      end
+
+      context 'when record is persisted' do
+        let(:instance) { klass.create!(author: { name: value }) }
+
+        it { is_expected.to eq({ key => value }) }
 
-        it 'uses the newly set shard key value' do
-          subject.should == { 'name' => new_value }
+        context 'changing shard key value' do
+          let(:new_value) { 'Jules Verne' }
+
+          before do
+            instance.author.name = new_value
+          end
+
+          it { is_expected.to eq({ 'author.name' => new_value }) }
         end
       end
     end
   end
 
   describe '#shard_key_selector_in_db' do
     subject { instance.shard_key_selector_in_db }
-    let(:klass) { Band }
-    let(:value) { 'a-brand-name' }
 
-    before { klass.shard_key(:name) }
+    context 'when key is an immediate attribute' do
+      let(:klass) { Band }
+      let(:value) { 'a-brand-name' }
 
-    context 'when record is new' do
-      let(:instance) { klass.new(name: value) }
+      before { klass.shard_key(:name) }
 
-      it { is_expected.to eq({ 'name' => value }) }
+      context 'when record is new' do
+        let(:instance) { klass.new(name: value) }
 
-      context 'changing shard key value' do
-        let(:new_value) { 'a-new-value' }
+        it { is_expected.to eq({ 'name' => value }) }
 
-        before do
-          instance.name = new_value
-        end
+        context 'changing shard key value' do
+          let(:new_value) { 'a-new-value' }
 
-        it 'uses the existing shard key value' do
-          subject.should == { 'name' => new_value }
+          before do
+            instance.name = new_value
+          end
+
+          it { is_expected.to eq({ 'name' => new_value }) }
         end
       end
-    end
 
-    context 'when record is persisted' do
-      let(:instance) { klass.create!(name: value) }
+      context 'when record is persisted' do
+        let(:instance) { klass.create!(name: value) }
+
+        it { is_expected.to eq({ 'name' => value }) }
 
-      it { is_expected.to eq({ 'name' => value }) }
+        context 'changing shard key value' do
+          let(:new_value) { 'a-new-value' }
 
-      context 'changing shard key value' do
-        let(:new_value) { 'a-new-value' }
+          before do
+            instance.name = new_value
+          end
+
+          it { is_expected.to eq({ 'name' => value }) }
+        end
+      end
+
+      context "when record is not found" do
+        let!(:instance) { klass.create!(name: value) }
 
         before do
-          instance.name = new_value
+          instance.destroy
         end
 
-        it { is_expected.to eq({ 'name' => value }) }
+        it "raises a DocumentNotFound error with the shard key in the description on reload" do
+          expect do
+            instance.reload
+          end.to raise_error(Mongoid::Errors::DocumentNotFound, /Document not found for class Band with id #{instance.id.to_s} and shard key name: a-brand-name./)
+        end
       end
     end
 
-    context "when record is not found" do
-      let!(:instance) { klass.create!(name: value) }
+    context 'when key is an embedded attribute' do
+      let(:klass) { SmReview }
+      let(:value) { 'Arthur Conan Doyle' }
+      let(:key)   { 'author.name' }
+
+      context 'when record is new' do
+        let(:instance) { klass.new(author: { name: value }) }
 
-      before do
-        instance.destroy
+        it { is_expected.to eq({ key => value }) }
+
+        context 'changing shard key value' do
+          let(:new_value) { 'Jules Verne' }
+
+          before do
+            instance.author.name = new_value
+          end
+
+          it { is_expected.to eq({ key => new_value }) }
+        end
       end
 
-      it "raises a DocumentNotFound error with the shard key in the description on reload" do
-        expect do
-          instance.reload
-        end.to raise_error(Mongoid::Errors::DocumentNotFound, /Document not found for class Band with id #{instance.id.to_s} and shard key name: a-brand-name./)
+      context 'when record is persisted' do
+        let(:instance) { klass.create!(author: { name: value }) }
+
+        it { is_expected.to eq({ key => value }) }
+
+        context 'changing shard key value' do
+          let(:new_value) { 'Jules Verne' }
+
+          before do
+            instance.author.name = new_value
+          end
+
+          it { is_expected.to eq({ key => value }) }
+        end
+
+        context "when record is not found" do
+          let!(:instance) { klass.create!(author: { name: value }) }
+
+          before do
+            instance.destroy
+          end
+
+          it "raises a DocumentNotFound error with the shard key in the description on reload" do
+            expect do
+              instance.reload
+            end.to raise_error(Mongoid::Errors::DocumentNotFound, /Document not found for class SmReview with id #{instance.id.to_s} and shard key author.name: Arthur Conan Doyle./)
+          end
+        end
       end
     end
   end