Line data Source code
1 : // Regression: cached registration catalogs must replay procedures registered
2 : // after the catalog was first constructed (CREATE PROCEDURE in setup, CALL in
3 : // a later script). See udf_registration_catalog.cc reuse branch.
4 :
5 : #include "backend/catalog/udf_registration_catalog.h"
6 :
7 : #include <cstdlib>
8 : #include <filesystem>
9 : #include <memory>
10 : #include <random>
11 : #include <string>
12 : #include <system_error>
13 : #include <utility>
14 :
15 : #include "absl/status/status.h"
16 : #include "absl/strings/str_cat.h"
17 : #include "backend/catalog/googlesql_catalog.h"
18 : #include "backend/catalog/procedure_registry.h"
19 : #include "backend/catalog/stored_procedure.h"
20 : #include "backend/catalog/udf_registry.h"
21 : #include "backend/schema/schema.h"
22 : #include "backend/storage/duckdb/duckdb_storage.h"
23 : #include "googlesql/public/analyzer.h"
24 : #include "googlesql/public/analyzer_options.h"
25 : #include "googlesql/public/analyzer_output.h"
26 : #include "googlesql/public/language_options.h"
27 : #include "googlesql/public/options.pb.h"
28 : #include "googlesql/public/simple_catalog.h"
29 : #include "googlesql/public/types/type_factory.h"
30 : #include "googlesql/resolved_ast/resolved_ast.h"
31 : #include "googlesql/resolved_ast/resolved_node_kind.pb.h"
32 : #include "gtest/gtest.h"
33 :
34 : namespace bigquery_emulator {
35 : namespace backend {
36 : namespace catalog {
37 : namespace {
38 :
39 : namespace fs = std::filesystem;
40 :
41 : const char* kProject = "proj_udf_registration_catalog_test";
42 :
43 3 : ::googlesql::LanguageOptions MakeLanguageOptions() {
44 3 : ::googlesql::LanguageOptions language;
45 3 : language.EnableMaximumLanguageFeatures();
46 3 : language.set_product_mode(::googlesql::PRODUCT_EXTERNAL);
47 3 : language.set_name_resolution_mode(::googlesql::NAME_RESOLUTION_DEFAULT);
48 3 : return language;
49 3 : }
50 :
51 1 : ::googlesql::AnalyzerOptions MakeAnalyzerOptions() {
52 1 : ::googlesql::AnalyzerOptions options(MakeLanguageOptions());
53 1 : options.set_error_message_mode(::googlesql::ERROR_MESSAGE_ONE_LINE);
54 1 : options.CreateDefaultArenasIfNotSet();
55 1 : options.mutable_language()->SetSupportsAllStatementKinds();
56 1 : return options;
57 1 : }
58 :
59 : class UdfRegistrationCatalogTest : public ::testing::Test {
60 : protected:
61 1 : void SetUp() override {
62 1 : const char* tmpdir_env = std::getenv("TMPDIR");
63 1 : const std::string tmpdir = tmpdir_env != nullptr ? tmpdir_env : "/tmp";
64 1 : std::random_device rd;
65 1 : std::seed_seq seed{rd(), rd()};
66 1 : std::mt19937_64 rng(seed);
67 1 : data_dir_ =
68 1 : fs::path(tmpdir) / absl::StrCat("bqemu-udf-reg-catalog-test-", rng());
69 1 : std::error_code ec;
70 1 : fs::remove_all(data_dir_, ec);
71 1 : auto opened = storage::duckdb::DuckDBStorage::Open(data_dir_.string());
72 2 : ASSERT_TRUE(opened.ok()) << opened.status();
73 1 : storage_ = std::move(opened).value();
74 1 : ASSERT_TRUE(storage_->CreateDataset({kProject, "ds"}, "US").ok());
75 1 : }
76 :
77 1 : void TearDown() override {
78 1 : storage_.reset();
79 1 : std::error_code ec;
80 1 : fs::remove_all(data_dir_, ec);
81 1 : }
82 :
83 1 : absl::Status RegisterProcedureFromSql(absl::string_view sql) {
84 1 : ::googlesql::TypeFactory analyze_tf;
85 1 : std::unique_ptr<GoogleSqlCatalog> catalog =
86 1 : std::make_unique<GoogleSqlCatalog>(
87 1 : kProject, storage_.get(), &analyze_tf, MakeLanguageOptions(), "ds");
88 1 : std::unique_ptr<const ::googlesql::AnalyzerOutput> output;
89 1 : absl::Status analyzed = ::googlesql::AnalyzeStatement(
90 1 : sql, MakeAnalyzerOptions(), catalog.get(), &analyze_tf, &output);
91 1 : if (!analyzed.ok()) return analyzed;
92 1 : const ::googlesql::ResolvedStatement* stmt = output->resolved_statement();
93 1 : if (stmt == nullptr) {
94 0 : return absl::InternalError("analyzer returned null statement");
95 0 : }
96 1 : if (stmt->node_kind() != ::googlesql::RESOLVED_CREATE_PROCEDURE_STMT) {
97 0 : return absl::InvalidArgumentError("expected CREATE PROCEDURE statement");
98 0 : }
99 1 : const auto* create_procedure =
100 1 : stmt->GetAs<::googlesql::ResolvedCreateProcedureStmt>();
101 1 : if (create_procedure == nullptr) {
102 0 : return absl::InternalError("CREATE PROCEDURE has null resolved stmt");
103 0 : }
104 1 : return RegisterProjectProcedure(
105 1 : kProject, *create_procedure, std::move(output));
106 1 : }
107 :
108 : fs::path data_dir_{};
109 : std::unique_ptr<storage::duckdb::DuckDBStorage> storage_{};
110 : };
111 :
112 : TEST_F(UdfRegistrationCatalogTest,
113 1 : ReusedRegistrationCatalogReplaysProceduresRegisteredAfterCreation) {
114 1 : ::googlesql::TypeFactory reg_tf;
115 1 : const ::googlesql::LanguageOptions language = MakeLanguageOptions();
116 :
117 1 : GoogleSqlCatalog* first = nullptr;
118 1 : first = GetOrCreateRegistrationCatalog(
119 1 : kProject, storage_.get(), ®_tf, language, "ds");
120 1 : ASSERT_NE(first, nullptr);
121 :
122 1 : const ::googlesql::Procedure* before = nullptr;
123 1 : EXPECT_FALSE(first->FindProcedure({"ds", "echo"}, &before).ok());
124 :
125 1 : ASSERT_TRUE(RegisterProcedureFromSql(
126 1 : "CREATE OR REPLACE PROCEDURE ds.echo(OUT arr ARRAY<INT64>) "
127 1 : "BEGIN SET arr = GENERATE_ARRAY(1, 3); END;")
128 1 : .ok());
129 :
130 1 : const StoredSQLProcedure* in_registry =
131 1 : FindProjectProcedure(kProject, "echo");
132 2 : ASSERT_NE(in_registry, nullptr)
133 2 : << "procedure must be in project registry after RegisterProjectProcedure";
134 :
135 1 : GoogleSqlCatalog* second = nullptr;
136 1 : second = GetOrCreateRegistrationCatalog(
137 1 : kProject, storage_.get(), ®_tf, language, "ds");
138 1 : ASSERT_NE(second, nullptr);
139 1 : EXPECT_EQ(second, first);
140 :
141 1 : const ::googlesql::Procedure* after = nullptr;
142 2 : ASSERT_TRUE(second->FindProcedure({"ds", "echo"}, &after).ok())
143 2 : << "reused registration catalog must replay procedures registered "
144 2 : "after first construction";
145 1 : ASSERT_NE(after, nullptr);
146 1 : EXPECT_EQ(after, in_registry);
147 1 : EXPECT_EQ(after->Name(), "echo");
148 1 : }
149 :
150 : } // namespace
151 : } // namespace catalog
152 : } // namespace backend
153 : } // namespace bigquery_emulator
|
