test/lib: add key_utils.hh

Contains methods to generate partition and clustering keys. In the case
of the former, one can specify the shard to generate keys for.
We currently have some methods to generate these but they are not
generic. Therefore the tests are littered by open-coded variants.
The methods introduced here are completely generic: they can generate
keys for any schema.

test/lib/key_utils.cc

@@ -0,0 +1,109 @@
+/*
+ * Copyright (C) 2023-present ScyllaDB
+ */
+
+/*
+ * SPDX-License-Identifier: AGPL-3.0-or-later
+ */
+
+#include "test/lib/key_utils.hh"
+#include "test/lib/random_schema.hh"
+#include "test/lib/random_utils.hh"
+#include "test/lib/test_utils.hh"
+
+namespace tests {
+
+namespace {
+
+template<typename RawKey, typename DecoratedKey, typename Comparator>
+std::vector<DecoratedKey> generate_keys(
+        size_t n,
+        schema_ptr s,
+        Comparator cmp,
+        const std::vector<data_type>& types,
+        std::function<std::optional<DecoratedKey>(const RawKey&)> decorate_fun,
+        bool allow_prefixes,
+        std::optional<key_size> size) {
+    auto keys = std::set<DecoratedKey, Comparator>(cmp);
+    const auto effective_size = size.value_or(tests::key_size{1, 128});
+
+    std::mt19937 engine(tests::random::get_int<uint32_t>());
+    std::uniform_int_distribution<size_t> component_count_dist(1, types.size());
+    tests::value_generator value_gen;
+
+    std::vector<data_value> components;
+    components.reserve(types.size());
+
+    while (keys.size() != n) {
+        components.clear();
+        auto component_count = allow_prefixes ? component_count_dist(engine) : types.size();
+        for (size_t i = 0; i < component_count; ++i) {
+            components.emplace_back(value_gen.generate_atomic_value(engine, *types.at(i), effective_size.min, effective_size.max));
+        }
+        auto raw_key = RawKey::from_deeply_exploded(*s, components);
+        // discard empty keys on the off chance that we generate one
+        if (raw_key.is_empty() || (types.size() == 1 && raw_key.begin(*s)->empty())) {
+            continue;
+        }
+        if constexpr (std::is_same_v<RawKey, DecoratedKey>) {
+            keys.emplace(std::move(raw_key));
+        } else if (auto decorated_key_opt = decorate_fun(raw_key); decorated_key_opt) {
+            keys.emplace(std::move(*decorated_key_opt));
+        }
+    }
+
+    return std::vector<DecoratedKey>(keys.begin(), keys.end());
+}
+
+}
+
+std::vector<dht::decorated_key> generate_partition_keys(size_t n, schema_ptr s, std::optional<shard_id> shard, std::optional<key_size> size) {
+    return generate_keys<partition_key, dht::decorated_key, dht::decorated_key::less_comparator>(
+            n,
+            s,
+            dht::decorated_key::less_comparator(s),
+            s->partition_key_type()->types(),
+            [s, shard, tokens = std::set<dht::token>()] (const partition_key& pkey) mutable -> std::optional<dht::decorated_key> {
+                auto dkey = dht::decorate_key(*s, pkey);
+                if (shard && *shard != dht::shard_of(*s, dkey.token())) {
+                    return {};
+                }
+                if (!tokens.insert(dkey.token()).second) {
+                    return {};
+                }
+                return dkey;
+            },
+            false,
+            size);
+}
+
+std::vector<dht::decorated_key> generate_partition_keys(size_t n, schema_ptr s, local_shard_only lso, std::optional<key_size> size) {
+    return generate_partition_keys(n, std::move(s), lso == local_shard_only::yes ? std::optional(this_shard_id()) : std::nullopt, size);
+}
+
+dht::decorated_key generate_partition_key(schema_ptr s, std::optional<shard_id> shard, std::optional<key_size> size) {
+    auto&& keys = generate_partition_keys(1, std::move(s), shard, size);
+    return std::move(keys.front());
+}
+
+dht::decorated_key generate_partition_key(schema_ptr s, local_shard_only lso, std::optional<key_size> size) {
+    return generate_partition_key(std::move(s), lso == local_shard_only::yes ? std::optional(this_shard_id()) : std::nullopt, size);
+}
+
+std::vector<clustering_key> generate_clustering_keys(size_t n, schema_ptr s, bool allow_prefixes, std::optional<key_size> size) {
+    return generate_keys<clustering_key, clustering_key, clustering_key::less_compare>(
+            n,
+            s,
+            clustering_key::less_compare(*s),
+            s->clustering_key_type()->types(),
+            {},
+            allow_prefixes,
+            size);
+}
+
+clustering_key generate_clustering_key(schema_ptr s, bool allow_prefix, std::optional<key_size> size) {
+    auto&& keys = generate_clustering_keys(1, std::move(s), allow_prefix, size);
+    return std::move(keys.front());
+}
+
+} // namespace tests

test/lib/key_utils.hh

@@ -0,0 +1,48 @@
+/*
+ * Copyright (C) 2023-present ScyllaDB
+ */
+
+/*
+ * SPDX-License-Identifier: AGPL-3.0-or-later
+ */
+
+#pragma once
+
+#include "dht/i_partitioner.hh"
+
+struct local_shard_only_tag { };
+using local_shard_only = bool_class<local_shard_only_tag>;
+
+namespace tests {
+
+struct key_size {
+    size_t min;
+    size_t max;
+};
+
+// Generate n partition keys for the given schema.
+//
+// Returned keys are unique (their token too), ordered and never empty.
+// Parameters:
+// * n - number of keys
+// * s - schema of the keys, used also to obtain the sharder
+// * shard - only generate keys for this shard (if engaged)
+// * size - the min and max size of the key in bytes, if disengaged default
+//          limits (1-128) are used. If you want exactly sized keys, use
+//          ascii or bytes types only as the key types.
+std::vector<dht::decorated_key> generate_partition_keys(size_t n, schema_ptr s, std::optional<shard_id> shard = this_shard_id(), std::optional<key_size> size = {});
+std::vector<dht::decorated_key> generate_partition_keys(size_t n, schema_ptr s, local_shard_only lso, std::optional<key_size> size = {});
+// Overload for a single key
+dht::decorated_key generate_partition_key(schema_ptr s, std::optional<shard_id> shard = this_shard_id(), std::optional<key_size> size = {});
+dht::decorated_key generate_partition_key(schema_ptr s, local_shard_only lso, std::optional<key_size> size = {});
+
+// Generate n clustering keys
+//
+// Returned keys are unique, ordered and never empty.
+// Parameters are the same as that of generate_partition_keys().
+// If allow_prefixes is true, prefix keys may be generated too.
+std::vector<clustering_key> generate_clustering_keys(size_t n, schema_ptr s, bool allow_prefixes = false, std::optional<key_size> size = {});
+// Overload for a single key
+clustering_key generate_clustering_key(schema_ptr s, bool allow_prefix = false, std::optional<key_size> size = {});
+
+} // namespace tests

test/lib/simple_schema.hh

@@ -20,11 +20,9 @@
 #include "schema_builder.hh"
 #include "reader_permit.hh"
 #include "types/map.hh"
+#include "test/lib/key_utils.hh"
 #include "atomic_cell_or_collection.hh"
 
-struct local_shard_only_tag { };
-using local_shard_only = bool_class<local_shard_only_tag>;
-
 //
 // Make set of keys sorted by token for current or remote shard.
 //