Support custom sharding key in delete, update and get

Part of #166

Author: ligurio

crud/common/sharding_key.lua

@@ -7,6 +7,7 @@ local dev_checks = require('crud.common.dev_checks')
 local cache = require('crud.common.sharding_key_cache')
 local utils = require('crud.common.utils')
 
+local ShardingKeyError = errors.new_class("ShardingKeyError", {capture_stack = false})
 local FetchShardingKeyError = errors.new_class('FetchShardingKeyError', {capture_stack = false})
 local GetShardingKeyFieldnosError = errors.new_class('GetShardingKeyFieldnosError',  {capture_stack = false})
 
@@ -115,6 +116,87 @@ function sharding_key_module.fetch_on_router(space_name)
     return cache.sharding_key_def
 end
 
+-- Make sure sharding key definition is a part of primary key.
+function sharding_key_module.is_part_of_pk(space_name, space_format,
+                                           primary_index, sharding_key_def)
+    dev_checks('string', 'table', 'table', 'table')
+
+    if cache.as_index_object[space_name] == nil then
+        cache.as_index_object[space_name] =
+            sharding_key_module.as_index_object(space_format, space_name,
+                                                sharding_key_def)
+    end
+    local sharding_key_as_index_obj = cache.as_index_object[space_name]
+
+    if cache.is_part_of_pk[space_name] ~= nil then
+        return cache.is_part_of_pk[space_name]
+    end
+
+    local is_part_of_pk = true
+    local pk_fieldno_map = utils.get_index_fieldno_map(primary_index.parts)
+    for _, part in ipairs(sharding_key_as_index_obj.parts) do
+        if pk_fieldno_map[part.fieldno] == nil then
+            is_part_of_pk = false
+            break
+        end
+    end
+    cache.is_part_of_pk[space_name] = is_part_of_pk
+
+    return is_part_of_pk
+end
+
+-- Build an array with sharding key values.
+function sharding_key_module.extract_from_index(sharding_key_def,
+                                                index_parts,
+                                                sharding_key_fieldno_map)
+    dev_checks('table', 'table', 'table')
+
+    local sharding_key = {}
+    for i, k in ipairs(sharding_key_def) do
+        local fieldno = index_parts[i].fieldno
+        if sharding_key_fieldno_map[fieldno] == true then
+            table.insert(sharding_key, k)
+        end
+    end
+
+    return sharding_key
+end
+
+-- Extract sharding key from pk.
+-- Returns a table with sharding key or pair of nil and error.
+function sharding_key_module.extract_from_pk(space_obj, sharding_key)
+    dev_checks('table', '?')
+
+    local space_name = space_obj.name
+    local sharding_key_def, err = sharding_key_module.fetch_on_router(space_name)
+    if err ~= nil then
+        return nil, err
+    end
+    if sharding_key_def == nil then
+        return sharding_key
+    end
+
+    local primary_index = space_obj.index[0]
+    local space_format = space_obj:format()
+
+    local is_part_of_pk = sharding_key_module.is_part_of_pk(space_name, space_format,
+                                                            primary_index, sharding_key_def)
+    if is_part_of_pk == false then
+        return nil, ShardingKeyError:new(
+            "Sharding key for space %q is missed in primary index, specify bucket_id",
+            space_name
+        )
+    end
+    if type(sharding_key) ~= 'table' then
+        sharding_key = {sharding_key}
+    end
+    local sharding_key_as_index_obj = cache.as_index_object[space_name]
+    local sharding_key_fieldno_map = utils.get_index_fieldno_map(sharding_key_as_index_obj.parts)
+
+    return sharding_key_module.extract_from_index(sharding_key, primary_index.parts,
+                                                  sharding_key_fieldno_map)
+end
+
 function sharding_key_module.init()
    _G._crud.fetch_on_storage = sharding_key_module.fetch_on_storage
 end

crud/common/sharding_key_cache.lua

@@ -1,5 +1,13 @@
 local sharding_key_cache = {}
 
 sharding_key_cache.sharding_key_def = nil
+sharding_key_cache.is_part_of_pk = {}
+sharding_key_cache.as_index_object = {}
+
+function sharding_key_cache.drop_caches()
+    sharding_key_cache.sharding_key_def = nil
+    sharding_key_cache.is_part_of_pk = {}
+    sharding_key_cache.as_index_object = {}
+end
 
 return sharding_key_cache

crud/common/utils.lua

@@ -375,6 +375,18 @@ function utils.get_bucket_id_fieldno(space, shard_index_name)
     return bucket_id_index.parts[1].fieldno
 end
 
