Line data Source code
1 : #ifndef BIGQUERY_EMULATOR_BACKEND_ENGINE_DUCKDB_DUCKDB_EXECUTOR_TEST_FIXTURE_H_
2 : #define BIGQUERY_EMULATOR_BACKEND_ENGINE_DUCKDB_DUCKDB_EXECUTOR_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/status/status.h"
15 : #include "absl/status/statusor.h"
16 : #include "absl/strings/str_cat.h"
17 : #include "absl/strings/string_view.h"
18 : #include "backend/catalog/googlesql_catalog.h"
19 : #include "backend/engine/duckdb/duckdb_executor.h"
20 : #include "backend/engine/engine.h"
21 : #include "backend/schema/schema.h"
22 : #include "backend/storage/duckdb/duckdb_storage.h"
23 : #include "backend/storage/storage.h"
24 : #include "googlesql/public/analyzer.h"
25 : #include "googlesql/public/analyzer_options.h"
26 : #include "googlesql/public/analyzer_output.h"
27 : #include "googlesql/public/language_options.h"
28 : #include "googlesql/public/options.pb.h"
29 : #include "googlesql/public/types/type_factory.h"
30 : #include "googlesql/resolved_ast/resolved_ast.h"
31 : #include "gtest/gtest.h"
32 :
33 : namespace bigquery_emulator {
34 : namespace backend {
35 : namespace engine {
36 : namespace duckdb {
37 :
38 : namespace fs = std::filesystem;
39 :
40 12 : inline ::googlesql::LanguageOptions MakeLanguageOptions() {
41 12 : ::googlesql::LanguageOptions language;
42 12 : language.EnableMaximumLanguageFeatures();
43 12 : language.set_product_mode(::googlesql::PRODUCT_EXTERNAL);
44 12 : language.set_name_resolution_mode(::googlesql::NAME_RESOLUTION_DEFAULT);
45 12 : return language;
46 12 : }
47 :
48 6 : inline ::googlesql::AnalyzerOptions MakeAnalyzerOptions(bool all_statements) {
49 6 : ::googlesql::AnalyzerOptions options(MakeLanguageOptions());
50 6 : options.set_error_message_mode(::googlesql::ERROR_MESSAGE_ONE_LINE);
51 6 : options.set_attach_error_location_payload(true);
52 6 : options.CreateDefaultArenasIfNotSet();
53 6 : if (all_statements) {
54 2 : options.mutable_language()->SetSupportsAllStatementKinds();
55 2 : }
56 6 : return options;
57 6 : }
58 :
59 : class DuckDbExecutorTest : public ::testing::Test {
60 : protected:
61 6 : void SetUp() override {
62 6 : const char* tmpdir_env = std::getenv("TMPDIR");
63 6 : const std::string tmpdir = tmpdir_env != nullptr ? tmpdir_env : "/tmp";
64 6 : std::random_device rd;
65 6 : std::seed_seq seed{rd(), rd()};
66 6 : std::mt19937_64 rng(seed);
67 6 : data_dir_ =
68 6 : fs::path(tmpdir) / absl::StrCat("bqemu-duckdb-executor-test-", rng());
69 6 : std::error_code ec;
70 6 : fs::remove_all(data_dir_, ec);
71 6 : auto opened = storage::duckdb::DuckDBStorage::Open(data_dir_.string());
72 12 : ASSERT_TRUE(opened.ok()) << opened.status();
73 6 : storage_ = std::move(opened).value();
74 6 : executor_ = std::make_unique<DuckDbExecutor>(storage_.get());
75 6 : }
76 :
77 6 : void TearDown() override {
78 6 : executor_.reset();
79 6 : storage_.reset();
80 6 : std::error_code ec;
81 6 : fs::remove_all(data_dir_, ec);
82 6 : }
83 :
84 6 : QueryRequest MakeRequest(absl::string_view sql) {
85 6 : QueryRequest req;
86 6 : req.project_id = "proj-test";
87 6 : req.sql = std::string(sql);
88 6 : return req;
89 6 : }
90 :
91 1 : void CreatePeopleTable() {
92 1 : schema::TableSchema bq_schema;
93 1 : schema::ColumnSchema id;
94 1 : id.name = "id";
95 1 : id.type = schema::ColumnType::kInt64;
96 1 : id.mode = schema::ColumnMode::kRequired;
97 1 : bq_schema.columns.push_back(id);
98 1 : schema::ColumnSchema name;
99 1 : name.name = "name";
100 1 : name.type = schema::ColumnType::kString;
101 1 : name.mode = schema::ColumnMode::kNullable;
102 1 : bq_schema.columns.push_back(name);
103 1 : ASSERT_TRUE(storage_->CreateDataset({"proj-test", "ds"}, "US").ok());
104 1 : ASSERT_TRUE(
105 1 : storage_->CreateTable({"proj-test", "ds", "people"}, bq_schema).ok());
106 :
107 3 : auto make_row = [](int64_t id_val, std::string name_val) {
108 3 : storage::Row r;
109 3 : r.cells = {
110 3 : storage::Value::Int64(id_val),
111 3 : storage::Value::String(std::move(name_val)),
112 3 : };
113 3 : return r;
114 3 : };
115 1 : std::vector<storage::Row> rows = {
116 1 : make_row(1, "ada"),
117 1 : make_row(2, "linus"),
118 1 : make_row(3, "grace"),
119 1 : };
120 1 : ASSERT_TRUE(storage_
121 1 : ->AppendRows({"proj-test", "ds", "people"},
122 1 : absl::MakeConstSpan(rows))
123 1 : .ok());
124 1 : }
125 :
126 : struct CatalogBundle {
127 : std::unique_ptr<::googlesql::TypeFactory> type_factory{};
128 : std::unique_ptr<catalog::GoogleSqlCatalog> catalog{};
129 : };
130 6 : CatalogBundle MakeCatalog() {
131 6 : auto type_factory = std::make_unique<::googlesql::TypeFactory>();
132 6 : auto catalog = std::make_unique<catalog::GoogleSqlCatalog>(
133 6 : "proj-test", storage_.get(), type_factory.get(), MakeLanguageOptions());
134 6 : return {std::move(type_factory), std::move(catalog)};
135 6 : }
136 :
137 : absl::StatusOr<std::unique_ptr<const ::googlesql::AnalyzerOutput>> Analyze(
138 : absl::string_view sql,
139 : ::googlesql::Catalog* catalog,
140 6 : bool all_statements) {
141 6 : ::googlesql::AnalyzerOptions options = MakeAnalyzerOptions(all_statements);
142 6 : ::googlesql::TypeFactory type_factory;
143 6 : std::unique_ptr<const ::googlesql::AnalyzerOutput> output;
144 6 : absl::Status s = ::googlesql::AnalyzeStatement(
145 6 : sql, options, catalog, &type_factory, &output);
146 6 : if (!s.ok()) return s;
147 6 : return output;
148 6 : }
149 :
150 : fs::path data_dir_{};
151 : std::unique_ptr<storage::duckdb::DuckDBStorage> storage_{};
152 : std::unique_ptr<DuckDbExecutor> executor_{};
153 : };
154 :
155 : } // namespace duckdb
156 : } // namespace engine
157 : } // namespace backend
158 : } // namespace bigquery_emulator
159 :
160 : #endif // BIGQUERY_EMULATOR_BACKEND_ENGINE_DUCKDB_DUCKDB_EXECUTOR_TEST_FIXTURE_H_
|
