// Licensed to the Apache Software Foundation (ASF) under one

// or more contributor license agreements. See the NOTICE file

// distributed with this work for additional information

// regarding copyright ownership. The ASF licenses this file

// to you under the Apache License, Version 2.0 (the

// "License"); you may not use this file except in compliance

// with the License. You may obtain a copy of the License at

//

// http://www.apache.org/licenses/LICENSE-2.0

//

// Unless required by applicable law or agreed to in writing,

// software distributed under the License is distributed on an

// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY

// KIND, either express or implied. See the License for the

// specific language governing permissions and limitations

// under the License.

#include "meta-store/document_message.h"

#include <bthread/bthread.h>

#include <fmt/format.h>

#include <foundationdb/fdb_c_options.g.h>

#include <gen_cpp/cloud.pb.h>

#include <gen_cpp/olap_common.pb.h>

#include <gen_cpp/olap_file.pb.h>

#include <gtest/gtest.h>

#include <iostream>

#include <memory>

#include <mutex>

#include <string>

#include <string_view>

#include "common/config.h"

#include "common/util.h"

#include "meta-store/document_message_get_range.h"

#include "meta-store/keys.h"

#include "meta-store/mem_txn_kv.h"

#include "meta-store/txn_kv.h"

#include "meta-store/txn_kv_error.h"

#include "meta-store/versionstamp.h"

using namespace doris::cloud;

static std::mutex config_mutex;

int main(int argc, char** argv) {

config::init(nullptr, true);

::testing::InitGoogleTest(&argc, argv);

return RUN_ALL_TESTS();

}

// Convert a string to a hex-escaped string.

// A non-displayed character is represented as \xHH where HH is the hexadecimal value of the character.

// A displayed character is represented as itself.

static std::string escape_hex(std::string_view data) {

std::string result;

for (char c : data) {

if (isprint(c)) {

result += c;

} else {

result += fmt::format("\\x{:02x}", static_cast<unsigned char>(c));

}

}

return result;

}

static size_t count_range(TxnKv* txn_kv, std::string_view begin = "",

std::string_view end = "\xFF") {

std::unique_ptr<Transaction> txn;

EXPECT_EQ(txn_kv->create_txn(&txn), TxnErrorCode::TXN_OK);

if (!txn) {

return 0; // Failed to create transaction

}

FullRangeGetOptions opts;

opts.txn = txn.get();

auto iter = txn_kv->full_range_get(std::string(begin), std::string(end), std::move(opts));

size_t total = 0;

for (auto&& kvp = iter->next(); kvp.has_value(); kvp = iter->next()) {

total += 1;

}

EXPECT_TRUE(iter->is_valid()); // The iterator should still be valid after the next call.

return total;

}

static bool is_empty_range(TxnKv* txn_kv, std::string_view begin = "",

std::string_view end = "\xFF") {

return count_range(txn_kv, begin, end) == 0;

}

static std::string dump_range(TxnKv* txn_kv, std::string_view begin = "",

std::string_view end = "\xFF") {

std::unique_ptr<Transaction> txn;

if (txn_kv->create_txn(&txn) != TxnErrorCode::TXN_OK) {

return "Failed to create dump range transaction";

}

FullRangeGetOptions opts;

opts.txn = txn.get();

auto iter = txn_kv->full_range_get(std::string(begin), std::string(end), std::move(opts));

std::string buffer;

for (auto&& kv = iter->next(); kv.has_value(); kv = iter->next()) {

buffer +=

fmt::format("Key: {}, Value: {}\n", escape_hex(kv->first), escape_hex(kv->second));

}

EXPECT_TRUE(iter->is_valid()); // The iterator should still be valid after the next call.

return buffer;

}

TEST(DocumentMessageTest, SplitSingleMessage) {

auto txn_kv = std::make_shared<MemTxnKv>();

{

// Write a message that does not need to be split.

std::unique_ptr<Transaction> txn;

ASSERT_EQ(txn_kv->create_txn(&txn), TxnErrorCode::TXN_OK);

SplitSingleMessagePB msg;

msg.set_other_fields(1);

ASSERT_TRUE(document_put(txn.get(), "split_single_message_key", std::move(msg)));

ASSERT_EQ(txn->commit(), TxnErrorCode::TXN_OK);

// Read the message back.

std::unique_ptr<Transaction> read_txn;

ASSERT_EQ(txn_kv->create_txn(&read_txn), TxnErrorCode::TXN_OK);

SplitSingleMessagePB read_msg;

ASSERT_EQ(document_get(read_txn.get(), "split_single_message_key", &read_msg),

TxnErrorCode::TXN_OK);

ASSERT_EQ(read_msg.other_fields(), 1);

ASSERT_EQ(count_range(txn_kv.get()), 1) << dump_range(txn_kv.get());

// Remove the message.

document_remove<SplitSingleMessagePB>(read_txn.get(), "split_single_message_key");

ASSERT_EQ(read_txn->commit(), TxnErrorCode::TXN_OK);

}

ASSERT_TRUE(is_empty_range(txn_kv.get())) << dump_range(txn_kv.get());

{

// Write a message that needs to be split.

std::unique_ptr<Transaction> txn;

ASSERT_EQ(txn_kv->create_txn(&txn), TxnErrorCode::TXN_OK);

SplitSingleMessagePB msg;

msg.set_other_fields(2);

auto* key_bounds = msg.mutable_segment_key_bounds();

key_bounds->set_min_key("min_key");

key_bounds->set_max_key("max_key");

ASSERT_TRUE(document_put(txn.get(), "split_single_message_key_split", std::move(msg)));

ASSERT_EQ(txn->commit(), TxnErrorCode::TXN_OK);

// Read the message back.

std::unique_ptr<Transaction> read_txn;

ASSERT_EQ(txn_kv->create_txn(&read_txn), TxnErrorCode::TXN_OK);

SplitSingleMessagePB read_msg;

ASSERT_EQ(document_get(read_txn.get(), "split_single_message_key_split", &read_msg),

TxnErrorCode::TXN_OK);

ASSERT_EQ(read_msg.other_fields(), 2);

ASSERT_EQ(read_msg.segment_key_bounds().min_key(), "min_key");

ASSERT_EQ(read_msg.segment_key_bounds().max_key(), "max_key");

ASSERT_EQ(count_range(txn_kv.get()), 2) << dump_range(txn_kv.get());

// Remove the message.

document_remove<SplitSingleMessagePB>(read_txn.get(), "split_single_message_key_split");

ASSERT_EQ(read_txn->commit(), TxnErrorCode::TXN_OK);

}

ASSERT_TRUE(is_empty_range(txn_kv.get())) << dump_range(txn_kv.get());

}