+function utils.get_index_fieldno_map(index_parts)
+    dev_checks('table')
+
+    local fieldno_map = {}
+    for _, part in ipairs(index_parts) do
+        local fieldno = part.fieldno
+        fieldno_map[fieldno] = true
+    end
+
+    return fieldno_map
+end
+
 -- Build a map with field names as a keys and fieldno's
 -- as a values using space format as a source.
 function utils.get_format_fieldno_map(space_format)

crud/delete.lua

@@ -5,6 +5,7 @@ local vshard = require('vshard')
 local call = require('crud.common.call')
 local utils = require('crud.common.utils')
 local sharding = require('crud.common.sharding')
+local sharding_key_module = require('crud.common.sharding_key')
 local dev_checks = require('crud.common.dev_checks')
 local schema = require('crud.common.schema')
 
@@ -55,7 +56,16 @@ local function call_delete_on_router(space_name, key, opts)
         key = key:totable()
     end
 
-    local bucket_id = sharding.key_get_bucket_id(key, opts.bucket_id)
+    local sharding_key = key
+    if opts.bucket_id == nil then
+        local err
+        sharding_key, err = sharding_key_module.extract_from_pk(space, key)
+        if err ~= nil then
+            return nil, err
+        end
+    end
+
+    local bucket_id = sharding.key_get_bucket_id(sharding_key, opts.bucket_id)
     local call_opts = {
         mode = 'write',
         timeout = opts.timeout,

crud/get.lua

@@ -5,6 +5,7 @@ local vshard = require('vshard')
 local call = require('crud.common.call')
 local utils = require('crud.common.utils')
 local sharding = require('crud.common.sharding')
+local sharding_key_module = require('crud.common.sharding_key')
 local dev_checks = require('crud.common.dev_checks')
 local schema = require('crud.common.schema')
 
@@ -58,7 +59,16 @@ local function call_get_on_router(space_name, key, opts)
         key = key:totable()
     end
 
-    local bucket_id = sharding.key_get_bucket_id(key, opts.bucket_id)
+    local sharding_key = key
+    if opts.bucket_id == nil then
+        local err
+        sharding_key, err = sharding_key_module.extract_from_pk(space, key)
+        if err ~= nil then
+            return nil, err
+        end
+    end
+
+    local bucket_id = sharding.key_get_bucket_id(sharding_key, opts.bucket_id)
     local call_opts = {
         mode = opts.mode or 'read',
         prefer_replica = opts.prefer_replica,

crud/update.lua

@@ -5,6 +5,7 @@ local vshard = require('vshard')
 local call = require('crud.common.call')
 local utils = require('crud.common.utils')
 local sharding = require('crud.common.sharding')
+local sharding_key_module = require('crud.common.sharding_key')
 local dev_checks = require('crud.common.dev_checks')
 local schema = require('crud.common.schema')
 
@@ -83,6 +84,15 @@ local function call_update_on_router(space_name, key, user_operations, opts)
         key = key:totable()
     end
 
+    local sharding_key = key
+    if opts.bucket_id == nil then
+        local err
+        sharding_key, err = sharding_key_module.extract_from_pk(space, key)
+        if err ~= nil then
+            return nil, err
+        end
+    end
+
     local operations = user_operations
     if not utils.tarantool_supports_fieldpaths() then
         operations, err = utils.convert_operations(user_operations, space_format)
@@ -91,7 +101,7 @@ local function call_update_on_router(space_name, key, user_operations, opts)
         end
     end
 
-    local bucket_id = sharding.key_get_bucket_id(key, opts.bucket_id)
+    local bucket_id = sharding.key_get_bucket_id(sharding_key, opts.bucket_id)
     local call_opts = {
         mode = 'write',
         timeout = opts.timeout,

test/entrypoint/srv_ddl.lua

@@ -45,6 +45,8 @@ package.preload['customers-storage'] = function()
                 sharding_key = {'name'},
             }
 
+            local customers_age_key_schema = table.deepcopy(customers_name_key_schema)
+            customers_age_key_schema.sharding_key = {'age'}
             local name_index = {
                 name = 'name',
                 type = 'TREE',
@@ -91,6 +93,7 @@ package.preload['customers-storage'] = function()
                     customers_name_key = customers_name_key_schema,
                     customers_name_key_uniq_index = customers_name_key_uniq_index_schema,
                     customers_name_key_non_uniq_index = customers_name_key_non_uniq_index_schema,
+                    customers_age_key = customers_age_key_schema,
                     customers_secondary_idx_name_key = customers_secondary_idx_name_key_schema,
                 }
             }
@@ -101,6 +104,11 @@ package.preload['customers-storage'] = function()
                     error(err)
                 end
             end
+
+            rawset(_G, 'set_sharding_key', function(space_name, sharding_key_def)
+                local fieldno_sharding_key = 2
+                box.space['_ddl_sharding_key']:update(space_name, {{'=', fieldno_sharding_key, sharding_key_def}})
+            end)
         end,
     }
 end