lib/tako/query_chain.rb in tako-0.2.2 vs lib/tako/query_chain.rb in tako-0.3.0

@@ -7,24 +7,40 @@
       @proxy = proxy
       @base_object = base_object
     end
 
     def method_missing(method, *args)
-      @proxy.in_proxy do
-        if block_given?
-          base_object.send(method, *args) do
-            yield
-          end
-        else
-          base_object.send(method, *args)
+      @proxy.with_shard do
+        result = if block_given?
+                    base_object.send(method, *args) do
+                      yield
+                    end
+                  else
+                    base_object.send(method, *args)
+                  end
+
+        if chain_available?(result)
+          @base_object = result
+          return self
         end
+
+        result
       end
     end
 
     def shard(shard_name)
       new(
-        Tako::Repository.shard(shard_name),
+        Tako::Repository.create_proxy(shard_name),
         self
       )
+    end
+
+    private
+
+    def chain_available?(obj)
+      [
+        ::ActiveRecord::Relation,
+        ::ActiveRecord::QueryMethods::WhereChain
+      ].any? { |anc| obj.is_a?(anc) }
     end
   end
 end