TEST(DocumentMessageTest, DocumentPutRowsetMeta) {

std::unique_lock lock(config_mutex);

config::enable_split_rowset_meta_pb = true;

config::split_rowset_meta_pb_size = 0; // Always split the rowset meta pb.

auto txn_kv = std::make_shared<MemTxnKv>();

doris::RowsetMetaCloudPB meta;

meta.set_rowset_id(123);

meta.set_rowset_id_v2("document_put_rowset_meta");

meta.set_start_version(10000);

meta.set_end_version(10005);

meta.set_num_rows(100);

meta.set_num_segments(1);

doris::KeyBoundsPB* key_bounds = meta.mutable_segments_key_bounds()->Add();

key_bounds->set_min_key("min_key");

key_bounds->set_max_key("max_key");

key_bounds = meta.mutable_segments_key_bounds()->Add();

key_bounds->set_min_key("min_key2");

key_bounds->set_max_key("max_key2");

std::string rowset_meta_key =

meta_rowset_key({meta.rowset_id_v2(), meta.end_version(), meta.num_segments()});

{

// create a txn , and put rowset meta

doris::RowsetMetaCloudPB meta_copy(meta);

std::unique_ptr<Transaction> txn;

ASSERT_EQ(txn_kv->create_txn(&txn), TxnErrorCode::TXN_OK);

ASSERT_TRUE(document_put(txn.get(), rowset_meta_key, std::move(meta_copy)));

ASSERT_EQ(txn->commit(), TxnErrorCode::TXN_OK);

}

{

// create a txn , and get rowset meta

std::unique_ptr<Transaction> txn;

ASSERT_EQ(txn_kv->create_txn(&txn), TxnErrorCode::TXN_OK);

doris::RowsetMetaCloudPB saved_meta;

ASSERT_EQ(document_get(txn.get(), rowset_meta_key, &saved_meta), TxnErrorCode::TXN_OK);

ASSERT_EQ(saved_meta.rowset_id_v2(), meta.rowset_id_v2());

ASSERT_EQ(saved_meta.end_version(), meta.end_version());

ASSERT_EQ(saved_meta.start_version(), meta.start_version());

ASSERT_EQ(saved_meta.num_rows(), meta.num_rows());

ASSERT_EQ(saved_meta.num_segments(), meta.num_segments());

ASSERT_EQ(saved_meta.segments_key_bounds_size(), 2);

ASSERT_EQ(saved_meta.segments_key_bounds(0).min_key(), "min_key");

ASSERT_EQ(saved_meta.segments_key_bounds(0).max_key(), "max_key");

ASSERT_EQ(saved_meta.segments_key_bounds(1).min_key(), "min_key2");

ASSERT_EQ(saved_meta.segments_key_bounds(1).max_key(), "max_key2");

}

ASSERT_GE(txn_kv->total_kvs(), 2) << dump_range(txn_kv.get());

{

// create a txn, and remove rowset meta

std::unique_ptr<Transaction> txn;

ASSERT_EQ(txn_kv->create_txn(&txn), TxnErrorCode::TXN_OK);

document_remove<doris::RowsetMetaCloudPB>(txn.get(), rowset_meta_key);

ASSERT_EQ(txn->commit(), TxnErrorCode::TXN_OK);

}

ASSERT_TRUE(is_empty_range(txn_kv.get())) << dump_range(txn_kv.get());

}

