// 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/blob_message.h"

#include <gen_cpp/cloud.pb.h>

#include <gen_cpp/olap_file.pb.h>

#include <gtest/gtest.h>

#include <limits>

#include <memory>

#include <string>

#include "meta-store/mem_txn_kv.h"

#include "meta-store/txn_kv.h"

#include "meta-store/txn_kv_error.h"

#include "meta-store/versioned_value.h"

using namespace doris::cloud;

// 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 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: {}\n", escape_hex(kv->first));

}

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

return buffer;

}

// Test blob_put and blob_get with small message (single KV)

TEST(BlobMessageTest, PutGetSmallMessage) {

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

ASSERT_EQ(txn_kv->init(), 0);

std::unique_ptr<Transaction> txn;

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

// Create a small protobuf message

doris::RowsetMetaCloudPB rowset_meta;

rowset_meta.set_rowset_id(12345);

rowset_meta.set_tablet_id(678);

rowset_meta.set_num_rows(100);

rowset_meta.set_data_disk_size(1024);

std::string key = "test_blob_small_message";

uint8_t version = 1;

// Put the message

blob_put(txn.get(), key, rowset_meta, version);

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

// Get the message

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

ValueBuf val;

ASSERT_EQ(blob_get(txn.get(), key, &val), TxnErrorCode::TXN_OK);

ASSERT_EQ(val.ver, version);

// Parse and verify

doris::RowsetMetaCloudPB retrieved_meta;

ASSERT_TRUE(val.to_pb(&retrieved_meta));

ASSERT_EQ(retrieved_meta.rowset_id(), rowset_meta.rowset_id());

ASSERT_EQ(retrieved_meta.tablet_id(), rowset_meta.tablet_id());

ASSERT_EQ(retrieved_meta.num_rows(), rowset_meta.num_rows());

ASSERT_EQ(retrieved_meta.data_disk_size(), rowset_meta.data_disk_size());

}

static OperationLogPB create_operation_log(int num_tablets, int num_partitions) {

OperationLogPB operation_log;

CommitTxnLogPB* commit_txn_log = operation_log.mutable_commit_txn();

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

int64_t tablet_id = static_cast<int64_t>(std::numeric_limits<uint32_t>::max()) + i;

commit_txn_log->mutable_tablet_to_partition_map()->insert({tablet_id, i + 1000});

}

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

int64_t partition_id = static_cast<int64_t>(std::numeric_limits<uint32_t>::max()) + i;

commit_txn_log->mutable_partition_version_map()->insert({partition_id, i});

}

return operation_log;

}

// Test blob_put and blob_get with large message (split into multiple KVs)

TEST(BlobMessageTest, PutGetLargeMessage) {

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

ASSERT_EQ(txn_kv->init(), 0);

std::unique_ptr<Transaction> txn;

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

// Create a large protobuf message that will be split, 100k tablet and 10k partitions

OperationLogPB operation_log = create_operation_log(100000, 10000);

std::string key = "test_blob_large_message";

uint8_t version = 2;

size_t split_size = 50000; // 50KB split size

// Put the large message

versioned::blob_put(txn.get(), key, operation_log, version, split_size);

txn->enable_get_versionstamp();

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

Versionstamp vs;

ASSERT_EQ(txn->get_versionstamp(&vs), TxnErrorCode::TXN_OK);

// Get the message

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

ValueBuf val;

std::string versioned_key = encode_versioned_key(key, vs);

ASSERT_EQ(blob_get(txn.get(), versioned_key, &val), TxnErrorCode::TXN_OK)

<< dump_range(txn_kv.get()) << ", versioned_key=" << escape_hex(versioned_key);

ASSERT_EQ(val.ver, version);

// Verify it was split into multiple KVs

std::vector<std::string> keys = val.keys();

ASSERT_GT(keys.size(), 1) << "Large message should be split into multiple KVs";

// Parse and verify

OperationLogPB retrieved_log;

ASSERT_TRUE(val.to_pb(&retrieved_log));

ASSERT_EQ(retrieved_log.commit_txn().tablet_to_partition_map_size(),

operation_log.commit_txn().tablet_to_partition_map_size());

ASSERT_EQ(retrieved_log.commit_txn().partition_version_map_size(),

operation_log.commit_txn().partition_version_map_size());

}

