Line data Source code
1 : #ifndef BIGQUERY_EMULATOR_BACKEND_ENGINE_DUCKDB_TRANSPILER_TRANSPILER_TEST_FIXTURE_H_
2 : #define BIGQUERY_EMULATOR_BACKEND_ENGINE_DUCKDB_TRANSPILER_TRANSPILER_TEST_FIXTURE_H_
3 :
4 : #include <memory>
5 : #include <string>
6 : #include <vector>
7 :
8 : #include "absl/status/status.h"
9 : #include "absl/status/statusor.h"
10 : #include "absl/strings/str_cat.h"
11 : #include "absl/strings/string_view.h"
12 : #include "backend/engine/disposition.h"
13 : #include "backend/engine/duckdb/transpiler/functions.h"
14 : #include "backend/engine/duckdb/transpiler/transpiler.h"
15 : #include "backend/engine/duckdb/udf/registrar.h"
16 : #include "duckdb.h"
17 : #include "googlesql/public/analyzer.h"
18 : #include "googlesql/public/analyzer_options.h"
19 : #include "googlesql/public/analyzer_output.h"
20 : #include "googlesql/public/builtin_function_options.h"
21 : #include "googlesql/public/catalog.h"
22 : #include "googlesql/public/id_string.h"
23 : #include "googlesql/public/language_options.h"
24 : #include "googlesql/public/options.pb.h"
25 : #include "googlesql/public/simple_catalog.h"
26 : #include "googlesql/public/types/type_factory.h"
27 : #include "googlesql/public/value.h"
28 : #include "googlesql/resolved_ast/resolved_ast.h"
29 : #include "googlesql/resolved_ast/resolved_column.h"
30 : #include "gtest/gtest.h"
31 :
32 : namespace bigquery_emulator {
33 : namespace backend {
34 : namespace engine {
35 : namespace duckdb {
36 : namespace transpiler {
37 :
38 : // Mirrors `duckdb_engine::MakeAnalyzerOptions` so the tests
39 : // resolve names through the same `LanguageOptions` snapshot the
40 : // engine itself uses. Drifting these two breaks function dispatch
41 : // (e.g. `IFNULL` resolves but `COALESCE` does not) in subtle ways
42 : // that only surface in the conformance harness.
43 169 : inline ::googlesql::AnalyzerOptions MakeAnalyzerOptions() {
44 169 : ::googlesql::LanguageOptions language;
45 169 : language.EnableMaximumLanguageFeatures();
46 169 : language.set_product_mode(::googlesql::PRODUCT_EXTERNAL);
47 169 : language.set_name_resolution_mode(::googlesql::NAME_RESOLUTION_DEFAULT);
48 : // Match the engine / route-classifier fixture so INSERT / CTAS bind
49 : // tests (R17 attribution materialization) can AnalyzeStatement.
50 169 : language.SetSupportsAllStatementKinds();
51 169 : ::googlesql::AnalyzerOptions options(language);
52 169 : options.set_error_message_mode(::googlesql::ERROR_MESSAGE_ONE_LINE);
53 : // Match the engine: keep PIVOT / UNPIVOT in their raw resolved-AST
54 : // forms so the transpiler `EmitPivotScan` / `EmitUnpivotScan`
55 : // emit paths are exercised. The engine itself disables these
56 : // rewriters (see `local_coordinator_engine.cc::MakeAnalyzerOptions`)
57 : // because the disposition table routes the raw nodes through
58 : // `duckdb_rewrite`.
59 169 : options.disable_rewrite(::googlesql::REWRITE_PIVOT);
60 169 : options.disable_rewrite(::googlesql::REWRITE_UNPIVOT);
61 169 : options.CreateDefaultArenasIfNotSet();
62 169 : return options;
63 169 : }
64 :
65 : // Helper: synthesize a `ResolvedWithExpr` directly so tests do not
66 : // depend on the analyzer preserving a `WITH(...)` expression against
67 : // constant-folding / inlining heuristics.
68 : struct TestWithExprBinding {
69 : std::string name;
70 : std::unique_ptr<const ::googlesql::ResolvedExpr> expr;
71 : };
72 :
73 : inline std::unique_ptr<::googlesql::ResolvedWithExpr> MakeTestWithExpr(
74 3 : std::vector<TestWithExprBinding> bindings) {
75 3 : if (bindings.empty()) return nullptr;
76 3 : std::vector<std::unique_ptr<const ::googlesql::ResolvedComputedColumn>>
77 3 : assignments;
78 3 : std::vector<::googlesql::ResolvedColumn> columns;
79 3 : int next_id = 1;
80 4 : for (auto& binding : bindings) {
81 4 : if (binding.expr == nullptr) return nullptr;
82 4 : const ::googlesql::Type* t = binding.expr->type();
83 4 : ::googlesql::ResolvedColumn col(
84 4 : next_id++,
85 4 : /*table_name=*/::googlesql::IdString::MakeGlobal("$with"),
86 4 : /*name=*/::googlesql::IdString::MakeGlobal(binding.name),
87 4 : t);
88 4 : columns.push_back(col);
89 4 : auto cc =
90 4 : ::googlesql::MakeResolvedComputedColumn(col, std::move(binding.expr));
91 4 : assignments.push_back(std::move(cc));
92 4 : }
93 3 : std::unique_ptr<const ::googlesql::ResolvedExpr> body =
94 3 : ::googlesql::MakeResolvedColumnRef(columns.front(),
95 3 : /*is_correlated=*/false);
96 3 : return ::googlesql::MakeResolvedWithExpr(
97 3 : columns.front().type(), std::move(assignments), std::move(body));
98 3 : }
99 :
100 : // One-stop test fixture. Owns the type factory, catalog, and a
101 : // people table; every test gets a fresh `Transpiler` so the
102 : // per-traversal accumulator (when one lands) starts clean.
103 : class TranspilerTest : public ::testing::Test {
104 : protected:
105 165 : void SetUp() override {
106 165 : type_factory_ = std::make_unique<::googlesql::TypeFactory>();
107 165 : catalog_ = std::make_unique<::googlesql::SimpleCatalog>(
108 165 : "test_catalog", type_factory_.get());
109 165 : ::googlesql::LanguageOptions language;
110 165 : language.EnableMaximumLanguageFeatures();
111 165 : language.set_product_mode(::googlesql::PRODUCT_EXTERNAL);
112 165 : ASSERT_TRUE(catalog_
113 165 : ->AddBuiltinFunctionsAndTypes(
114 165 : ::googlesql::BuiltinFunctionOptions(language))
115 165 : .ok());
116 :
117 165 : auto people = std::make_unique<::googlesql::SimpleTable>(
118 165 : "people",
119 165 : std::vector<::googlesql::SimpleTable::NameAndType>{
120 165 : {"id", type_factory_->get_int64()},
121 165 : {"name", type_factory_->get_string()},
122 165 : });
123 165 : catalog_->AddOwnedTable(std::move(people));
124 :
125 165 : const ::googlesql::Type* int64_array_type = nullptr;
126 165 : EXPECT_TRUE(
127 165 : type_factory_
128 165 : ->MakeArrayType(type_factory_->get_int64(), &int64_array_type)
129 165 : .ok());
130 165 : auto arr_table = std::make_unique<::googlesql::SimpleTable>(
131 165 : "arr_table",
132 165 : std::vector<::googlesql::SimpleTable::NameAndType>{
133 165 : {"id", type_factory_->get_int64()},
134 165 : {"arr", int64_array_type},
135 165 : });
136 165 : catalog_->AddOwnedTable(std::move(arr_table));
137 :
138 : // The join tests need a second table with disjoint column names so
139 : // the analyzer doesn't have to disambiguate references in the ON
140 : // expression; the transpiler doesn't know how to disambiguate yet
141 : // (the per-column emit goes through `ResolvedColumn::name()`).
142 165 : auto orders = std::make_unique<::googlesql::SimpleTable>(
143 165 : "orders",
144 165 : std::vector<::googlesql::SimpleTable::NameAndType>{
145 165 : {"order_id", type_factory_->get_int64()},
146 165 : {"amount", type_factory_->get_int64()},
147 165 : });
148 165 : catalog_->AddOwnedTable(std::move(orders));
149 :
150 : // A table with a string discriminator + numeric value column so the
151 : // PIVOT / UNPIVOT tests have something the analyzer accepts for
152 : // `FOR <expr> IN (<literals>)` (PIVOT) and
153 : // `UNPIVOT(<value_cols> FOR <label_col> IN (<col_groups>))`
154 : // (UNPIVOT).
155 165 : auto sales = std::make_unique<::googlesql::SimpleTable>(
156 165 : "sales",
157 165 : std::vector<::googlesql::SimpleTable::NameAndType>{
158 165 : {"region", type_factory_->get_string()},
159 165 : {"kind", type_factory_->get_string()},
160 165 : {"amount", type_factory_->get_int64()},
161 165 : });
162 165 : catalog_->AddOwnedTable(std::move(sales));
163 :
164 : // Wide table for UNPIVOT: each column is one of the unpivot
165 : // arguments the analyzer threads through `unpivot_arg_list`.
166 165 : auto wide = std::make_unique<::googlesql::SimpleTable>(
167 165 : "wide",
168 165 : std::vector<::googlesql::SimpleTable::NameAndType>{
169 165 : {"region", type_factory_->get_string()},
170 165 : {"q1", type_factory_->get_int64()},
171 165 : {"q2", type_factory_->get_int64()},
172 165 : });
173 165 : catalog_->AddOwnedTable(std::move(wide));
174 :
175 165 : auto org = std::make_unique<::googlesql::SimpleTable>(
176 165 : "org",
177 165 : std::vector<::googlesql::SimpleTable::NameAndType>{
178 165 : {"employee", type_factory_->get_string()},
179 165 : {"manager", type_factory_->get_string()},
180 165 : });
181 165 : catalog_->AddOwnedTable(std::move(org));
182 :
183 165 : auto transactions = std::make_unique<::googlesql::SimpleTable>(
184 165 : "transactions",
185 165 : std::vector<::googlesql::SimpleTable::NameAndType>{
186 165 : {"timestamp", type_factory_->get_timestamp()},
187 165 : {"origin", type_factory_->get_string()},
188 165 : {"destination", type_factory_->get_string()},
189 165 : {"amount", type_factory_->get_numeric()},
190 165 : });
191 165 : catalog_->AddOwnedTable(std::move(transactions));
192 :
193 165 : transpiler_ = std::make_unique<Transpiler>();
194 165 : }
195 :
196 : // Analyze `sql` against the fixture catalog and return the
197 : // resolved AST. The `AnalyzerOutput` lives in `last_output_` so
198 : // the `ResolvedStatement` (and the `Type*` / `Function*` pointers
199 : // it references) stays alive for the duration of the test.
200 163 : const ::googlesql::ResolvedStatement* Analyze(absl::string_view sql) {
201 163 : ::googlesql::AnalyzerOptions options = MakeAnalyzerOptions();
202 163 : return AnalyzeWith(sql, options);
203 163 : }
204 :
205 : // Analyze `sql` with `options` already configured -- handy for the
206 : // parameter-emit tests that need `AddQueryParameter` /
207 : // `AddPositionalQueryParameter` calls before analysis. Same
208 : // ownership contract as `Analyze`: the resolved AST lives in
209 : // `last_output_` for the duration of the test.
210 : const ::googlesql::ResolvedStatement* AnalyzeWith(
211 169 : absl::string_view sql, const ::googlesql::AnalyzerOptions& options) {
212 169 : last_output_.reset();
213 169 : absl::Status s = ::googlesql::AnalyzeStatement(
214 169 : sql, options, catalog_.get(), type_factory_.get(), &last_output_);
215 338 : EXPECT_TRUE(s.ok()) << s;
216 169 : if (!s.ok() || last_output_ == nullptr) return nullptr;
217 169 : return last_output_->resolved_statement();
218 169 : }
219 :
220 : // Convenience: pluck the inner `ResolvedScan` out of a
221 : // `SELECT ... FROM ...` statement. We unwrap the ResolvedQueryStmt
222 : // (and the ResolvedProjectScan the analyzer wraps around any
223 : // explicit SELECT list) so the per-shape `Emit*` assertion below
224 : // sees the exact subtree it covers.
225 : const ::googlesql::ResolvedScan* QueryInputScan(
226 51 : const ::googlesql::ResolvedStatement* stmt) {
227 51 : EXPECT_NE(stmt, nullptr);
228 51 : if (stmt == nullptr) return nullptr;
229 51 : const auto* q = stmt->GetAs<::googlesql::ResolvedQueryStmt>();
230 51 : EXPECT_NE(q, nullptr);
231 51 : if (q == nullptr) return nullptr;
232 51 : const ::googlesql::ResolvedScan* scan = q->query();
233 88 : while (scan != nullptr &&
234 88 : scan->node_kind() == ::googlesql::RESOLVED_PROJECT_SCAN) {
235 37 : scan = scan->GetAs<::googlesql::ResolvedProjectScan>()->input_scan();
236 37 : }
237 51 : return scan;
238 51 : }
239 :
240 : // Walk down to the first ResolvedExpr we can find inside a SELECT
241 : // list -- handy for testing literal / function / column-ref emit
242 : // without having to also implement EmitProjectScan.
243 : const ::googlesql::ResolvedExpr* QueryFirstSelectExpr(
244 19 : const ::googlesql::ResolvedStatement* stmt) {
245 19 : EXPECT_NE(stmt, nullptr);
246 19 : if (stmt == nullptr) return nullptr;
247 19 : const auto* q = stmt->GetAs<::googlesql::ResolvedQueryStmt>();
248 19 : if (q == nullptr || q->query() == nullptr) return nullptr;
249 19 : const ::googlesql::ResolvedScan* scan = q->query();
250 19 : if (scan->node_kind() != ::googlesql::RESOLVED_PROJECT_SCAN) return nullptr;
251 19 : const auto* project = scan->GetAs<::googlesql::ResolvedProjectScan>();
252 19 : if (project->expr_list_size() == 0) return nullptr;
253 19 : return project->expr_list(0)->expr();
254 19 : }
255 :
256 : std::unique_ptr<::googlesql::TypeFactory> type_factory_{};
257 : std::unique_ptr<::googlesql::SimpleCatalog> catalog_{};
258 : std::unique_ptr<const ::googlesql::AnalyzerOutput> last_output_{};
259 : std::unique_ptr<Transpiler> transpiler_{};
260 : };
261 :
262 : // Subclass that publishes the protected `Emit*` family so the tests
263 : // can assert on individual emits without having to drive a full
264 : // query through `Transpile`. The class doesn't override anything --
265 : // it just widens the visibility.
266 : class TestTranspiler : public Transpiler {
267 : public:
268 : using Transpiler::EmitAggregateScan;
269 : using Transpiler::EmitAnalyticScan;
270 : using Transpiler::EmitArrayScan;
271 : using Transpiler::EmitCast;
272 : using Transpiler::EmitColumnRef;
273 : using Transpiler::EmitComputedColumn;
274 : using Transpiler::EmitFilterScan;
275 : using Transpiler::EmitFunctionArgument;
276 : using Transpiler::EmitFunctionCall;
277 : using Transpiler::EmitGetJsonField;
278 : using Transpiler::EmitGetStructField;
279 : using Transpiler::EmitJoinScan;
280 : using Transpiler::EmitLimitOffsetScan;
281 : using Transpiler::EmitLiteral;
282 : using Transpiler::EmitMakeStruct;
283 : using Transpiler::EmitOrderByScan;
284 : using Transpiler::EmitOutputColumn;
285 : using Transpiler::EmitParameter;
286 : using Transpiler::EmitPivotScan;
287 : using Transpiler::EmitProjectScan;
288 : using Transpiler::EmitQueryStmt;
289 : using Transpiler::EmitRecursiveRefScan;
290 : using Transpiler::EmitRecursiveScan;
291 : using Transpiler::EmitSampleScan;
292 : using Transpiler::EmitSetOperationScan;
293 : using Transpiler::EmitSingleRowScan;
294 : using Transpiler::EmitSubqueryExpr;
295 : using Transpiler::EmitTableScan;
296 : using Transpiler::EmitUnpivotScan;
297 : using Transpiler::EmitWithExpr;
298 : using Transpiler::EmitWithRefScan;
299 : using Transpiler::EmitWithScan;
300 : };
301 :
302 : // DuckDB-backed binding checker for composition / property tests. Opens an
303 : // in-memory connection, registers polyfill UDFs, and asserts transpiled SQL
304 : // binds (via duckdb_query, which runs parse + bind + plan).
305 : class TranspilerBindFixture : public TranspilerTest {
306 : protected:
307 16 : void SetUp() override {
308 16 : TranspilerTest::SetUp();
309 16 : ASSERT_EQ(::duckdb_open(nullptr, &db_), ::DuckDBSuccess);
310 16 : ASSERT_EQ(::duckdb_connect(db_, &conn_), ::DuckDBSuccess);
311 16 : absl::Status reg = udf::RegisterAll(conn_);
312 32 : ASSERT_TRUE(reg.ok()) << reg;
313 16 : }
314 :
315 16 : void TearDown() override {
316 16 : if (conn_ != nullptr) ::duckdb_disconnect(&conn_);
317 16 : if (db_ != nullptr) ::duckdb_close(&db_);
318 16 : conn_ = nullptr;
319 16 : db_ = nullptr;
320 16 : TranspilerTest::TearDown();
321 16 : }
322 :
323 112 : void ExecDdl(absl::string_view sql) {
324 112 : ::duckdb_result result;
325 224 : ASSERT_EQ(::duckdb_query(conn_, std::string(sql).c_str(), &result),
326 224 : ::DuckDBSuccess)
327 224 : << ::duckdb_result_error(&result);
328 112 : ::duckdb_destroy_result(&result);
329 112 : }
330 :
331 : void AssertTranspileBinds(const ::googlesql::ResolvedStatement* stmt,
332 : absl::string_view source_sql,
333 47 : TestTranspiler* t) {
334 94 : ASSERT_NE(stmt, nullptr) << "analyze failed for:\n" << source_sql;
335 47 : std::string emitted = t->Transpile(stmt);
336 94 : ASSERT_FALSE(emitted.empty()) << "transpiler returned empty SQL for:\n"
337 94 : << source_sql;
338 47 : SCOPED_TRACE(emitted);
339 47 : ::duckdb_result result{};
340 47 : const auto rc = ::duckdb_query(conn_, emitted.c_str(), &result);
341 47 : if (rc != ::DuckDBSuccess) {
342 0 : const char* err = ::duckdb_result_error(&result);
343 0 : FAIL() << "DuckDB rejected transpiled SQL\n"
344 0 : << "source_sql:\n"
345 0 : << source_sql << "\n"
346 0 : << "emitted_sql:\n"
347 0 : << emitted << "\n"
348 0 : << "duckdb_error:\n"
349 0 : << (err == nullptr ? "(null)" : err);
350 0 : }
351 47 : ::duckdb_destroy_result(&result);
352 47 : }
353 :
354 47 : void AssertSqlTranspileBinds(absl::string_view sql) {
355 47 : const ::googlesql::ResolvedStatement* stmt = Analyze(sql);
356 47 : TestTranspiler t;
357 47 : AssertTranspileBinds(stmt, sql, &t);
358 47 : }
359 :
360 : ::duckdb_database db_ = nullptr;
361 : ::duckdb_connection conn_ = nullptr;
362 : };
363 :
364 : } // namespace transpiler
365 : } // namespace duckdb
366 : } // namespace engine
367 : } // namespace backend
368 : } // namespace bigquery_emulator
369 :
370 : #endif // BIGQUERY_EMULATOR_BACKEND_ENGINE_DUCKDB_TRANSPILER_TRANSPILER_TEST_FIXTURE_H_
|