TEST(DocumentMessageTest, DocumentPutRowsetMetaWithoutSplit) {

std::unique_lock lock(config_mutex);

config::enable_split_rowset_meta_pb = false;

auto txn_kv = std::make_shared<MemTxnKv>();

doris::RowsetMetaCloudPB meta;

meta.set_rowset_id(123);

meta.set_rowset_id_v2("document_put_rowset_meta");

meta.set_start_version(10000);

meta.set_end_version(10005);

meta.set_num_rows(100);

meta.set_num_segments(1);

doris::KeyBoundsPB* key_bounds = meta.mutable_segments_key_bounds()->Add();

key_bounds->set_min_key("min_key");

key_bounds->set_max_key("max_key");

std::string rowset_meta_key =

meta_rowset_key({meta.rowset_id_v2(), meta.end_version(), meta.num_segments()});

{

// create a txn , and put rowset meta

doris::RowsetMetaCloudPB meta_copy(meta);

std::unique_ptr<Transaction> txn;

ASSERT_EQ(txn_kv->create_txn(&txn), TxnErrorCode::TXN_OK);

ASSERT_TRUE(document_put(txn.get(), rowset_meta_key, std::move(meta_copy)));

ASSERT_EQ(txn->commit(), TxnErrorCode::TXN_OK);

}

{

// create a txn , and get rowset meta

std::unique_ptr<Transaction> txn;

ASSERT_EQ(txn_kv->create_txn(&txn), TxnErrorCode::TXN_OK);

doris::RowsetMetaCloudPB saved_meta;

ASSERT_EQ(document_get(txn.get(), rowset_meta_key, &saved_meta), TxnErrorCode::TXN_OK);

ASSERT_EQ(saved_meta.rowset_id_v2(), meta.rowset_id_v2());

ASSERT_EQ(saved_meta.end_version(), meta.end_version());

ASSERT_EQ(saved_meta.start_version(), meta.start_version());

ASSERT_EQ(saved_meta.num_rows(), meta.num_rows());

ASSERT_EQ(saved_meta.num_segments(), meta.num_segments());

ASSERT_EQ(saved_meta.segments_key_bounds_size(), 1);

ASSERT_EQ(saved_meta.segments_key_bounds(0).min_key(), "min_key");

ASSERT_EQ(saved_meta.segments_key_bounds(0).max_key(), "max_key");

}

ASSERT_EQ(txn_kv->total_kvs(), 1) << dump_range(txn_kv.get());

{

// create a txn, and remove rowset meta

std::unique_ptr<Transaction> txn;

ASSERT_EQ(txn_kv->create_txn(&txn), TxnErrorCode::TXN_OK);

document_remove<doris::RowsetMetaCloudPB>(txn.get(), rowset_meta_key);

ASSERT_EQ(txn->commit(), TxnErrorCode::TXN_OK);

}

ASSERT_TRUE(is_empty_range(txn_kv.get())) << dump_range(txn_kv.get());

{

// Allow split but no thing to split, because the rowset meta pb is small enough.

config::enable_split_rowset_meta_pb = true;

config::split_rowset_meta_pb_size = 100000000; // Big enough to not split.

std::unique_ptr<Transaction> txn;

ASSERT_EQ(txn_kv->create_txn(&txn), TxnErrorCode::TXN_OK);

doris::RowsetMetaCloudPB meta_copy(meta);

ASSERT_TRUE(document_put(txn.get(), rowset_meta_key, std::move(meta_copy)));

ASSERT_EQ(txn->commit(), TxnErrorCode::TXN_OK);

// Read the message back.

std::unique_ptr<Transaction> read_txn;

ASSERT_EQ(txn_kv->create_txn(&read_txn), TxnErrorCode::TXN_OK);

doris::RowsetMetaCloudPB read_meta;

ASSERT_EQ(document_get(read_txn.get(), rowset_meta_key, &read_meta), TxnErrorCode::TXN_OK);

ASSERT_EQ(read_meta.rowset_id_v2(), meta.rowset_id_v2());

ASSERT_EQ(read_meta.end_version(), meta.end_version());

ASSERT_EQ(read_meta.start_version(), meta.start_version());

ASSERT_EQ(read_meta.num_rows(), meta.num_rows());

ASSERT_EQ(read_meta.num_segments(), meta.num_segments());

ASSERT_EQ(read_meta.segments_key_bounds_size(), 1);

ASSERT_EQ(read_meta.segments_key_bounds(0).min_key(), "min_key");

ASSERT_EQ(read_meta.segments_key_bounds(0).max_key(), "max_key");

// Remove the message.

document_remove<doris::RowsetMetaCloudPB>(read_txn.get(), rowset_meta_key);

ASSERT_EQ(read_txn->commit(), TxnErrorCode::TXN_OK);

}

{

// Allow split but no thing to split, because the split fields is empty.

config::enable_split_rowset_meta_pb = true;

config::split_rowset_meta_pb_size = 0; // Always split the rowset meta pb.

std::unique_ptr<Transaction> txn;

ASSERT_EQ(txn_kv->create_txn(&txn), TxnErrorCode::TXN_OK);

doris::RowsetMetaCloudPB meta_copy(meta);

meta_copy.clear_segments_key_bounds(); // Clear the segments key bounds to make it empty.

ASSERT_TRUE(document_put(txn.get(), rowset_meta_key, std::move(meta_copy)));

ASSERT_EQ(txn->commit(), TxnErrorCode::TXN_OK);

// If you meat a failed here, it means the rowset meta pb is split into multiple keys,

// you should clear the split fields you added.

ASSERT_EQ(txn_kv->total_kvs(), 1) << dump_range(txn_kv.get());

// Read the message back.

std::unique_ptr<Transaction> read_txn;

ASSERT_EQ(txn_kv->create_txn(&read_txn), TxnErrorCode::TXN_OK);

doris::RowsetMetaCloudPB read_meta;

ASSERT_EQ(document_get(read_txn.get(), rowset_meta_key, &read_meta), TxnErrorCode::TXN_OK);

ASSERT_EQ(read_meta.rowset_id_v2(), meta.rowset_id_v2());

ASSERT_EQ(read_meta.end_version(), meta.end_version());

ASSERT_EQ(read_meta.start_version(), meta.start_version());

ASSERT_EQ(read_meta.num_rows(), meta.num_rows());

ASSERT_EQ(read_meta.num_segments(), meta.num_segments());

ASSERT_EQ(read_meta.segments_key_bounds_size(), 0); // No segments key bounds

// Remove the message.

document_remove<doris::RowsetMetaCloudPB>(read_txn.get(), rowset_meta_key);

ASSERT_EQ(read_txn->commit(), TxnErrorCode::TXN_OK);

}

ASSERT_TRUE(is_empty_range(txn_kv.get())) << dump_range(txn_kv.get());

}

