Line data    Source code

       1              : #ifndef BIGQUERY_EMULATOR_FRONTEND_HANDLERS_QUERY_TEST_FIXTURE_H_
       2              : #define BIGQUERY_EMULATOR_FRONTEND_HANDLERS_QUERY_TEST_FIXTURE_H_
       3              : 
       4              : #include <cstdint>
       5              : #include <cstdlib>
       6              : #include <filesystem>
       7              : #include <functional>
       8              : #include <memory>
       9              : #include <random>
      10              : #include <string>
      11              : #include <system_error>
      12              : #include <utility>
      13              : #include <vector>
      14              : 
      15              : #include "absl/strings/str_cat.h"
      16              : #include "absl/strings/string_view.h"
      17              : #include "absl/types/span.h"
      18              : #include "backend/engine/coordinator/local_coordinator_engine.h"
      19              : #include "backend/schema/schema.h"
      20              : #include "backend/storage/duckdb/duckdb_storage.h"
      21              : #include "backend/storage/storage.h"
      22              : #include "frontend/handlers/query.h"
      23              : #include "grpcpp/grpcpp.h"
      24              : #include "gtest/gtest.h"
      25              : #include "proto/emulator.pb.h"
      26              : 
      27              : namespace bigquery_emulator {
      28              : namespace frontend {
      29              : 
      30              : namespace fs = std::filesystem;
      31              : 
      32              : class QueryServiceTest : public ::testing::Test {
      33              :  protected:
      34           20 :   void SetUp() override {
      35           20 :     const char* tmpdir_env = std::getenv("TMPDIR");
      36           20 :     const std::string tmpdir = tmpdir_env != nullptr ? tmpdir_env : "/tmp";
      37           20 :     std::random_device rd;
      38           20 :     std::seed_seq seed{rd(), rd()};
      39           20 :     std::mt19937_64 rng(seed);
      40           20 :     data_dir_ = fs::path(tmpdir) / absl::StrCat("bqemu-query-test-", rng());
      41           20 :     std::error_code ec;
      42           20 :     fs::remove_all(data_dir_, ec);
      43           20 :     auto opened =
      44           20 :         backend::storage::duckdb::DuckDBStorage::Open(data_dir_.string());
      45           40 :     ASSERT_TRUE(opened.ok()) << opened.status();
      46           20 :     storage_ = std::move(opened).value();
      47           20 :     engine_ =
      48           20 :         std::make_unique<backend::engine::coordinator::LocalCoordinatorEngine>(
      49           20 :             storage_.get());
      50           20 :     service_ = std::make_unique<QueryService>(storage_.get(), engine_.get());
      51           20 :   }
      52              : 
      53           20 :   void TearDown() override {
      54           20 :     service_.reset();
      55           20 :     engine_.reset();
      56           20 :     storage_.reset();
      57           20 :     std::error_code ec;
      58           20 :     fs::remove_all(data_dir_, ec);
      59           20 :   }
      60              : 
      61           19 :   v1::QueryRequest MakeRequest(absl::string_view sql) {
      62           19 :     v1::QueryRequest req;
      63           19 :     req.set_project_id("proj-test");
      64           19 :     req.set_sql(std::string(sql));
      65           19 :     return req;
      66           19 :   }
      67              : 
      68            5 :   void CreatePeopleTable() {
      69            5 :     backend::schema::TableSchema schema;
      70            5 :     backend::schema::ColumnSchema id;
      71            5 :     id.name = "id";
      72            5 :     id.type = backend::schema::ColumnType::kInt64;
      73            5 :     id.mode = backend::schema::ColumnMode::kRequired;
      74            5 :     schema.columns.push_back(id);
      75            5 :     backend::schema::ColumnSchema name;
      76            5 :     name.name = "name";
      77            5 :     name.type = backend::schema::ColumnType::kString;
      78            5 :     name.mode = backend::schema::ColumnMode::kNullable;
      79            5 :     schema.columns.push_back(name);
      80            5 :     backend::schema::ColumnSchema tags;
      81            5 :     tags.name = "tags";
      82            5 :     tags.type = backend::schema::ColumnType::kString;
      83            5 :     tags.mode = backend::schema::ColumnMode::kRepeated;
      84            5 :     schema.columns.push_back(tags);
      85            5 :     ASSERT_TRUE(storage_->CreateDataset({"proj-test", "ds"}, "US").ok());
      86            5 :     ASSERT_TRUE(storage_->CreateTable({"proj-test", "ds", "t"}, schema).ok());
      87            5 :   }
      88              : 
      89              :   fs::path data_dir_{};
      90              :   std::unique_ptr<backend::storage::duckdb::DuckDBStorage> storage_{};
      91              :   std::unique_ptr<backend::engine::coordinator::LocalCoordinatorEngine>
      92              :       engine_{};
      93              :   std::unique_ptr<QueryService> service_{};
      94              : };
      95              : 
      96              : class MessageCollector {
      97              :  public:
      98           12 :   std::function<bool(const v1::QueryResultRow&)> Writer() {
      99           35 :     return [this](const v1::QueryResultRow& msg) {
     100           35 :       messages_.push_back(msg);
     101           35 :       return true;
     102           35 :     };
     103           12 :   }
     104           10 :   const std::vector<v1::QueryResultRow>& messages() const {
     105           10 :     return messages_;
     106           10 :   }
     107              : 
     108            6 :   const v1::QueryResultRow* PhaseTimingsTrailer() const {
     109           18 :     for (auto it = messages_.rbegin(); it != messages_.rend(); ++it) {
     110           18 :       if (it->has_phase_timings() && !it->phase_timings().phases().empty()) {
     111            6 :         return &(*it);
     112            6 :       }
     113           18 :     }
     114            0 :     return nullptr;
     115            6 :   }
     116              : 
     117            6 :   const v1::QueryResultRow* StatementTypeTrailer() const {
     118           12 :     for (auto it = messages_.rbegin(); it != messages_.rend(); ++it) {
     119           12 :       if (!it->statement_type().empty()) {
     120            6 :         return &(*it);
     121            6 :       }
     122           12 :     }
     123            0 :     return nullptr;
     124            6 :   }
     125              : 
     126            6 :   const v1::QueryResultRow* EmulatorRouteTrailer() const {
     127            6 :     for (auto it = messages_.rbegin(); it != messages_.rend(); ++it) {
     128            6 :       if (!it->emulator_route().empty()) {
     129            6 :         return &(*it);
     130            6 :       }
     131            6 :     }
     132            0 :     return nullptr;
     133            6 :   }
     134              : 
     135              :  private:
     136              :   std::vector<v1::QueryResultRow> messages_{};
     137              : };
     138              : 
     139              : }  // namespace frontend
     140              : }  // namespace bigquery_emulator
     141              : 
     142              : #endif  // BIGQUERY_EMULATOR_FRONTEND_HANDLERS_QUERY_TEST_FIXTURE_H_