Line data Source code
1 : #include "backend/sqltools/sql_tools.h"
2 :
3 : #include "absl/strings/match.h"
4 : #include "backend/sqltools/sql_references.h"
5 : #include "googlesql/public/builtin_function_options.h"
6 : #include "googlesql/public/function.h"
7 : #include "googlesql/public/function_signature.h"
8 : #include "googlesql/public/simple_catalog.h"
9 : #include "googlesql/public/types/type_factory.h"
10 : #include "gtest/gtest.h"
11 :
12 : namespace bigquery_emulator {
13 : namespace backend {
14 : namespace sqltools {
15 : namespace {
16 :
17 : class SqlToolsTest : public ::testing::Test {
18 : protected:
19 17 : void SetUp() override {
20 17 : language_ = MakeSqlToolsLanguageOptions();
21 17 : catalog_ = std::make_unique<::googlesql::SimpleCatalog>("test-project",
22 17 : &type_factory_);
23 17 : catalog_->AddBuiltinFunctionsAndTypes(
24 17 : ::googlesql::BuiltinFunctionOptions(language_));
25 17 : }
26 :
27 1 : void AddScalarFunction(const std::string& name) {
28 1 : ::googlesql::FunctionSignature signature(
29 1 : ::googlesql::FunctionArgumentType(::googlesql::types::Int64Type()),
30 1 : /*arguments=*/{},
31 1 : /*context_id=*/static_cast<int64_t>(0));
32 1 : auto function = std::make_unique<::googlesql::Function>(
33 1 : std::vector<std::string>{name},
34 1 : /*group=*/"External_function",
35 1 : ::googlesql::Function::SCALAR,
36 1 : std::vector<::googlesql::FunctionSignature>{signature});
37 1 : catalog_->AddFunction(function.get());
38 1 : owned_functions_.push_back(std::move(function));
39 1 : }
40 :
41 : ::googlesql::LanguageOptions language_;
42 : ::googlesql::TypeFactory type_factory_;
43 : std::unique_ptr<::googlesql::SimpleCatalog> catalog_;
44 : std::vector<std::unique_ptr<const ::googlesql::Function>> owned_functions_;
45 : };
46 :
47 1 : TEST_F(SqlToolsTest, FormatLenientProducesIndentedSql) {
48 1 : const absl::StatusOr<FormatResult> result =
49 1 : FormatSqlText("select 1", FormatOptions{});
50 2 : ASSERT_TRUE(result.ok()) << result.status();
51 1 : EXPECT_NE(result->formatted_sql.find("SELECT"), std::string::npos);
52 1 : EXPECT_NE(result->formatted_sql.find("1"), std::string::npos);
53 1 : }
54 :
55 1 : TEST_F(SqlToolsTest, ParseValidSelectReturnsStatementKind) {
56 1 : const absl::StatusOr<ParseResult> result =
57 1 : ParseSqlText("SELECT 1", language_);
58 2 : ASSERT_TRUE(result.ok()) << result.status();
59 1 : EXPECT_TRUE(result->diagnostics.empty());
60 1 : ASSERT_EQ(result->statement_kinds.size(), 1u);
61 1 : EXPECT_EQ(result->statement_kinds[0], "QueryStatement");
62 1 : }
63 :
64 1 : TEST_F(SqlToolsTest, ParseInvalidSqlReturnsDiagnostic) {
65 1 : const absl::StatusOr<ParseResult> result = ParseSqlText("SELEC 1", language_);
66 2 : ASSERT_TRUE(result.ok()) << result.status();
67 1 : EXPECT_FALSE(result->diagnostics.empty());
68 1 : }
69 :
70 1 : TEST_F(SqlToolsTest, TokenizeSelectReturnsKeywords) {
71 1 : TokenizeOptions options;
72 1 : const absl::StatusOr<TokenizeResult> result =
73 1 : TokenizeSqlText("SELECT 1", language_, options);
74 2 : ASSERT_TRUE(result.ok()) << result.status();
75 1 : ASSERT_GE(result->tokens.size(), 2u);
76 1 : EXPECT_EQ(result->tokens[0].kind, "keyword");
77 1 : EXPECT_EQ(result->tokens[0].image, "SELECT");
78 1 : }
79 :
80 1 : TEST_F(SqlToolsTest, CompleteAfterSelectIncludesKeywords) {
81 1 : CatalogNames names;
82 1 : const std::string sql = "SELECT ";
83 1 : const absl::StatusOr<CompleteResult> result =
84 1 : CompleteSqlText(sql, sql.size(), language_, catalog_.get(), names, "");
85 2 : ASSERT_TRUE(result.ok()) << result.status();
86 1 : bool found_from = false;
87 252 : for (const CompletionCandidate& candidate : result->candidates) {
88 252 : if (candidate.label == "FROM") {
89 1 : found_from = true;
90 1 : EXPECT_EQ(candidate.kind, "keyword");
91 1 : break;
92 1 : }
93 252 : }
94 1 : EXPECT_TRUE(found_from);
95 1 : }
96 :
97 1 : TEST_F(SqlToolsTest, CompleteStatementStartSuggestsOnlyClauseWords) {
98 1 : CatalogNames names;
99 1 : const std::string sql = "S";
100 1 : const absl::StatusOr<CompleteResult> result =
101 1 : CompleteSqlText(sql, sql.size(), language_, catalog_.get(), names, "");
102 2 : ASSERT_TRUE(result.ok()) << result.status();
103 :
104 1 : bool found_select = false;
105 1 : bool found_set = false;
106 1 : bool found_safe_cast = false;
107 1 : bool found_some = false;
108 1 : bool found_struct = false;
109 1 : bool found_function = false;
110 2 : for (const CompletionCandidate& candidate : result->candidates) {
111 2 : if (candidate.label == "SELECT") found_select = true;
112 2 : if (candidate.label == "SET") found_set = true;
113 2 : if (candidate.label == "SAFE_CAST") found_safe_cast = true;
114 2 : if (candidate.label == "SOME") found_some = true;
115 2 : if (candidate.label == "STRUCT") found_struct = true;
116 2 : if (candidate.kind == "function") found_function = true;
117 2 : }
118 1 : EXPECT_TRUE(found_select);
119 1 : EXPECT_TRUE(found_set);
120 1 : EXPECT_FALSE(found_safe_cast);
121 1 : EXPECT_FALSE(found_some);
122 1 : EXPECT_FALSE(found_struct);
123 1 : EXPECT_FALSE(found_function);
124 1 : }
125 :
126 1 : TEST_F(SqlToolsTest, CompleteAfterSelectExpressionContextUsesCuratedFunctions) {
127 1 : CatalogNames names;
128 1 : const std::string sql = "SELECT S";
129 1 : const absl::StatusOr<CompleteResult> result =
130 1 : CompleteSqlText(sql, sql.size(), language_, catalog_.get(), names, "");
131 2 : ASSERT_TRUE(result.ok()) << result.status();
132 :
133 1 : bool found_safe_add = false;
134 1 : bool found_schema = false;
135 1 : bool found_search = false;
136 1 : bool found_s2_function = false;
137 87 : for (const CompletionCandidate& candidate : result->candidates) {
138 87 : if (candidate.label == "SAFE_ADD" && candidate.kind == "function") {
139 1 : found_safe_add = true;
140 1 : EXPECT_EQ(candidate.detail, "(X, Y)");
141 1 : EXPECT_EQ(candidate.insert_text, "SAFE_ADD(");
142 1 : }
143 87 : if (candidate.label == "SCHEMA") found_schema = true;
144 87 : if (candidate.label == "SEARCH") found_search = true;
145 87 : if (absl::StartsWithIgnoreCase(candidate.label, "s2_") &&
146 87 : candidate.kind == "function") {
147 0 : found_s2_function = true;
148 0 : }
149 87 : }
150 1 : EXPECT_TRUE(found_safe_add);
151 1 : EXPECT_TRUE(found_schema);
152 1 : EXPECT_TRUE(found_search);
153 1 : EXPECT_FALSE(found_s2_function);
154 1 : }
155 :
156 1 : TEST_F(SqlToolsTest, CompleteAfterSemicolonUsesStatementStartBehavior) {
157 1 : CatalogNames names;
158 1 : const std::string sql = "SELECT 1 FROM t; S";
159 1 : const absl::StatusOr<CompleteResult> result =
160 1 : CompleteSqlText(sql, sql.size(), language_, catalog_.get(), names, "");
161 2 : ASSERT_TRUE(result.ok()) << result.status();
162 :
163 1 : bool found_select = false;
164 1 : bool found_set = false;
165 1 : bool found_function = false;
166 2 : for (const CompletionCandidate& candidate : result->candidates) {
167 2 : if (candidate.label == "SELECT") found_select = true;
168 2 : if (candidate.label == "SET") found_set = true;
169 2 : if (candidate.kind == "function") found_function = true;
170 2 : }
171 1 : EXPECT_TRUE(found_select);
172 1 : EXPECT_TRUE(found_set);
173 1 : EXPECT_FALSE(found_function);
174 1 : }
175 :
176 1 : TEST_F(SqlToolsTest, ParseInvalidSqlReturnsDiagnosticWithSpan) {
177 1 : const absl::StatusOr<ParseResult> result = ParseSqlText("SELEC 1", language_);
178 2 : ASSERT_TRUE(result.ok()) << result.status();
179 1 : ASSERT_FALSE(result->diagnostics.empty());
180 1 : const SqlDiagnostic& diag = result->diagnostics.front();
181 1 : EXPECT_GE(diag.start_byte, 0);
182 1 : EXPECT_GT(diag.end_byte, diag.start_byte);
183 1 : }
184 :
185 1 : TEST_F(SqlToolsTest, CompleteEmptyEditorAtCursorZero) {
186 1 : CatalogNames names;
187 1 : names.datasets = {"analytics"};
188 1 : const absl::StatusOr<CompleteResult> result =
189 1 : CompleteSqlText("", 0, language_, catalog_.get(), names, "analytics");
190 2 : ASSERT_TRUE(result.ok()) << result.status();
191 1 : bool found_select = false;
192 166 : for (const CompletionCandidate& candidate : result->candidates) {
193 166 : if (candidate.label == "SELECT") {
194 1 : found_select = true;
195 1 : break;
196 1 : }
197 166 : }
198 1 : EXPECT_TRUE(found_select);
199 1 : }
200 :
201 1 : TEST_F(SqlToolsTest, CompleteAfterFromUsesCatalogTables) {
202 1 : CatalogNames names;
203 1 : names.tables.push_back(
204 1 : CatalogTableEntry{"analytics.events", "p.analytics.events", "table", ""});
205 1 : names.tables.push_back(
206 1 : CatalogTableEntry{"events", "p.analytics.events", "table", ""});
207 1 : const std::string sql = "SELECT * FROM ev";
208 1 : const absl::StatusOr<CompleteResult> result = CompleteSqlText(
209 1 : sql, sql.size(), language_, catalog_.get(), names, "analytics");
210 2 : ASSERT_TRUE(result.ok()) << result.status();
211 1 : bool found_events = false;
212 1 : for (const CompletionCandidate& candidate : result->candidates) {
213 1 : if (candidate.label == "events") {
214 1 : found_events = true;
215 1 : break;
216 1 : }
217 1 : }
218 1 : EXPECT_TRUE(found_events);
219 1 : }
220 :
221 1 : TEST_F(SqlToolsTest, CompleteProjectQualifiedTableCandidate) {
222 1 : CatalogNames names;
223 1 : names.tables.push_back(
224 1 : CatalogTableEntry{"proj.ds.events", "proj.ds.events", "table", "table"});
225 1 : const std::string sql = "SELECT * FROM proj.d";
226 1 : const absl::StatusOr<CompleteResult> result =
227 1 : CompleteSqlText(sql, sql.size(), language_, catalog_.get(), names, "ds");
228 2 : ASSERT_TRUE(result.ok()) << result.status();
229 1 : bool found_fqn = false;
230 1 : for (const CompletionCandidate& candidate : result->candidates) {
231 1 : if (candidate.label == "proj.ds.events") {
232 1 : found_fqn = true;
233 1 : break;
234 1 : }
235 1 : }
236 1 : EXPECT_TRUE(found_fqn);
237 1 : }
238 :
239 1 : TEST_F(SqlToolsTest, CompleteColumnContextExcludesOutOfScopeColumns) {
240 1 : CatalogNames names;
241 1 : names.columns_by_table["sales_dataset.transactions"] = {
242 1 : CatalogColumnEntry{"transaction_id", "STRING"},
243 1 : CatalogColumnEntry{"customer_id", "INT64"},
244 1 : CatalogColumnEntry{"total_amount", "NUMERIC"},
245 1 : CatalogColumnEntry{"purchase_date", "TIMESTAMP"},
246 1 : };
247 1 : names.columns_by_table["test-dataset.table_a"] = {
248 1 : CatalogColumnEntry{"id", "INT64"},
249 1 : CatalogColumnEntry{"name", "STRING"},
250 1 : CatalogColumnEntry{"skillNum", "NUMERIC"},
251 1 : CatalogColumnEntry{"structarr", "ARRAY<STRUCT<key STRING, value JSON>>"},
252 1 : };
253 1 : names.columns_by_table["products_dataset.inventory"] = {
254 1 : CatalogColumnEntry{"product_id", "STRING"},
255 1 : CatalogColumnEntry{"specs", "STRUCT<weight FLOAT64>"},
256 1 : };
257 1 : PopulateInScopeTablesFromHeuristic(
258 1 : "SELECT s FROM sales_dataset.transactions", language_, "", &names);
259 1 : ASSERT_EQ(names.in_scope_tables.size(), 1u);
260 :
261 1 : const std::string sql = "SELECT s FROM sales_dataset.transactions";
262 1 : const absl::StatusOr<CompleteResult> result =
263 1 : CompleteSqlText(sql, 7, language_, catalog_.get(), names, "");
264 2 : ASSERT_TRUE(result.ok()) << result.status();
265 :
266 613 : for (const CompletionCandidate& candidate : result->candidates) {
267 613 : if (candidate.kind != "column") continue;
268 4 : EXPECT_NE(candidate.label, "skillNum");
269 4 : EXPECT_NE(candidate.label, "structarr");
270 4 : EXPECT_NE(candidate.label, "specs");
271 4 : }
272 1 : }
273 :
274 1 : TEST_F(SqlToolsTest, CompleteColumnContextIncludesInScopeColumns) {
275 1 : CatalogNames names;
276 1 : names.columns_by_table["sales_dataset.transactions"] = {
277 1 : CatalogColumnEntry{"transaction_id", "STRING"},
278 1 : CatalogColumnEntry{"customer_id", "INT64"},
279 1 : };
280 1 : names.columns_by_table["test-dataset.table_a"] = {
281 1 : CatalogColumnEntry{"skillNum", "NUMERIC"},
282 1 : };
283 1 : PopulateInScopeTablesFromHeuristic(
284 1 : "SELECT t FROM sales_dataset.transactions", language_, "", &names);
285 :
286 1 : const std::string sql = "SELECT t FROM sales_dataset.transactions";
287 1 : const absl::StatusOr<CompleteResult> result =
288 1 : CompleteSqlText(sql, 7, language_, catalog_.get(), names, "");
289 2 : ASSERT_TRUE(result.ok()) << result.status();
290 :
291 1 : bool found_transaction_id = false;
292 611 : for (const CompletionCandidate& candidate : result->candidates) {
293 611 : if (candidate.label == "transaction_id" && candidate.kind == "column") {
294 1 : found_transaction_id = true;
295 1 : }
296 611 : if (candidate.kind == "column" && candidate.label == "skillNum") {
297 0 : FAIL()
298 0 : << "skillNum should not be suggested for sales_dataset.transactions";
299 0 : }
300 611 : }
301 1 : EXPECT_TRUE(found_transaction_id);
302 1 : }
303 :
304 1 : TEST_F(SqlToolsTest, CompleteIncompleteSqlUsesHeuristicColumns) {
305 1 : CatalogNames names;
306 1 : names.columns_by_table["analytics.events"] = {
307 1 : CatalogColumnEntry{"id", "INT64"},
308 1 : CatalogColumnEntry{"name", "STRING"},
309 1 : };
310 1 : PopulateInScopeTablesFromHeuristic(
311 1 : "SELECT na FROM analytics.events WHERE ", language_, "analytics", &names);
312 1 : ASSERT_EQ(names.in_scope_tables.size(), 1u);
313 1 : ASSERT_EQ(names.in_scope_tables[0].columns.size(), 2u);
314 :
315 1 : const std::string sql = "SELECT na FROM analytics.events WHERE ";
316 1 : const absl::StatusOr<CompleteResult> result = CompleteSqlText(
317 1 : sql, sql.size(), language_, catalog_.get(), names, "analytics");
318 2 : ASSERT_TRUE(result.ok()) << result.status();
319 1 : bool found_name = false;
320 611 : for (const CompletionCandidate& candidate : result->candidates) {
321 611 : if (candidate.label == "name") {
322 1 : found_name = true;
323 1 : EXPECT_EQ(candidate.kind, "column");
324 1 : break;
325 1 : }
326 611 : }
327 1 : EXPECT_TRUE(found_name);
328 1 : }
329 :
330 1 : TEST_F(SqlToolsTest, CompleteUserRoutineNotDuplicatedAsFunction) {
331 1 : AddScalarFunction("add_one");
332 1 : CatalogNames names;
333 1 : names.routines.push_back(CatalogRoutineEntry{"ds.add_one",
334 1 : "test-project.ds.add_one",
335 1 : "routine",
336 1 : "SQL scalar function"});
337 1 : names.routines.push_back(CatalogRoutineEntry{
338 1 : "add_one", "test-project.ds.add_one", "routine", "SQL scalar function"});
339 :
340 1 : const std::string sql = "SELECT add_";
341 1 : const absl::StatusOr<CompleteResult> result =
342 1 : CompleteSqlText(sql, sql.size(), language_, catalog_.get(), names, "ds");
343 2 : ASSERT_TRUE(result.ok()) << result.status();
344 :
345 1 : bool found_routine = false;
346 1 : bool found_function = false;
347 1 : for (const CompletionCandidate& candidate : result->candidates) {
348 1 : if (candidate.label == "add_one" && candidate.kind == "routine") {
349 1 : found_routine = true;
350 1 : EXPECT_EQ(candidate.fqn, "test-project.ds.add_one");
351 1 : EXPECT_EQ(candidate.insert_text, "add_one(");
352 1 : }
353 1 : if (candidate.label == "add_one" && candidate.kind == "function") {
354 0 : found_function = true;
355 0 : }
356 1 : }
357 1 : EXPECT_TRUE(found_routine);
358 1 : EXPECT_FALSE(found_function);
359 1 : }
360 :
361 1 : TEST_F(SqlToolsTest, CompleteRoutineCandidateIncludesFqn) {
362 1 : CatalogNames names;
363 1 : names.routines.push_back(CatalogRoutineEntry{
364 1 : "proj.ds.my_fn",
365 1 : "proj.ds.my_fn",
366 1 : "routine",
367 1 : "SQL scalar function",
368 1 : });
369 1 : const std::string sql = "CREATE FUNCTION ";
370 1 : const absl::StatusOr<CompleteResult> result =
371 1 : CompleteSqlText(sql, sql.size(), language_, catalog_.get(), names, "ds");
372 2 : ASSERT_TRUE(result.ok()) << result.status();
373 1 : bool found = false;
374 341 : for (const CompletionCandidate& candidate : result->candidates) {
375 341 : if (candidate.label == "proj.ds.my_fn") {
376 1 : found = true;
377 1 : EXPECT_EQ(candidate.kind, "routine");
378 1 : EXPECT_EQ(candidate.fqn, "proj.ds.my_fn");
379 1 : break;
380 1 : }
381 341 : }
382 1 : EXPECT_TRUE(found);
383 1 : }
384 :
385 : } // namespace
386 : } // namespace sqltools
387 : } // namespace backend
388 : } // namespace bigquery_emulator
|
