TNTP-2109: Add bucket_ref to map calls

Author: ita-sammann

crud/common/call.lua

@@ -186,10 +186,10 @@ function call.map(vshard_router, func_name, func_args, opts)
         request_timeout = opts.mode == 'read' and opts.request_timeout or nil,
     }
     while iter:has_next() do
-        local args, replicaset, replicaset_id = iter:get()
+        local args, replicaset, replicaset_id, bucket_ids = iter:get()
 
         local future, err = retry_call_with_master_discovery(replicaset, vshard_call_name,
-            func_name, args, call_opts, opts.mode, {}) -- TODO: provide bucket_ids
+            func_name, args, call_opts, opts.mode, bucket_ids)
 
         if err ~= nil then
             local result_info = {

crud/common/map_call_cases/base_iter.lua

@@ -72,7 +72,7 @@ function BaseIterator:get()
     local replicaset = self.next_replicaset
     self.next_index, self.next_replicaset = next(self.replicasets, self.next_index)
 
-    return self.func_args, replicaset, replicaset_id
+    return self.func_args, replicaset, replicaset_id, {}
 end
 
 return BaseIterator

crud/common/map_call_cases/batch_insert_iter.lua

@@ -76,10 +76,11 @@ function BatchInsertIterator:get()
         self.next_batch.tuples,
         self.opts,
     }
+    local bucket_ids = self.next_batch.bucket_ids
 
     self.next_index, self.next_batch = next(self.batches_by_replicasets, self.next_index)
 
-    return func_args, replicaset, replicaset_id
+    return func_args, replicaset, replicaset_id, bucket_ids
 end
 
 return BatchInsertIterator

crud/common/map_call_cases/batch_upsert_iter.lua

@@ -85,10 +85,11 @@ function BatchUpsertIterator:get()
         self.next_batch.operations,
         self.opts,
     }
+    local bucket_ids = self.next_batch.bucket_ids
 
     self.next_index, self.next_batch = next(self.batches_by_replicasets, self.next_index)
 
-    return func_args, replicaset, replicaset_id
+    return func_args, replicaset, replicaset_id, bucket_ids
 end
 
 return BatchUpsertIterator

crud/common/sharding/init.lua

@@ -324,8 +324,10 @@ function sharding.split_tuples_by_replicaset(vshard_router, tuples, space, opts)
         local record_by_replicaset = batches[replicaset_id] or {
             replicaset = replicaset,
             tuples = {},
+            bucket_ids = {},
         }
         table.insert(record_by_replicaset.tuples, tuple)
+        record_by_replicaset.bucket_ids[sharding_data.bucket_id] = true
 
         if opts.operations ~= nil then
             record_by_replicaset.operations = record_by_replicaset.operations or {}
@@ -335,6 +337,14 @@ function sharding.split_tuples_by_replicaset(vshard_router, tuples, space, opts)
         batches[replicaset_id] = record_by_replicaset
     end
 
+    for _, rbr in pairs(batches) do
+        local bucket_ids = {}
+        for bid, _ in pairs(rbr.bucket_ids) do
+            table.insert(bucket_ids, bid)
+        end
+        rbr.bucket_ids = bucket_ids
+    end
+
     return {
         batches = batches,
         sharding_func_hash = sharding_func_hash,