Line data    Source code

       1              : #ifndef BIGQUERY_EMULATOR_BACKEND_ENGINE_COORDINATOR_ROUTE_CLASSIFIER_TEST_FIXTURE_H_
       2              : #define BIGQUERY_EMULATOR_BACKEND_ENGINE_COORDINATOR_ROUTE_CLASSIFIER_TEST_FIXTURE_H_
       3              : 
       4              : #include <memory>
       5              : #include <string>
       6              : #include <vector>
       7              : 
       8              : #include "absl/status/status.h"
       9              : #include "absl/strings/string_view.h"
      10              : #include "backend/catalog/emulator_builtin_extensions.h"
      11              : #include "backend/catalog/emulator_ml_test_catalog.h"
      12              : #include "backend/catalog/emulator_ml_tvf_extensions.h"
      13              : #include "backend/engine/coordinator/route_classifier.h"
      14              : #include "googlesql/public/analyzer.h"
      15              : #include "googlesql/public/analyzer_options.h"
      16              : #include "googlesql/public/analyzer_output.h"
      17              : #include "googlesql/public/builtin_function_options.h"
      18              : #include "googlesql/public/catalog.h"
      19              : #include "googlesql/public/language_options.h"
      20              : #include "googlesql/public/options.pb.h"
      21              : #include "googlesql/public/types/type_factory.h"
      22              : #include "googlesql/resolved_ast/resolved_ast.h"
      23              : #include "gtest/gtest.h"
      24              : 
      25              : namespace bigquery_emulator {
      26              : namespace backend {
      27              : namespace engine {
      28              : namespace coordinator {
      29              : 
      30           42 : inline ::googlesql::AnalyzerOptions MakeAnalyzerOptions() {
      31           42 :   ::googlesql::LanguageOptions language;
      32           42 :   language.EnableMaximumLanguageFeaturesForDevelopment();
      33           42 :   language.EnableLanguageFeature(::googlesql::FEATURE_REMOTE_MODEL);
      34           42 :   language.set_product_mode(::googlesql::PRODUCT_EXTERNAL);
      35           42 :   language.set_name_resolution_mode(::googlesql::NAME_RESOLUTION_DEFAULT);
      36           42 :   language.SetSupportsAllStatementKinds();
      37           42 :   ::googlesql::AnalyzerOptions options(language);
      38           42 :   options.set_error_message_mode(::googlesql::ERROR_MESSAGE_ONE_LINE);
      39           42 :   options.disable_rewrite(::googlesql::REWRITE_PIVOT);
      40           42 :   options.disable_rewrite(::googlesql::REWRITE_UNPIVOT);
      41           42 :   options.CreateDefaultArenasIfNotSet();
      42           42 :   return options;
      43           42 : }
      44              : 
      45              : class RouteClassifierTest : public ::testing::Test {
      46              :  protected:
      47           44 :   void SetUp() override {
      48           44 :     type_factory_ = std::make_unique<::googlesql::TypeFactory>();
      49           44 :     catalog_ = std::make_unique<catalog::EmulatorMlTestCatalog>(
      50           44 :         "test_catalog", type_factory_.get());
      51           44 :     ::googlesql::LanguageOptions language;
      52           44 :     language.EnableMaximumLanguageFeaturesForDevelopment();
      53           44 :     language.EnableLanguageFeature(::googlesql::FEATURE_REMOTE_MODEL);
      54           44 :     language.set_product_mode(::googlesql::PRODUCT_EXTERNAL);
      55           44 :     ASSERT_TRUE(catalog_
      56           44 :                     ->AddBuiltinFunctionsAndTypes(
      57           44 :                         ::googlesql::BuiltinFunctionOptions(language))
      58           44 :                     .ok());
      59              : 
      60           44 :     auto people = std::make_unique<::googlesql::SimpleTable>(
      61           44 :         "people",
      62           44 :         std::vector<::googlesql::SimpleTable::NameAndType>{
      63           44 :             {"id", type_factory_->get_int64()},
      64           44 :             {"name", type_factory_->get_string()},
      65           44 :         });
      66           44 :     catalog_->AddOwnedTable(std::move(people));
      67              : 
      68           44 :     const ::googlesql::Type* int64_array_type = nullptr;
      69           44 :     ASSERT_TRUE(
      70           44 :         type_factory_
      71           44 :             ->MakeArrayType(type_factory_->get_int64(), &int64_array_type)
      72           44 :             .ok());
      73           44 :     auto arr_table = std::make_unique<::googlesql::SimpleTable>(
      74           44 :         "arr_table",
      75           44 :         std::vector<::googlesql::SimpleTable::NameAndType>{
      76           44 :             {"id", type_factory_->get_int64()},
      77           44 :             {"arr", int64_array_type},
      78           44 :         });
      79           44 :     catalog_->AddOwnedTable(std::move(arr_table));
      80           44 :     catalog::RegisterEmulatorMlTvfStubs(*catalog_);
      81           44 :   }
      82              : 
      83           42 :   const ::googlesql::ResolvedStatement* Analyze(absl::string_view sql) {
      84           42 :     last_output_.reset();
      85           42 :     ::googlesql::AnalyzerOptions options = MakeAnalyzerOptions();
      86           42 :     absl::Status s = ::googlesql::AnalyzeStatement(
      87           42 :         sql, options, catalog_.get(), type_factory_.get(), &last_output_);
      88           84 :     EXPECT_TRUE(s.ok()) << s;
      89           42 :     if (!s.ok() || last_output_ == nullptr) return nullptr;
      90           42 :     return last_output_->resolved_statement();
      91           42 :   }
      92              : 
      93              :   std::unique_ptr<::googlesql::TypeFactory> type_factory_{};
      94              :   std::unique_ptr<catalog::EmulatorMlTestCatalog> catalog_{};
      95              :   std::unique_ptr<const ::googlesql::AnalyzerOutput> last_output_{};
      96              :   RouteClassifier classifier_{};
      97              : };
      98              : 
      99              : }  // namespace coordinator
     100              : }  // namespace engine
     101              : }  // namespace backend
     102              : }  // namespace bigquery_emulator
     103              : 
     104              : #endif  // BIGQUERY_EMULATOR_BACKEND_ENGINE_COORDINATOR_ROUTE_CLASSIFIER_TEST_FIXTURE_H_