Line data Source code
1 : // Positional STRUCT cast regression tests for the DuckDB executor.
2 : // Kept in a separate file so `duckdb_executor_test.cc` stays under the
3 : // cpp-lint file-length cap.
4 :
5 : #include <cstdlib>
6 : #include <filesystem>
7 : #include <memory>
8 : #include <random>
9 : #include <string>
10 : #include <system_error>
11 : #include <utility>
12 :
13 : #include "absl/status/status.h"
14 : #include "absl/status/statusor.h"
15 : #include "absl/strings/str_cat.h"
16 : #include "backend/catalog/googlesql_catalog.h"
17 : #include "backend/engine/duckdb/duckdb_executor.h"
18 : #include "backend/engine/engine.h"
19 : #include "backend/storage/duckdb/duckdb_storage.h"
20 : #include "googlesql/public/analyzer.h"
21 : #include "googlesql/public/analyzer_options.h"
22 : #include "googlesql/public/analyzer_output.h"
23 : #include "googlesql/public/language_options.h"
24 : #include "googlesql/public/options.pb.h"
25 : #include "googlesql/public/types/type_factory.h"
26 : #include "googlesql/resolved_ast/resolved_ast.h"
27 : #include "gtest/gtest.h"
28 :
29 : namespace bigquery_emulator {
30 : namespace backend {
31 : namespace engine {
32 : namespace duckdb {
33 : namespace {
34 :
35 : namespace fs = std::filesystem;
36 :
37 2 : ::googlesql::LanguageOptions MakeLanguageOptions() {
38 2 : ::googlesql::LanguageOptions language;
39 2 : language.EnableMaximumLanguageFeatures();
40 2 : language.set_product_mode(::googlesql::PRODUCT_EXTERNAL);
41 2 : language.set_name_resolution_mode(::googlesql::NAME_RESOLUTION_DEFAULT);
42 2 : return language;
43 2 : }
44 :
45 1 : ::googlesql::AnalyzerOptions MakeAnalyzerOptions(bool all_statements) {
46 1 : ::googlesql::AnalyzerOptions options(MakeLanguageOptions());
47 1 : options.set_error_message_mode(::googlesql::ERROR_MESSAGE_ONE_LINE);
48 1 : options.set_attach_error_location_payload(true);
49 1 : options.CreateDefaultArenasIfNotSet();
50 1 : if (all_statements) {
51 0 : options.mutable_language()->SetSupportsAllStatementKinds();
52 0 : }
53 1 : return options;
54 1 : }
55 :
56 : class DuckDbExecutorStructCastTest : public ::testing::Test {
57 : protected:
58 1 : void SetUp() override {
59 1 : const char* tmpdir_env = std::getenv("TMPDIR");
60 1 : const std::string tmpdir = tmpdir_env != nullptr ? tmpdir_env : "/tmp";
61 1 : std::random_device rd;
62 1 : std::seed_seq seed{rd(), rd()};
63 1 : std::mt19937_64 rng(seed);
64 1 : data_dir_ = fs::path(tmpdir) /
65 1 : absl::StrCat("bqemu-duckdb-struct-cast-test-", rng());
66 1 : std::error_code ec;
67 1 : fs::remove_all(data_dir_, ec);
68 1 : auto opened = storage::duckdb::DuckDBStorage::Open(data_dir_.string());
69 2 : ASSERT_TRUE(opened.ok()) << opened.status();
70 1 : storage_ = std::move(opened).value();
71 1 : executor_ = std::make_unique<DuckDbExecutor>(storage_.get());
72 1 : }
73 :
74 1 : void TearDown() override {
75 1 : executor_.reset();
76 1 : storage_.reset();
77 1 : std::error_code ec;
78 1 : fs::remove_all(data_dir_, ec);
79 1 : }
80 :
81 1 : QueryRequest MakeRequest(absl::string_view sql) {
82 1 : QueryRequest req;
83 1 : req.project_id = "proj-test";
84 1 : req.sql = std::string(sql);
85 1 : return req;
86 1 : }
87 :
88 : struct CatalogBundle {
89 : std::unique_ptr<::googlesql::TypeFactory> type_factory{};
90 : std::unique_ptr<catalog::GoogleSqlCatalog> catalog{};
91 : };
92 1 : CatalogBundle MakeCatalog() {
93 1 : auto type_factory = std::make_unique<::googlesql::TypeFactory>();
94 1 : auto catalog = std::make_unique<catalog::GoogleSqlCatalog>(
95 1 : "proj-test", storage_.get(), type_factory.get(), MakeLanguageOptions());
96 1 : return {std::move(type_factory), std::move(catalog)};
97 1 : }
98 :
99 : absl::StatusOr<std::unique_ptr<const ::googlesql::AnalyzerOutput>> Analyze(
100 : absl::string_view sql,
101 : ::googlesql::Catalog* catalog,
102 1 : bool all_statements) {
103 1 : ::googlesql::AnalyzerOptions options = MakeAnalyzerOptions(all_statements);
104 1 : ::googlesql::TypeFactory type_factory;
105 1 : std::unique_ptr<const ::googlesql::AnalyzerOutput> output;
106 1 : absl::Status s = ::googlesql::AnalyzeStatement(
107 1 : sql, options, catalog, &type_factory, &output);
108 1 : if (!s.ok()) return s;
109 1 : return output;
110 1 : }
111 :
112 : fs::path data_dir_{};
113 : std::unique_ptr<storage::duckdb::DuckDBStorage> storage_{};
114 : std::unique_ptr<DuckDbExecutor> executor_{};
115 : };
116 :
117 1 : TEST_F(DuckDbExecutorStructCastTest, ExecuteUnnestArrayStructPositionalCast) {
118 1 : static constexpr char kSql[] = R"sql(
119 1 : SELECT t.x, t.y
120 1 : FROM UNNEST(ARRAY<STRUCT<x INT64, y STRING>>[STRUCT(1, 'a'), STRUCT(2, 'b')]) AS t
121 1 : )sql";
122 1 : const std::string sql(kSql);
123 1 : CatalogBundle bundle = MakeCatalog();
124 1 : auto analyzed = Analyze(sql, bundle.catalog.get(), /*all_statements=*/false);
125 2 : ASSERT_TRUE(analyzed.ok()) << analyzed.status();
126 1 : const ::googlesql::ResolvedStatement* stmt =
127 1 : (*analyzed)->resolved_statement();
128 1 : ASSERT_NE(stmt, nullptr);
129 :
130 1 : absl::StatusOr<std::unique_ptr<RowSource>> source =
131 1 : executor_->ExecuteQuery(MakeRequest(sql), *stmt, bundle.catalog.get());
132 2 : ASSERT_TRUE(source.ok()) << source.status();
133 :
134 1 : storage::Row row;
135 1 : auto has = (*source)->Next(&row);
136 2 : ASSERT_TRUE(has.ok()) << has.status();
137 1 : ASSERT_TRUE(*has);
138 1 : ASSERT_EQ(row.cells.size(), 2u);
139 1 : EXPECT_EQ(row.cells[0].int64_value(), 1);
140 1 : EXPECT_EQ(row.cells[1].string_value(), "a");
141 :
142 1 : has = (*source)->Next(&row);
143 2 : ASSERT_TRUE(has.ok()) << has.status();
144 1 : ASSERT_TRUE(*has);
145 1 : EXPECT_EQ(row.cells[0].int64_value(), 2);
146 1 : EXPECT_EQ(row.cells[1].string_value(), "b");
147 :
148 1 : has = (*source)->Next(&row);
149 2 : ASSERT_TRUE(has.ok()) << has.status();
150 1 : EXPECT_FALSE(*has);
151 1 : }
152 :
153 : } // namespace
154 : } // namespace duckdb
155 : } // namespace engine
156 : } // namespace backend
157 : } // namespace bigquery_emulator
|
