Line data    Source code

       1              : // Routing-bug defense tests for the DuckDB executor. Kept in a
       2              : // 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/strings/str_cat.h"
      15              : #include "backend/catalog/googlesql_catalog.h"
      16              : #include "backend/engine/duckdb/duckdb_executor.h"
      17              : #include "backend/engine/engine.h"
      18              : #include "backend/storage/duckdb/duckdb_storage.h"
      19              : #include "googlesql/public/analyzer.h"
      20              : #include "googlesql/public/analyzer_options.h"
      21              : #include "googlesql/public/analyzer_output.h"
      22              : #include "googlesql/public/language_options.h"
      23              : #include "googlesql/public/options.pb.h"
      24              : #include "googlesql/public/types/type_factory.h"
      25              : #include "googlesql/resolved_ast/resolved_ast.h"
      26              : #include "gtest/gtest.h"
      27              : 
      28              : namespace bigquery_emulator {
      29              : namespace backend {
      30              : namespace engine {
      31              : namespace duckdb {
      32              : namespace {
      33              : 
      34              : namespace fs = std::filesystem;
      35              : 
      36            2 : ::googlesql::LanguageOptions MakeLanguageOptions() {
      37            2 :   ::googlesql::LanguageOptions language;
      38            2 :   language.EnableMaximumLanguageFeatures();
      39            2 :   language.set_product_mode(::googlesql::PRODUCT_EXTERNAL);
      40            2 :   language.set_name_resolution_mode(::googlesql::NAME_RESOLUTION_DEFAULT);
      41            2 :   return language;
      42            2 : }
      43              : 
      44            1 : ::googlesql::AnalyzerOptions MakeAnalyzerOptions(bool all_statements) {
      45            1 :   ::googlesql::AnalyzerOptions options(MakeLanguageOptions());
      46            1 :   options.set_error_message_mode(::googlesql::ERROR_MESSAGE_ONE_LINE);
      47            1 :   options.set_attach_error_location_payload(true);
      48            1 :   options.CreateDefaultArenasIfNotSet();
      49            1 :   if (all_statements) {
      50            1 :     options.mutable_language()->SetSupportsAllStatementKinds();
      51            1 :   }
      52            1 :   return options;
      53            1 : }
      54              : 
      55              : class DuckDbExecutorRoutingTest : public ::testing::Test {
      56              :  protected:
      57            1 :   void SetUp() override {
      58            1 :     const char* tmpdir_env = std::getenv("TMPDIR");
      59            1 :     const std::string tmpdir = tmpdir_env != nullptr ? tmpdir_env : "/tmp";
      60            1 :     std::random_device rd;
      61            1 :     std::seed_seq seed{rd(), rd()};
      62            1 :     std::mt19937_64 rng(seed);
      63            1 :     data_dir_ =
      64            1 :         fs::path(tmpdir) / absl::StrCat("bqemu-duckdb-routing-test-", rng());
      65            1 :     std::error_code ec;
      66            1 :     fs::remove_all(data_dir_, ec);
      67            1 :     auto opened = storage::duckdb::DuckDBStorage::Open(data_dir_.string());
      68            2 :     ASSERT_TRUE(opened.ok()) << opened.status();
      69            1 :     storage_ = std::move(opened).value();
      70            1 :     executor_ = std::make_unique<DuckDbExecutor>(storage_.get());
      71            1 :   }
      72              : 
      73            1 :   void TearDown() override {
      74            1 :     executor_.reset();
      75            1 :     storage_.reset();
      76            1 :     std::error_code ec;
      77            1 :     fs::remove_all(data_dir_, ec);
      78            1 :   }
      79              : 
      80            1 :   QueryRequest MakeRequest(absl::string_view sql) {
      81            1 :     QueryRequest req;
      82            1 :     req.project_id = "proj-test";
      83            1 :     req.sql = std::string(sql);
      84            1 :     return req;
      85            1 :   }
      86              : 
      87              :   struct CatalogBundle {
      88              :     std::unique_ptr<::googlesql::TypeFactory> type_factory{};
      89              :     std::unique_ptr<catalog::GoogleSqlCatalog> catalog{};
      90              :   };
      91            1 :   CatalogBundle MakeCatalog() {
      92            1 :     auto type_factory = std::make_unique<::googlesql::TypeFactory>();
      93            1 :     auto catalog = std::make_unique<catalog::GoogleSqlCatalog>(
      94            1 :         "proj-test", storage_.get(), type_factory.get(), MakeLanguageOptions());
      95            1 :     return {std::move(type_factory), std::move(catalog)};
      96            1 :   }
      97              : 
      98              :   absl::StatusOr<std::unique_ptr<const ::googlesql::AnalyzerOutput>> Analyze(
      99              :       absl::string_view sql,
     100              :       ::googlesql::Catalog* catalog,
     101            1 :       bool all_statements) {
     102            1 :     ::googlesql::AnalyzerOptions options = MakeAnalyzerOptions(all_statements);
     103            1 :     ::googlesql::TypeFactory type_factory;
     104            1 :     std::unique_ptr<const ::googlesql::AnalyzerOutput> output;
     105            1 :     absl::Status s = ::googlesql::AnalyzeStatement(
     106            1 :         sql, options, catalog, &type_factory, &output);
     107            1 :     if (!s.ok()) return s;
     108            1 :     return output;
     109            1 :   }
     110              : 
     111              :   fs::path data_dir_{};
     112              :   std::unique_ptr<storage::duckdb::DuckDBStorage> storage_{};
     113              :   std::unique_ptr<DuckDbExecutor> executor_{};
     114              : };
     115              : 
     116              : TEST_F(DuckDbExecutorRoutingTest,
     117            1 :        ExecuteDdlRejectsCreateTableAfterControlOpMigration) {
     118            1 :   CatalogBundle bundle = MakeCatalog();
     119            1 :   auto analyzed = Analyze("CREATE TABLE ds.t (id INT64, name STRING)",
     120            1 :                           bundle.catalog.get(),
     121            1 :                           /*all_statements=*/true);
     122            2 :   ASSERT_TRUE(analyzed.ok()) << analyzed.status();
     123            1 :   const ::googlesql::ResolvedStatement* stmt =
     124            1 :       (*analyzed)->resolved_statement();
     125            1 :   ASSERT_NE(stmt, nullptr);
     126              : 
     127            1 :   absl::Status s = executor_->ExecuteDdl(
     128            1 :       MakeRequest("CREATE TABLE ds.t (id INT64, name STRING)"),
     129            1 :       *stmt,
     130            1 :       bundle.catalog.get());
     131            1 :   ASSERT_FALSE(s.ok());
     132            2 :   EXPECT_EQ(s.code(), absl::StatusCode::kUnimplemented) << s;
     133            2 :   EXPECT_NE(std::string(s.message()).find("ControlOpExecutor"),
     134            2 :             std::string::npos)
     135            2 :       << s.message();
     136            1 : }
     137              : 
     138              : }  // namespace
     139              : }  // namespace duckdb
     140              : }  // namespace engine
     141              : }  // namespace backend
     142              : }  // namespace bigquery_emulator