TEST(DocumentMessageTest, DocumentPutTabletSchemaCloudPB) {

std::unique_lock lock(config_mutex);

config::enable_split_tablet_schema_pb = true;

config::split_tablet_schema_pb_size = 0; // Always split the tablet schema pb.

auto txn_kv = std::make_shared<MemTxnKv>();

doris::TabletSchemaCloudPB schema;

schema.set_keys_type(doris::DUP_KEYS);

schema.set_num_short_key_columns(2);

schema.set_num_rows_per_row_block(1024);

schema.set_compress_kind(doris::COMPRESS_LZ4);

schema.set_next_column_unique_id(10);

schema.set_schema_version(1);

// Add multiple columns to test splitting

for (int i = 0; i < 3; ++i) {

doris::ColumnPB* column = schema.add_column();

column->set_name(fmt::format("col_{}", i));

column->set_unique_id(i + 1);

column->set_type("INT");

column->set_is_key(i < 2); // First 2 columns are key columns

column->set_is_nullable(false);

}

std::string schema_key = versioned::meta_schema_key({"instance_id", 10, 0});

{

// create a txn and put tablet schema

doris::TabletSchemaCloudPB schema_copy(schema);

std::unique_ptr<Transaction> txn;

ASSERT_EQ(txn_kv->create_txn(&txn), TxnErrorCode::TXN_OK);

ASSERT_TRUE(document_put(txn.get(), schema_key, std::move(schema_copy)));

ASSERT_EQ(txn->commit(), TxnErrorCode::TXN_OK);

}

{

// create a txn and get tablet schema

std::unique_ptr<Transaction> txn;

ASSERT_EQ(txn_kv->create_txn(&txn), TxnErrorCode::TXN_OK);

doris::TabletSchemaCloudPB saved_schema;

ASSERT_EQ(document_get(txn.get(), schema_key, &saved_schema), TxnErrorCode::TXN_OK);

ASSERT_EQ(saved_schema.keys_type(), schema.keys_type());

ASSERT_EQ(saved_schema.num_short_key_columns(), schema.num_short_key_columns());

ASSERT_EQ(saved_schema.num_rows_per_row_block(), schema.num_rows_per_row_block());

ASSERT_EQ(saved_schema.compress_kind(), schema.compress_kind());

ASSERT_EQ(saved_schema.next_column_unique_id(), schema.next_column_unique_id());

ASSERT_EQ(saved_schema.schema_version(), schema.schema_version());

ASSERT_EQ(saved_schema.column_size(), 3);

for (int i = 0; i < 3; ++i) {

ASSERT_EQ(saved_schema.column(i).name(), fmt::format("col_{}", i));

ASSERT_EQ(saved_schema.column(i).unique_id(), i + 1);

ASSERT_EQ(saved_schema.column(i).type(), "INT");

ASSERT_EQ(saved_schema.column(i).is_key(), i < 2);

ASSERT_EQ(saved_schema.column(i).is_nullable(), false);

}

}

ASSERT_GE(txn_kv->total_kvs(), 2) << dump_range(txn_kv.get());

{

std::unique_ptr<Transaction> txn;

ASSERT_EQ(txn_kv->create_txn(&txn), TxnErrorCode::TXN_OK);

std::string begin_key = versioned::meta_schema_key({"instance_id", 10, 0});

std::string end_key = versioned::meta_schema_key({"instance_id", 10, INT64_MAX});

FullRangeGetOptions options;

std::unique_ptr<FullRangeGetIterator> iter =

txn->full_range_get(begin_key, end_key, std::move(options));

for (auto kvp = iter->next(); kvp.has_value(); kvp = iter->next()) {

auto&& [key, _] = *kvp;

int64_t decode_index_id = -1;

int64_t schema_version = -1;

std::string_view key_view(key);

ASSERT_TRUE(versioned::decode_meta_schema_key(&key_view, &decode_index_id,

&schema_version));

ASSERT_EQ(10, decode_index_id);

ASSERT_EQ(0, schema_version);

}

}

{

// create a txn and remove tablet schema

std::unique_ptr<Transaction> txn;

ASSERT_EQ(txn_kv->create_txn(&txn), TxnErrorCode::TXN_OK);

document_remove<doris::TabletSchemaCloudPB>(txn.get(), schema_key);

ASSERT_EQ(txn->commit(), TxnErrorCode::TXN_OK);

}

ASSERT_TRUE(is_empty_range(txn_kv.get())) << dump_range(txn_kv.get());

}

TEST(DocumentMessageTest, DocumentPutTabletSchemaCloudPBWithoutSplit) {

std::unique_lock lock(config_mutex);

config::enable_split_tablet_schema_pb = false;

auto txn_kv = std::make_shared<MemTxnKv>();

doris::TabletSchemaCloudPB schema;

schema.set_keys_type(doris::DUP_KEYS);

schema.set_num_short_key_columns(1);

schema.set_num_rows_per_row_block(1024);

schema.set_compress_kind(doris::COMPRESS_LZ4);

schema.set_next_column_unique_id(5);

schema.set_schema_version(1);

// Add one column

doris::ColumnPB* column = schema.add_column();

column->set_name("test_col");

column->set_unique_id(1);

column->set_type("VARCHAR");

column->set_is_key(true);

column->set_is_nullable(false);

std::string schema_key = "tablet_schema_key_no_split";

{

// create a txn and put tablet schema

doris::TabletSchemaCloudPB schema_copy(schema);

std::unique_ptr<Transaction> txn;

ASSERT_EQ(txn_kv->create_txn(&txn), TxnErrorCode::TXN_OK);

ASSERT_TRUE(document_put(txn.get(), schema_key, std::move(schema_copy)));

ASSERT_EQ(txn->commit(), TxnErrorCode::TXN_OK);

}

{

// create a txn and get tablet schema

std::unique_ptr<Transaction> txn;

ASSERT_EQ(txn_kv->create_txn(&txn), TxnErrorCode::TXN_OK);

doris::TabletSchemaCloudPB saved_schema;

ASSERT_EQ(document_get(txn.get(), schema_key, &saved_schema), TxnErrorCode::TXN_OK);

ASSERT_EQ(saved_schema.keys_type(), schema.keys_type());

ASSERT_EQ(saved_schema.num_short_key_columns(), schema.num_short_key_columns());

ASSERT_EQ(saved_schema.column_size(), 1);

ASSERT_EQ(saved_schema.column(0).name(), "test_col");

ASSERT_EQ(saved_schema.column(0).unique_id(), 1);

ASSERT_EQ(saved_schema.column(0).type(), "VARCHAR");

}

ASSERT_EQ(txn_kv->total_kvs(), 1) << dump_range(txn_kv.get());

{

// create a txn and remove tablet schema

std::unique_ptr<Transaction> txn;

ASSERT_EQ(txn_kv->create_txn(&txn), TxnErrorCode::TXN_OK);

document_remove<doris::TabletSchemaCloudPB>(txn.get(), schema_key);

ASSERT_EQ(txn->commit(), TxnErrorCode::TXN_OK);

}

ASSERT_TRUE(is_empty_range(txn_kv.get())) << dump_range(txn_kv.get());

{

// Allow split but nothing to split, because the split fields is empty.

config::enable_split_tablet_schema_pb = true;

config::split_tablet_schema_pb_size = 0; // Always split

std::unique_ptr<Transaction> txn;

ASSERT_EQ(txn_kv->create_txn(&txn), TxnErrorCode::TXN_OK);

doris::TabletSchemaCloudPB schema_copy(schema);

schema_copy.clear_column(); // Clear the columns to make it empty.

ASSERT_TRUE(document_put(txn.get(), schema_key, std::move(schema_copy)));

ASSERT_EQ(txn->commit(), TxnErrorCode::TXN_OK);

// If you meet a failure here, it means the tablet schema pb is split into multiple keys,

// you should clear the split fields you added.

ASSERT_EQ(txn_kv->total_kvs(), 1) << dump_range(txn_kv.get());

// Read the message back.

std::unique_ptr<Transaction> read_txn;

ASSERT_EQ(txn_kv->create_txn(&read_txn), TxnErrorCode::TXN_OK);

doris::TabletSchemaCloudPB read_schema;

ASSERT_EQ(document_get(read_txn.get(), schema_key, &read_schema), TxnErrorCode::TXN_OK);

ASSERT_EQ(read_schema.keys_type(), schema.keys_type());

ASSERT_EQ(read_schema.num_short_key_columns(), schema.num_short_key_columns());

ASSERT_EQ(read_schema.column_size(), 0); // No columns

// Remove the message.

document_remove<doris::TabletSchemaCloudPB>(read_txn.get(), schema_key);

ASSERT_EQ(read_txn->commit(), TxnErrorCode::TXN_OK);

}

ASSERT_TRUE(is_empty_range(txn_kv.get())) << dump_range(txn_kv.get());

}

