Line data Source code
1 : // Deferred control-op shape tests for the executor. Kept in a separate
2 : // file so `control_op_executor_test.cc` stays under the cpp-lint
3 : // file-length cap.
4 :
5 : #include <cstdint>
6 : #include <cstdlib>
7 : #include <filesystem>
8 : #include <fstream>
9 : #include <memory>
10 : #include <random>
11 : #include <string>
12 : #include <system_error>
13 : #include <utility>
14 : #include <vector>
15 :
16 : #include "absl/status/status.h"
17 : #include "absl/strings/str_cat.h"
18 : #include "absl/strings/string_view.h"
19 : #include "backend/catalog/googlesql_catalog.h"
20 : #include "backend/engine/control/control_op_executor.h"
21 : #include "backend/engine/control/control_op_internal.h"
22 : #include "backend/engine/engine.h"
23 : #include "backend/schema/schema.h"
24 : #include "backend/storage/duckdb/duckdb_storage.h"
25 : #include "backend/storage/storage.h"
26 : #include "googlesql/public/analyzer.h"
27 : #include "googlesql/public/analyzer_options.h"
28 : #include "googlesql/public/analyzer_output.h"
29 : #include "googlesql/public/language_options.h"
30 : #include "googlesql/public/options.pb.h"
31 : #include "googlesql/public/types/type_factory.h"
32 : #include "googlesql/resolved_ast/resolved_ast.h"
33 : #include "gtest/gtest.h"
34 :
35 : namespace bigquery_emulator {
36 : namespace backend {
37 : namespace engine {
38 : namespace control {
39 : namespace {
40 :
41 : namespace fs = std::filesystem;
42 :
43 8 : ::googlesql::LanguageOptions MakeLanguageOptions() {
44 8 : ::googlesql::LanguageOptions language;
45 8 : language.EnableMaximumLanguageFeatures();
46 8 : language.set_product_mode(::googlesql::PRODUCT_EXTERNAL);
47 8 : language.set_name_resolution_mode(::googlesql::NAME_RESOLUTION_DEFAULT);
48 8 : return language;
49 8 : }
50 :
51 4 : ::googlesql::AnalyzerOptions MakeAnalyzerOptions() {
52 4 : ::googlesql::AnalyzerOptions options(MakeLanguageOptions());
53 4 : options.set_error_message_mode(::googlesql::ERROR_MESSAGE_ONE_LINE);
54 4 : options.set_attach_error_location_payload(true);
55 4 : options.CreateDefaultArenasIfNotSet();
56 4 : options.mutable_language()->SetSupportsAllStatementKinds();
57 4 : return options;
58 4 : }
59 :
60 : class ControlOpExecutorDeferredTest : public ::testing::Test {
61 : protected:
62 4 : void SetUp() override {
63 4 : const char* tmpdir_env = std::getenv("TMPDIR");
64 4 : const std::string tmpdir = tmpdir_env != nullptr ? tmpdir_env : "/tmp";
65 4 : std::random_device rd;
66 4 : std::seed_seq seed{rd(), rd()};
67 4 : std::mt19937_64 rng(seed);
68 4 : data_dir_ = fs::path(tmpdir) /
69 4 : absl::StrCat("bqemu-control-op-deferred-test-", rng());
70 4 : std::error_code ec;
71 4 : fs::remove_all(data_dir_, ec);
72 4 : auto opened = storage::duckdb::DuckDBStorage::Open(data_dir_.string());
73 8 : ASSERT_TRUE(opened.ok()) << opened.status();
74 4 : storage_ = std::move(opened).value();
75 4 : executor_ = std::make_unique<ControlOpExecutor>(storage_.get());
76 4 : ASSERT_TRUE(storage_->CreateDataset({"proj-test", "ds"}, "US").ok());
77 4 : }
78 :
79 4 : void TearDown() override {
80 4 : executor_.reset();
81 4 : storage_.reset();
82 4 : std::error_code ec;
83 4 : fs::remove_all(data_dir_, ec);
84 4 : }
85 :
86 3 : QueryRequest MakeRequest(absl::string_view sql) {
87 3 : QueryRequest req;
88 3 : req.project_id = "proj-test";
89 3 : req.sql = std::string(sql);
90 3 : return req;
91 3 : }
92 :
93 3 : void CreatePeopleTable() {
94 3 : schema::TableSchema bq_schema;
95 3 : schema::ColumnSchema id;
96 3 : id.name = "id";
97 3 : id.type = schema::ColumnType::kInt64;
98 3 : id.mode = schema::ColumnMode::kRequired;
99 3 : bq_schema.columns.push_back(id);
100 3 : schema::ColumnSchema name;
101 3 : name.name = "name";
102 3 : name.type = schema::ColumnType::kString;
103 3 : name.mode = schema::ColumnMode::kNullable;
104 3 : bq_schema.columns.push_back(name);
105 3 : ASSERT_TRUE(
106 3 : storage_->CreateTable({"proj-test", "ds", "people"}, bq_schema).ok());
107 :
108 9 : auto make_row = [](int64_t id_val, std::string name_val) {
109 9 : storage::Row r;
110 9 : r.cells = {
111 9 : storage::Value::Int64(id_val),
112 9 : storage::Value::String(std::move(name_val)),
113 9 : };
114 9 : return r;
115 9 : };
116 3 : std::vector<storage::Row> rows = {
117 3 : make_row(1, "ada"),
118 3 : make_row(2, "linus"),
119 3 : make_row(3, "grace"),
120 3 : };
121 3 : ASSERT_TRUE(storage_
122 3 : ->AppendRows({"proj-test", "ds", "people"},
123 3 : absl::MakeConstSpan(rows))
124 3 : .ok());
125 3 : }
126 :
127 : struct CatalogBundle {
128 : std::unique_ptr<::googlesql::TypeFactory> type_factory{};
129 : std::unique_ptr<catalog::GoogleSqlCatalog> catalog{};
130 : };
131 4 : CatalogBundle MakeCatalog() {
132 4 : auto type_factory = std::make_unique<::googlesql::TypeFactory>();
133 4 : auto catalog = std::make_unique<catalog::GoogleSqlCatalog>(
134 4 : "proj-test", storage_.get(), type_factory.get(), MakeLanguageOptions());
135 4 : return {std::move(type_factory), std::move(catalog)};
136 4 : }
137 :
138 3 : absl::Status RunDdl(absl::string_view sql) {
139 3 : CatalogBundle bundle = MakeCatalog();
140 3 : ::googlesql::AnalyzerOptions options = MakeAnalyzerOptions();
141 3 : ::googlesql::TypeFactory type_factory;
142 3 : std::unique_ptr<const ::googlesql::AnalyzerOutput> output;
143 3 : absl::Status analyze = ::googlesql::AnalyzeStatement(
144 3 : sql, options, bundle.catalog.get(), &type_factory, &output);
145 3 : if (!analyze.ok()) return analyze;
146 3 : if (output == nullptr || output->resolved_statement() == nullptr) {
147 0 : return absl::InternalError(
148 0 : "ControlOpExecutorDeferredTest::RunDdl: analyzer produced no "
149 0 : "resolved statement");
150 0 : }
151 3 : return executor_->ExecuteDdl(
152 3 : MakeRequest(sql), *output->resolved_statement(), bundle.catalog.get());
153 3 : }
154 :
155 : fs::path data_dir_{};
156 : std::unique_ptr<storage::duckdb::DuckDBStorage> storage_{};
157 : std::unique_ptr<ControlOpExecutor> executor_{};
158 : };
159 :
160 1 : TEST_F(ControlOpExecutorDeferredTest, CreateViewRegisteredByCoordinator) {
161 1 : absl::Status s = RunDdl("CREATE VIEW ds.v AS SELECT 1 AS id");
162 2 : EXPECT_TRUE(s.ok()) << s;
163 1 : }
164 :
165 1 : TEST_F(ControlOpExecutorDeferredTest, CreateMaterializedViewMaterializesRows) {
166 1 : CreatePeopleTable();
167 1 : absl::Status s = RunDdl(
168 1 : "CREATE MATERIALIZED VIEW ds.mv AS SELECT id, name FROM ds.people");
169 2 : ASSERT_TRUE(s.ok()) << s;
170 1 : auto scan = storage_->ScanRows({"proj-test", "ds", "mv"});
171 2 : ASSERT_TRUE(scan.ok()) << scan.status();
172 1 : int rows = 0;
173 1 : storage::Row row;
174 4 : while (true) {
175 4 : auto has = (*scan)->Next(&row);
176 8 : ASSERT_TRUE(has.ok()) << has.status();
177 4 : if (!*has) break;
178 3 : ASSERT_EQ(row.cells.size(), 2u);
179 3 : ++rows;
180 3 : }
181 1 : EXPECT_EQ(rows, 3);
182 1 : }
183 :
184 1 : TEST_F(ControlOpExecutorDeferredTest, ExportDataWritesLocalCsv) {
185 1 : CreatePeopleTable();
186 1 : const std::string path = absl::StrCat(std::getenv("TEST_TMPDIR") != nullptr
187 1 : ? std::getenv("TEST_TMPDIR")
188 1 : : "/tmp",
189 1 : "/bqemu_export_test.csv");
190 1 : absl::Status s =
191 1 : RunDdl(absl::StrCat("EXPORT DATA OPTIONS(uri='file://",
192 1 : path,
193 1 : "', format='CSV') AS ",
194 1 : "SELECT id, name FROM ds.people ORDER BY id"));
195 2 : ASSERT_TRUE(s.ok()) << s;
196 1 : std::ifstream in(path);
197 2 : ASSERT_TRUE(in.good()) << "expected export file at " << path;
198 1 : std::string line;
199 1 : ASSERT_TRUE(std::getline(in, line));
200 1 : EXPECT_EQ(line, "id,name");
201 1 : }
202 :
203 1 : TEST_F(ControlOpExecutorDeferredTest, TruncateTableClearsRows) {
204 1 : CreatePeopleTable();
205 1 : CatalogBundle bundle = MakeCatalog();
206 1 : ::googlesql::AnalyzerOptions options = MakeAnalyzerOptions();
207 1 : ::googlesql::TypeFactory type_factory;
208 1 : std::unique_ptr<const ::googlesql::AnalyzerOutput> output;
209 1 : absl::Status analyze =
210 1 : ::googlesql::AnalyzeStatement("TRUNCATE TABLE ds.people",
211 1 : options,
212 1 : bundle.catalog.get(),
213 1 : &type_factory,
214 1 : &output);
215 2 : ASSERT_TRUE(analyze.ok()) << analyze;
216 1 : ASSERT_NE(output, nullptr);
217 1 : ASSERT_NE(output->resolved_statement(), nullptr);
218 1 : const auto* truncate =
219 1 : output->resolved_statement()->GetAs<::googlesql::ResolvedTruncateStmt>();
220 1 : ASSERT_NE(truncate, nullptr);
221 1 : absl::StatusOr<int64_t> deleted =
222 1 : internal::RunTruncateTable(*storage_, truncate);
223 2 : ASSERT_TRUE(deleted.ok()) << deleted.status();
224 1 : EXPECT_EQ(*deleted, 3);
225 1 : absl::StatusOr<std::int64_t> after =
226 1 : storage_->CountRows({"proj-test", "ds", "people"});
227 2 : ASSERT_TRUE(after.ok()) << after.status();
228 1 : EXPECT_EQ(*after, 0);
229 1 : auto schema = storage_->GetSchema({"proj-test", "ds", "people"});
230 2 : ASSERT_TRUE(schema.ok()) << schema.status();
231 1 : EXPECT_EQ(schema->columns.size(), 2u);
232 1 : }
233 :
234 : } // namespace
235 : } // namespace control
236 : } // namespace engine
237 : } // namespace backend
238 : } // namespace bigquery_emulator
|
