lib/tako/query_chain.rb in tako-0.1.0 vs lib/tako/query_chain.rb in tako-0.2.0

@@ -8,10 +8,16 @@
       @base_object = base_object
     end
 
     def method_missing(method, *args)
       @proxy.in_proxy do
-        base_object.send(method, *args)
+        if block_given?
+          base_object.send(method, *args) do
+            yield
+          end
+        else
+          base_object.send(method, *args)
+        end
       end
     end
 
     def shard(shard_name)
       new(