TEST(DocumentMessageTest, DocumentPutVersionPB) {

auto txn_kv = std::make_shared<MemTxnKv>();

VersionPB version_pb;

version_pb.set_version(123);

{

// create a txn , and put version pb

std::unique_ptr<Transaction> txn;

ASSERT_EQ(txn_kv->create_txn(&txn), TxnErrorCode::TXN_OK);

ASSERT_TRUE(document_put(txn.get(), "version_key", std::move(version_pb)));

ASSERT_EQ(txn->commit(), TxnErrorCode::TXN_OK);

}

{

// create a txn , and get version pb

std::unique_ptr<Transaction> txn;

ASSERT_EQ(txn_kv->create_txn(&txn), TxnErrorCode::TXN_OK);

VersionPB saved_version_pb;

ASSERT_EQ(document_get(txn.get(), "version_key", &saved_version_pb), TxnErrorCode::TXN_OK);

ASSERT_EQ(saved_version_pb.version(), 123);

}

ASSERT_EQ(txn_kv->total_kvs(), 1) << dump_range(txn_kv.get());

{

// create a txn, and remove version pb

std::unique_ptr<Transaction> txn;

ASSERT_EQ(txn_kv->create_txn(&txn), TxnErrorCode::TXN_OK);

document_remove<VersionPB>(txn.get(), "version_key");

ASSERT_EQ(txn->commit(), TxnErrorCode::TXN_OK);

}

ASSERT_TRUE(is_empty_range(txn_kv.get())) << dump_range(txn_kv.get());

}

TEST(DocumentMessageTest, DocumentGetAllRowsetMeta) {

std::unique_lock lock(config_mutex);

config::enable_split_rowset_meta_pb = true;

config::split_rowset_meta_pb_size = 0; // Always split the rowset meta pb.

auto txn_kv = std::make_shared<MemTxnKv>();

doris::RowsetMetaCloudPB meta;

meta.set_rowset_id(123);

meta.set_rowset_id_v2("document_put_rowset_meta");

meta.set_start_version(10000);

meta.set_end_version(10005);

meta.set_num_rows(100);

meta.set_num_segments(1);

doris::KeyBoundsPB* key_bounds = meta.mutable_segments_key_bounds()->Add();

key_bounds->set_min_key("min_key");

key_bounds->set_max_key("max_key");

std::string rowset_meta_key_prefix = "rowset_meta_key_";

for (int i = 0; i < 10; ++i) {

// create a txn , and put rowset meta

std::string rowset_meta_key = rowset_meta_key_prefix + std::to_string(i);

doris::RowsetMetaCloudPB meta_copy(meta);

std::unique_ptr<Transaction> txn;

ASSERT_EQ(txn_kv->create_txn(&txn), TxnErrorCode::TXN_OK);

ASSERT_TRUE(document_put(txn.get(), rowset_meta_key, std::move(meta_copy)));

ASSERT_EQ(txn->commit(), TxnErrorCode::TXN_OK);

}

{

// scan all rowset meta keys

std::unique_ptr<Transaction> txn;

ASSERT_EQ(txn_kv->create_txn(&txn), TxnErrorCode::TXN_OK);

FullRangeGetOptions opts;

opts.txn = txn.get();

opts.batch_limit = 2;

auto iter = txn_kv->full_range_get(rowset_meta_key_prefix, rowset_meta_key_prefix + '\xFF',

std::move(opts));

int count = 0;

while (true) {

doris::RowsetMetaCloudPB saved_meta;

TxnErrorCode err = document_get(iter.get(), &saved_meta);

if (err == TxnErrorCode::TXN_KEY_NOT_FOUND) {

break; // No more keys to read

}

ASSERT_EQ(err, TxnErrorCode::TXN_OK);

ASSERT_EQ(saved_meta.rowset_id_v2(), meta.rowset_id_v2());

ASSERT_EQ(saved_meta.end_version(), meta.end_version());

ASSERT_EQ(saved_meta.start_version(), meta.start_version());

ASSERT_EQ(saved_meta.num_rows(), meta.num_rows());

ASSERT_EQ(saved_meta.num_segments(), meta.num_segments());

ASSERT_EQ(saved_meta.segments_key_bounds_size(), 1);

ASSERT_EQ(saved_meta.segments_key_bounds(0).min_key(), "min_key");

ASSERT_EQ(saved_meta.segments_key_bounds(0).max_key(), "max_key");

count += 1;

}

ASSERT_EQ(count, 10) << dump_range(txn_kv.get());

}

}

