Skip to content

Instantly share code, notes, and snippets.

@Kiterai
Created May 14, 2024 12:47
Show Gist options
  • Save Kiterai/d0696eca57dd2e002d9563c27ac65672 to your computer and use it in GitHub Desktop.
Save Kiterai/d0696eca57dd2e002d9563c27ac65672 to your computer and use it in GitHub Desktop.
test for mongodb gridfs with seek
#include <cassert>
#include <cstdint>
#include <iostream>
#include <random>
#include <vector>
#include <bsoncxx/builder/basic/document.hpp>
#include <bsoncxx/json.hpp>
#include <mongocxx/client.hpp>
#include <mongocxx/instance.hpp>
#include <mongocxx/stdx.hpp>
#include <mongocxx/uri.hpp>
using bsoncxx::builder::basic::kvp;
using bsoncxx::builder::basic::make_array;
using bsoncxx::builder::basic::make_document;
class gridfs_file {
mongocxx::gridfs::bucket bucket;
bsoncxx::types::bson_value::view file_id_hold;
mongocxx::gridfs::downloader downloader;
int64_t len, pos;
std::vector<uint8_t> buf;
public:
gridfs_file(const mongocxx::database &db, bsoncxx::types::bson_value::view file_id) {
bucket = db.gridfs_bucket();
file_id_hold = file_id;
downloader = bucket.open_download_stream(file_id_hold);
len = downloader.file_length();
pos = 0;
buf.resize(downloader.chunk_size());
}
auto read(uint8_t *buffer, size_t len) {
return pos += downloader.read(buffer, len);
}
void seek(const int64_t target_pos) {
assert(target_pos > 0);
if (target_pos < pos) {
downloader.close();
downloader = bucket.open_download_stream(file_id_hold);
pos = 0;
}
while (target_pos - pos > buf.size())
read(buf.data(), buf.size());
assert(target_pos - pos <= buf.size());
read(buf.data(), target_pos - pos);
}
};
int main() {
mongocxx::instance instance{};
// const auto host = safe_getenv("NS_MONGODB_HOSTNAME");
// const auto port = safe_getenv("NS_MONGODB_PORT");
// const auto user = safe_getenv("NS_MONGODB_USER");
// const auto pass = safe_getenv("NS_MONGODB_PASSWORD");
// const auto dbname = safe_getenv("NS_MONGODB_DATABASE");
const std::string host = "127.0.0.1";
const std::string port = "27017";
const std::string user = "root";
const std::string pass = "password";
const std::string dbname = "db";
mongocxx::uri uri{"mongodb://" + user + ":" + pass + "@" + host + ":" + port};
mongocxx::client client{uri};
auto db = client[dbname];
auto col = db["test"];
auto doc_value = make_document(
kvp("name", "MongoDB"),
kvp("type", "database"),
kvp("count", 1),
kvp("versions", make_array("v6.0", "v5.0", "v4.4", "v4.2", "v4.0", "v3.6")),
kvp("info", make_document(kvp("x", 203), kvp("y", 102))));
const auto res = col.insert_one(doc_value.view());
if (!res) {
std::cout << "err" << std::endl;
return -1;
}
const auto id = res->inserted_id();
std::cout << id.get_oid().value.to_string() << std::endl;
std::mt19937 rnd{std::random_device{}()};
std::vector<uint8_t> test_data;
test_data.resize(16 * 1024 * 1024);
// test_data.resize(1024);
for (auto &val : test_data) {
val = rnd();
}
auto bucket = db.gridfs_bucket();
bsoncxx::types::bson_value::view f_id;
{
auto uploader = bucket.open_upload_stream("test.txt");
uploader.write(test_data.data(), test_data.size());
const auto res = uploader.close();
f_id = res.id();
std::cout << f_id.get_oid().value.to_string() << std::endl;
}
{
auto downloader = gridfs_file(db, f_id);
constexpr int chk_size = 256;
for (int i = 0; i < 100; i++) {
auto pos = rnd();
pos %= (test_data.size() - chk_size);
assert(0 <= pos);
assert(pos < test_data.size() - chk_size);
downloader.seek(pos);
uint8_t buf[chk_size];
downloader.read(buf, chk_size);
bool ok = true;
for (int j = 0; j < chk_size; j++) {
if (test_data[pos + j] != buf[j])
ok = false;
}
if (!ok) {
std::cout << "error!" << std::endl;
for (int j = 0; j < 16; j++) {
std::cout << (int)test_data[pos + j] << ' ';
}
std::cout << std::endl;
for (int j = 0; j < 16; j++) {
std::cout << (int)buf[j] << ' ';
}
std::cout << std::endl;
} else {
std::cout << "test " << i + 1 << " passed" << std::endl;
}
}
}
return 0;
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment