Line data    Source code

       1              : #ifndef BIGQUERY_EMULATOR_FRONTEND_HANDLERS_STORAGE_READ_TEST_FIXTURE_H_
       2              : #define BIGQUERY_EMULATOR_FRONTEND_HANDLERS_STORAGE_READ_TEST_FIXTURE_H_
       3              : 
       4              : #include <cstdint>
       5              : #include <cstdlib>
       6              : #include <filesystem>
       7              : #include <memory>
       8              : #include <random>
       9              : #include <string>
      10              : #include <system_error>
      11              : #include <utility>
      12              : #include <vector>
      13              : 
      14              : #include "absl/strings/str_cat.h"
      15              : #include "backend/schema/schema.h"
      16              : #include "backend/storage/duckdb/duckdb_storage.h"
      17              : #include "backend/storage/storage.h"
      18              : #include "frontend/handlers/storage_read.h"
      19              : #include "gtest/gtest.h"
      20              : #include "proto/storage_read.pb.h"
      21              : 
      22              : namespace bigquery_emulator {
      23              : namespace frontend {
      24              : 
      25              : namespace fs = std::filesystem;
      26              : 
      27           27 : inline fs::path MakeTempDataDir(absl::string_view prefix) {
      28           27 :   const char* tmpdir_env = std::getenv("TMPDIR");
      29           27 :   const std::string tmpdir = tmpdir_env != nullptr ? tmpdir_env : "/tmp";
      30           27 :   std::random_device rd;
      31           27 :   std::seed_seq seed{rd(), rd()};
      32           27 :   std::mt19937_64 rng(seed);
      33           27 :   fs::path out = fs::path(tmpdir) /
      34           27 :                  absl::StrCat("bqemu-", std::string(prefix), "-", rng());
      35           27 :   std::error_code ec;
      36           27 :   fs::remove_all(out, ec);
      37           27 :   return out;
      38           27 : }
      39              : 
      40              : class StorageReadServiceTest : public ::testing::Test {
      41              :  protected:
      42           17 :   void SetUp() override {
      43           17 :     data_dir_ = MakeTempDataDir("storage-read-test");
      44           17 :     auto opened =
      45           17 :         backend::storage::duckdb::DuckDBStorage::Open(data_dir_.string());
      46           34 :     ASSERT_TRUE(opened.ok()) << opened.status();
      47           17 :     storage_ = std::move(opened).value();
      48           17 :     service_ = std::make_unique<StorageReadService>(storage_.get());
      49           17 :   }
      50              : 
      51           17 :   void TearDown() override {
      52           17 :     service_.reset();
      53           17 :     storage_.reset();
      54           17 :     std::error_code ec;
      55           17 :     fs::remove_all(data_dir_, ec);
      56           17 :   }
      57              : 
      58           12 :   void CreatePeopleTable() {
      59           12 :     backend::schema::TableSchema schema;
      60           12 :     backend::schema::ColumnSchema id;
      61           12 :     id.name = "id";
      62           12 :     id.type = backend::schema::ColumnType::kInt64;
      63           12 :     id.mode = backend::schema::ColumnMode::kRequired;
      64           12 :     schema.columns.push_back(id);
      65           12 :     backend::schema::ColumnSchema name;
      66           12 :     name.name = "name";
      67           12 :     name.type = backend::schema::ColumnType::kString;
      68           12 :     name.mode = backend::schema::ColumnMode::kNullable;
      69           12 :     schema.columns.push_back(name);
      70           12 :     backend::schema::ColumnSchema tags;
      71           12 :     tags.name = "tags";
      72           12 :     tags.type = backend::schema::ColumnType::kString;
      73           12 :     tags.mode = backend::schema::ColumnMode::kRepeated;
      74           12 :     schema.columns.push_back(tags);
      75           12 :     ASSERT_TRUE(storage_->CreateDataset({"proj-test", "ds"}, "US").ok());
      76           12 :     ASSERT_TRUE(storage_->CreateTable({"proj-test", "ds", "t"}, schema).ok());
      77           12 :   }
      78              : 
      79           14 :   v1::CreateReadSessionRequest MakePeopleRequest() {
      80           14 :     v1::CreateReadSessionRequest req;
      81           14 :     req.set_parent("projects/proj-test");
      82           14 :     req.mutable_read_session()->set_table(
      83           14 :         "projects/proj-test/datasets/ds/tables/t");
      84           14 :     return req;
      85           14 :   }
      86              : 
      87              :   fs::path data_dir_{};
      88              :   std::unique_ptr<backend::storage::duckdb::DuckDBStorage> storage_{};
      89              :   std::unique_ptr<StorageReadService> service_{};
      90              : };
      91              : 
      92              : }  // namespace frontend
      93              : }  // namespace bigquery_emulator
      94              : 
      95              : #endif  // BIGQUERY_EMULATOR_FRONTEND_HANDLERS_STORAGE_READ_TEST_FIXTURE_H_