TEST(DocumentMessageTest, DocumentGetAllTabletSchemaCloudPB) {

std::unique_lock lock(config_mutex);

config::enable_split_tablet_schema_pb = true;

config::split_tablet_schema_pb_size = 0; // Always split the tablet schema pb.

auto txn_kv = std::make_shared<MemTxnKv>();

doris::TabletSchemaCloudPB schema;

schema.set_keys_type(doris::DUP_KEYS);

schema.set_num_short_key_columns(1);

schema.set_num_rows_per_row_block(1024);

schema.set_compress_kind(doris::COMPRESS_LZ4);

schema.set_next_column_unique_id(5);

schema.set_schema_version(1);

doris::ColumnPB* column = schema.add_column();

column->set_name("test_col");

column->set_unique_id(1);

column->set_type("VARCHAR");

column->set_is_key(true);

column->set_is_nullable(false);

std::string schema_key_prefix = "tablet_schema_key_";

for (int i = 0; i < 10; ++i) {

// create a txn and put tablet schema

doris::TabletSchemaCloudPB schema_copy(schema);

schema_copy.set_schema_version(i + 1);

std::unique_ptr<Transaction> txn;

ASSERT_EQ(txn_kv->create_txn(&txn), TxnErrorCode::TXN_OK);

ASSERT_TRUE(document_put(txn.get(), schema_key_prefix + std::to_string(i),

std::move(schema_copy)));

ASSERT_EQ(txn->commit(), TxnErrorCode::TXN_OK);

}

{

// scan all tablet schema keys

std::unique_ptr<Transaction> txn;

ASSERT_EQ(txn_kv->create_txn(&txn), TxnErrorCode::TXN_OK);

FullRangeGetOptions opts;

opts.txn = txn.get();

opts.batch_limit = 2;

auto iter = txn_kv->full_range_get(schema_key_prefix, schema_key_prefix + '\xFF',

std::move(opts));

int count = 0;

while (true) {

doris::TabletSchemaCloudPB saved_schema;

TxnErrorCode code = document_get(iter.get(), &saved_schema);

if (code == TxnErrorCode::TXN_KEY_NOT_FOUND) {

break;

}

ASSERT_EQ(code, TxnErrorCode::TXN_OK);

ASSERT_EQ(saved_schema.keys_type(), schema.keys_type());

ASSERT_EQ(saved_schema.num_short_key_columns(), schema.num_short_key_columns());

ASSERT_EQ(saved_schema.column_size(), 1);

ASSERT_EQ(saved_schema.column(0).name(), "test_col");

ASSERT_EQ(saved_schema.column(0).unique_id(), 1);

ASSERT_EQ(saved_schema.column(0).type(), "VARCHAR");

ASSERT_EQ(saved_schema.schema_version(), count + 1);

count += 1;

}

ASSERT_EQ(count, 10) << dump_range(txn_kv.get());

}

}

TEST(DocumentMessageTest, VersionedDocumentPut) {

std::unique_lock lock(config_mutex);

config::enable_split_rowset_meta_pb = true;

config::split_rowset_meta_pb_size = 0; // Always split the rowset meta pb.

auto txn_kv = std::make_shared<MemTxnKv>();

doris::RowsetMetaCloudPB meta;

meta.set_rowset_id(123);

meta.set_rowset_id_v2("versioned_document_test");

meta.set_start_version(10000);

meta.set_end_version(10005);

meta.set_num_rows(100);

meta.set_num_segments(1);

doris::KeyBoundsPB* key_bounds = meta.mutable_segments_key_bounds()->Add();

key_bounds->set_min_key("min_key");

key_bounds->set_max_key("max_key");

std::string key_prefix = "versioned_rowset_meta_key_";

{

// Create a txn and put versioned rowset meta with initial version

doris::RowsetMetaCloudPB meta_copy(meta);

std::unique_ptr<Transaction> txn;

ASSERT_EQ(txn_kv->create_txn(&txn), TxnErrorCode::TXN_OK);

ASSERT_TRUE(versioned::document_put(txn.get(), key_prefix, std::move(meta_copy)));

ASSERT_EQ(txn->commit(), TxnErrorCode::TXN_OK);

}

Versionstamp version;

{

// Get the rowset meta with version

std::unique_ptr<Transaction> txn;

ASSERT_EQ(txn_kv->create_txn(&txn), TxnErrorCode::TXN_OK);

doris::RowsetMetaCloudPB read_meta;

ASSERT_EQ(versioned::document_get(txn.get(), key_prefix, &read_meta, &version),

TxnErrorCode::TXN_OK);

}

// Remove the range

{

// Create a txn and remove the versioned rowset meta

std::unique_ptr<Transaction> txn;

ASSERT_EQ(txn_kv->create_txn(&txn), TxnErrorCode::TXN_OK);

versioned::document_remove<doris::RowsetMetaCloudPB>(txn.get(), key_prefix, version);

ASSERT_EQ(txn->commit(), TxnErrorCode::TXN_OK);

}

ASSERT_TRUE(is_empty_range(txn_kv.get())) << dump_range(txn_kv.get());

{

// Create a txn and try to get the removed versioned rowset meta

std::unique_ptr<Transaction> txn;

ASSERT_EQ(txn_kv->create_txn(&txn), TxnErrorCode::TXN_OK);

doris::RowsetMetaCloudPB read_meta;

ASSERT_EQ(versioned::document_get(txn.get(), key_prefix, &read_meta, &version),

TxnErrorCode::TXN_KEY_NOT_FOUND);

}

}

TEST(DocumentMessageTest, VersionedDocumentPutAndGet) {

std::unique_lock lock(config_mutex);

config::enable_split_rowset_meta_pb = true;

config::split_rowset_meta_pb_size = 0; // Always split the rowset meta pb.

auto txn_kv = std::make_shared<MemTxnKv>();

doris::RowsetMetaCloudPB meta;

meta.set_rowset_id(123);

meta.set_rowset_id_v2("versioned_document_test");

meta.set_start_version(10000);

meta.set_end_version(10005);

meta.set_num_rows(100);

meta.set_num_segments(1);

doris::KeyBoundsPB* key_bounds = meta.mutable_segments_key_bounds()->Add();

key_bounds->set_min_key("min_key");

key_bounds->set_max_key("max_key");

key_bounds = meta.mutable_segments_key_bounds()->Add();

key_bounds->set_min_key("min_key2");

key_bounds->set_max_key("max_key2");

std::string key_prefix = "versioned_rowset_meta_key_";

{

// Create a txn and put versioned rowset meta with initial version

doris::RowsetMetaCloudPB meta_copy(meta);

std::unique_ptr<Transaction> txn;

ASSERT_EQ(txn_kv->create_txn(&txn), TxnErrorCode::TXN_OK);

ASSERT_TRUE(versioned::document_put(txn.get(), key_prefix, std::move(meta_copy)));

ASSERT_EQ(txn->commit(), TxnErrorCode::TXN_OK);

}

{

// Get the rowset meta with version

std::unique_ptr<Transaction> txn;

ASSERT_EQ(txn_kv->create_txn(&txn), TxnErrorCode::TXN_OK);

doris::RowsetMetaCloudPB read_meta;

Versionstamp version;

ASSERT_EQ(versioned::document_get(txn.get(), key_prefix, &read_meta, &version),

TxnErrorCode::TXN_OK);

// Verify the metadata

ASSERT_EQ(read_meta.rowset_id_v2(), meta.rowset_id_v2());

ASSERT_EQ(read_meta.end_version(), meta.end_version());

ASSERT_EQ(read_meta.start_version(), meta.start_version());

ASSERT_EQ(read_meta.num_rows(), meta.num_rows());

ASSERT_EQ(read_meta.num_segments(), meta.num_segments());

ASSERT_EQ(read_meta.segments_key_bounds_size(), 2);

ASSERT_EQ(read_meta.segments_key_bounds(0).min_key(), "min_key");

ASSERT_EQ(read_meta.segments_key_bounds(0).max_key(), "max_key");

ASSERT_EQ(read_meta.segments_key_bounds(1).min_key(), "min_key2");

ASSERT_EQ(read_meta.segments_key_bounds(1).max_key(), "max_key2");

}

{

// Update the rowset meta with a new version

doris::RowsetMetaCloudPB updated_meta(meta);

updated_meta.set_num_rows(200); // Change num_rows

std::unique_ptr<Transaction> txn;

ASSERT_EQ(txn_kv->create_txn(&txn), TxnErrorCode::TXN_OK);

ASSERT_TRUE(versioned::document_put(txn.get(), key_prefix, std::move(updated_meta)));

ASSERT_EQ(txn->commit(), TxnErrorCode::TXN_OK);

}

{

// Get the latest version of rowset meta

std::unique_ptr<Transaction> txn;

ASSERT_EQ(txn_kv->create_txn(&txn), TxnErrorCode::TXN_OK);

doris::RowsetMetaCloudPB read_meta;

Versionstamp version;

ASSERT_EQ(versioned::document_get(txn.get(), key_prefix, &read_meta, &version),

TxnErrorCode::TXN_OK);

// Verify the updated metadata

ASSERT_EQ(read_meta.rowset_id_v2(), meta.rowset_id_v2());

ASSERT_EQ(read_meta.end_version(), meta.end_version());

ASSERT_EQ(read_meta.start_version(), meta.start_version());

ASSERT_EQ(read_meta.num_rows(), 200); // This should now be 200

ASSERT_EQ(read_meta.num_segments(), meta.num_segments());

ASSERT_EQ(read_meta.segments_key_bounds_size(), 2);

ASSERT_EQ(read_meta.segments_key_bounds(0).min_key(), "min_key");

ASSERT_EQ(read_meta.segments_key_bounds(0).max_key(), "max_key");

ASSERT_EQ(read_meta.segments_key_bounds(1).min_key(), "min_key2");

ASSERT_EQ(read_meta.segments_key_bounds(1).max_key(), "max_key2");

}

// Testing with a non-splittable message

{

std::string simple_key_prefix = "versioned_version_pb_key_";

VersionPB version_pb;

version_pb.set_version(123);

// Put the simple message

std::unique_ptr<Transaction> txn;

ASSERT_EQ(txn_kv->create_txn(&txn), TxnErrorCode::TXN_OK);

ASSERT_TRUE(versioned::document_put(txn.get(), simple_key_prefix, std::move(version_pb)));

ASSERT_EQ(txn->commit(), TxnErrorCode::TXN_OK);

// Get the simple message

std::unique_ptr<Transaction> read_txn;

ASSERT_EQ(txn_kv->create_txn(&read_txn), TxnErrorCode::TXN_OK);

VersionPB read_version_pb;

Versionstamp version;

ASSERT_EQ(versioned::document_get(read_txn.get(), simple_key_prefix, &read_version_pb,

&version),

TxnErrorCode::TXN_OK);

// Verify the simple message

ASSERT_EQ(read_version_pb.version(), 123);

}

}