// Test blob_put with string value

TEST(BlobMessageTest, PutGetStringValue) {

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

ASSERT_EQ(txn_kv->init(), 0);

std::unique_ptr<Transaction> txn;

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

std::string key = "test_blob_string_value";

std::string value = "This is a test string value";

uint8_t version = 3;

// Put the string value

blob_put(txn.get(), key, value, version);

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

// Get the value

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

ValueBuf val;

ASSERT_EQ(blob_get(txn.get(), key, &val), TxnErrorCode::TXN_OK);

ASSERT_EQ(val.ver, version);

// Verify the value

std::string retrieved_value = val.value();

ASSERT_EQ(retrieved_value, value);

}

// Test blob_get with non-existent key

TEST(BlobMessageTest, GetNonExistentKey) {

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

ASSERT_EQ(txn_kv->init(), 0);

std::unique_ptr<Transaction> txn;

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

std::string key = "test_blob_non_existent_key";

ValueBuf val;

ASSERT_EQ(blob_get(txn.get(), key, &val), TxnErrorCode::TXN_KEY_NOT_FOUND);

}

// Test blob_get_range with TxnKv

TEST(BlobMessageTest, GetRangeWithTxnKv) {

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

ASSERT_EQ(txn_kv->init(), 0);

std::vector<std::string> expected_keys;

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

std::unique_ptr<Transaction> txn;

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

std::string key = "test_blob_range_test_" + std::to_string(i);

expected_keys.push_back(key);

OperationLogPB operation_log = create_operation_log(100000, 10000);

operation_log.set_min_timestamp(i);

versioned::blob_put(txn.get(), key, operation_log, 2, 50000);

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

}

// Get range

std::string begin_key = "test_blob_range_test_";

std::string end_key = begin_key + "\xff";

auto iter = blob_get_range(txn_kv, begin_key, end_key);

ASSERT_NE(iter, nullptr);

std::unique_ptr<Transaction> txn;

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

// Iterate and verify

int count = 0;

std::vector<std::string> retrieved_keys;

while (iter->valid()) {

ASSERT_EQ(iter->version(), 2);

Versionstamp vs;

std::string_view key(iter->key());

ASSERT_TRUE(decode_versioned_key(&key, &vs));

retrieved_keys.push_back(std::string(key));

OperationLogPB retrieved_log;

ASSERT_TRUE(iter->parse_value(&retrieved_log));

ASSERT_EQ(retrieved_log.min_timestamp(), count);

for (auto&& raw_key : iter->raw_keys()) {

txn->remove(raw_key);

}

count++;

iter->next();

}

ASSERT_EQ(count, 5);

ASSERT_EQ(retrieved_keys, expected_keys);

ASSERT_EQ(iter->error_code(), TxnErrorCode::TXN_OK);

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

{

// The range should be empty now

auto iter2 = blob_get_range(txn_kv, begin_key, end_key);

ASSERT_NE(iter2, nullptr);

ASSERT_FALSE(iter2->valid());

ASSERT_EQ(iter2->error_code(), TxnErrorCode::TXN_OK);

}

}

// Test blob_get_range with empty range

TEST(BlobMessageTest, GetRangeEmpty) {

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

ASSERT_EQ(txn_kv->init(), 0);

std::string begin_key = "test_blob_empty_range_";

std::string end_key = begin_key + "\xff";

auto iter = blob_get_range(txn_kv, begin_key, end_key);

ASSERT_NE(iter, nullptr);

ASSERT_FALSE(iter->valid());

ASSERT_EQ(iter->error_code(), TxnErrorCode::TXN_OK);

}