TEST(DocumentMessageTest, VersionedDocumentPutAndGetWithVersionstamp) {

std::unique_lock lock(config_mutex);

config::enable_split_rowset_meta_pb = true;

config::split_rowset_meta_pb_size = 0; // Always split the rowset meta pb.

auto txn_kv = std::make_shared<MemTxnKv>();

doris::RowsetMetaCloudPB meta;

meta.set_rowset_id(123);

meta.set_rowset_id_v2("versioned_document_test");

meta.set_start_version(10000);

meta.set_end_version(10005);

meta.set_num_rows(100);

meta.set_num_segments(1);

doris::KeyBoundsPB* key_bounds = meta.mutable_segments_key_bounds()->Add();

key_bounds->set_min_key("min_key");

key_bounds->set_max_key("max_key");

key_bounds = meta.mutable_segments_key_bounds()->Add();

key_bounds->set_min_key("min_key2");

key_bounds->set_max_key("max_key2");

std::string key_prefix = "versioned_rowset_meta_key_";

{

// Create a txn and put versioned rowset meta with initial version

doris::RowsetMetaCloudPB meta_copy(meta);

std::unique_ptr<Transaction> txn;

ASSERT_EQ(txn_kv->create_txn(&txn), TxnErrorCode::TXN_OK);

ASSERT_TRUE(versioned::document_put(txn.get(), key_prefix, std::move(meta_copy)));

ASSERT_EQ(txn->commit(), TxnErrorCode::TXN_OK);

}

Versionstamp first_versionstamp;

{

// Get the rowset meta with version

std::unique_ptr<Transaction> txn;

ASSERT_EQ(txn_kv->create_txn(&txn), TxnErrorCode::TXN_OK);

doris::RowsetMetaCloudPB read_meta;

ASSERT_EQ(versioned::document_get(txn.get(), key_prefix, &read_meta, &first_versionstamp),

TxnErrorCode::TXN_OK);

// Verify the metadata

ASSERT_EQ(read_meta.rowset_id_v2(), meta.rowset_id_v2());

ASSERT_EQ(read_meta.end_version(), meta.end_version());

ASSERT_EQ(read_meta.start_version(), meta.start_version());

ASSERT_EQ(read_meta.num_rows(), meta.num_rows());

ASSERT_EQ(read_meta.num_segments(), meta.num_segments());

ASSERT_EQ(read_meta.segments_key_bounds_size(), 2);

ASSERT_EQ(read_meta.segments_key_bounds(0).min_key(), "min_key");

ASSERT_EQ(read_meta.segments_key_bounds(0).max_key(), "max_key");

ASSERT_EQ(read_meta.segments_key_bounds(1).min_key(), "min_key2");

ASSERT_EQ(read_meta.segments_key_bounds(1).max_key(), "max_key2");

}

{

// Update the rowset meta with a new version

doris::RowsetMetaCloudPB updated_meta(meta);

updated_meta.set_num_rows(200); // Change num_rows

std::unique_ptr<Transaction> txn;

ASSERT_EQ(txn_kv->create_txn(&txn), TxnErrorCode::TXN_OK);

ASSERT_TRUE(versioned::document_put(txn.get(), key_prefix, std::move(updated_meta)));

ASSERT_EQ(txn->commit(), TxnErrorCode::TXN_OK);

}

{

// Get the latest version of rowset meta

std::unique_ptr<Transaction> txn;

ASSERT_EQ(txn_kv->create_txn(&txn), TxnErrorCode::TXN_OK);

doris::RowsetMetaCloudPB read_meta;

Versionstamp version;

ASSERT_EQ(versioned::document_get(txn.get(), key_prefix, &read_meta, &version),

TxnErrorCode::TXN_OK);

// Verify the updated metadata

ASSERT_EQ(read_meta.rowset_id_v2(), meta.rowset_id_v2());

ASSERT_EQ(read_meta.end_version(), meta.end_version());

ASSERT_EQ(read_meta.start_version(), meta.start_version());

ASSERT_EQ(read_meta.num_rows(), 200); // This should now be 200

ASSERT_EQ(read_meta.num_segments(), meta.num_segments());

ASSERT_EQ(read_meta.segments_key_bounds_size(), 2);

ASSERT_EQ(read_meta.segments_key_bounds(0).min_key(), "min_key");

ASSERT_EQ(read_meta.segments_key_bounds(0).max_key(), "max_key");

ASSERT_EQ(read_meta.segments_key_bounds(1).min_key(), "min_key2");

ASSERT_EQ(read_meta.segments_key_bounds(1).max_key(), "max_key2");

}

{

// Get the previous version of rowset meta

std::unique_ptr<Transaction> txn;

ASSERT_EQ(txn_kv->create_txn(&txn), TxnErrorCode::TXN_OK);

doris::RowsetMetaCloudPB read_meta;

Versionstamp snapshot_version(first_versionstamp.version() + 1, 0);

Versionstamp version;

ASSERT_EQ(versioned::document_get(txn.get(), key_prefix, snapshot_version, &read_meta,

&version),

TxnErrorCode::TXN_OK);

// Verify the updated metadata

ASSERT_EQ(read_meta.rowset_id_v2(), meta.rowset_id_v2());

ASSERT_EQ(read_meta.end_version(), meta.end_version());

ASSERT_EQ(read_meta.start_version(), meta.start_version());

ASSERT_EQ(read_meta.num_rows(), meta.num_rows());

ASSERT_EQ(read_meta.num_segments(), meta.num_segments());

ASSERT_EQ(read_meta.segments_key_bounds_size(), 2);

ASSERT_EQ(read_meta.segments_key_bounds(0).min_key(), "min_key");

ASSERT_EQ(read_meta.segments_key_bounds(0).max_key(), "max_key");

ASSERT_EQ(read_meta.segments_key_bounds(1).min_key(), "min_key2");

ASSERT_EQ(read_meta.segments_key_bounds(1).max_key(), "max_key2");

}

{

// Not found

std::unique_ptr<Transaction> txn;

ASSERT_EQ(txn_kv->create_txn(&txn), TxnErrorCode::TXN_OK);

doris::RowsetMetaCloudPB read_meta;

Versionstamp snapshot_version(Versionstamp::min());

Versionstamp version;

ASSERT_EQ(versioned::document_get(txn.get(), key_prefix, snapshot_version, &read_meta,

&version),

TxnErrorCode::TXN_KEY_NOT_FOUND)

<< dump_range(txn_kv.get());

}

}

TEST(DocumentMessageTest, VersionedDocumentPutTabletSchemaCloudPB) {

std::unique_lock lock(config_mutex);

config::enable_split_tablet_schema_pb = true;

config::split_tablet_schema_pb_size = 0; // Always split the tablet schema pb.

auto txn_kv = std::make_shared<MemTxnKv>();

doris::TabletSchemaCloudPB schema;

schema.set_keys_type(doris::DUP_KEYS);

schema.set_num_short_key_columns(2);

schema.set_num_rows_per_row_block(1024);

schema.set_compress_kind(doris::COMPRESS_LZ4);

schema.set_next_column_unique_id(10);

schema.set_schema_version(1);

// Add multiple columns to test splitting

for (int i = 0; i < 3; ++i) {

doris::ColumnPB* column = schema.add_column();

column->set_name(fmt::format("col_{}", i));

column->set_unique_id(i + 1);

column->set_type("INT");

column->set_is_key(i < 2); // First 2 columns are key columns

column->set_is_nullable(false);

}

std::string key_prefix = "versioned_tablet_schema_key_";