Line data    Source code

       1              : #include <memory>
       2              : #include <utility>
       3              : #include <vector>
       4              : 
       5              : #include "backend/engine/semantic/eval_context.h"
       6              : #include "backend/engine/semantic/scan_eval.h"
       7              : #include "backend/engine/semantic/value.h"
       8              : #include "googlesql/public/types/type_factory.h"
       9              : #include "googlesql/public/value.h"
      10              : #include "googlesql/resolved_ast/resolved_ast.h"
      11              : #include "googlesql/resolved_ast/resolved_node_kind.pb.h"
      12              : #include "gtest/gtest.h"
      13              : 
      14              : namespace bigquery_emulator {
      15              : namespace backend {
      16              : namespace engine {
      17              : namespace semantic {
      18              : namespace {
      19              : 
      20              : std::unique_ptr<const ::googlesql::ResolvedScan> MakePrivacyAggregateScan(
      21              :     ::googlesql::ResolvedNodeKind kind,
      22              :     const ::googlesql::ResolvedColumn& out_col,
      23            3 :     const ::googlesql::ResolvedColumn& side_col) {
      24            3 :   std::vector<std::unique_ptr<const ::googlesql::ResolvedComputedColumnBase>>
      25            3 :       aggregate_list;
      26            3 :   aggregate_list.push_back(::googlesql::MakeResolvedDeferredComputedColumn(
      27            3 :       out_col,
      28            3 :       ::googlesql::MakeResolvedLiteral(::googlesql::Value::Int64(7)),
      29            3 :       side_col));
      30              : 
      31            3 :   switch (kind) {
      32            1 :     case ::googlesql::RESOLVED_ANONYMIZED_AGGREGATE_SCAN:
      33            1 :       return ::googlesql::MakeResolvedAnonymizedAggregateScan(
      34            1 :           /*column_list=*/{out_col, side_col},
      35            1 :           ::googlesql::MakeResolvedSingleRowScan(),
      36            1 :           /*group_by_list=*/{},
      37            1 :           std::move(aggregate_list),
      38            1 :           /*grouping_set_list=*/{},
      39            1 :           /*rollup_column_list=*/{},
      40            1 :           ::googlesql::MakeResolvedLiteral(::googlesql::Value::Int64(1)),
      41            1 :           /*anonymization_option_list=*/{});
      42            1 :     case ::googlesql::RESOLVED_DIFFERENTIAL_PRIVACY_AGGREGATE_SCAN:
      43            1 :       return ::googlesql::MakeResolvedDifferentialPrivacyAggregateScan(
      44            1 :           /*column_list=*/{out_col, side_col},
      45            1 :           ::googlesql::MakeResolvedSingleRowScan(),
      46            1 :           /*group_by_list=*/{},
      47            1 :           std::move(aggregate_list),
      48            1 :           /*grouping_set_list=*/{},
      49            1 :           /*rollup_column_list=*/{},
      50            1 :           /*group_selection_threshold_expr=*/nullptr,
      51            1 :           /*option_list=*/{});
      52            1 :     case ::googlesql::RESOLVED_AGGREGATION_THRESHOLD_AGGREGATE_SCAN:
      53            1 :       return ::googlesql::MakeResolvedAggregationThresholdAggregateScan(
      54            1 :           /*column_list=*/{out_col, side_col},
      55            1 :           ::googlesql::MakeResolvedSingleRowScan(),
      56            1 :           /*group_by_list=*/{},
      57            1 :           std::move(aggregate_list),
      58            1 :           /*grouping_set_list=*/{},
      59            1 :           /*rollup_column_list=*/{},
      60            1 :           /*option_list=*/{});
      61            0 :     default:
      62            0 :       return nullptr;
      63            3 :   }
      64            3 : }
      65              : 
      66              : class PrivacyAggregateStubTest
      67              :     : public ::testing::TestWithParam<::googlesql::ResolvedNodeKind> {};
      68              : 
      69            3 : TEST_P(PrivacyAggregateStubTest, MaterializeScanEvaluatesPlainAggregate) {
      70            3 :   ::googlesql::TypeFactory type_factory;
      71            3 :   const ::googlesql::Type* int64 = type_factory.get_int64();
      72            3 :   const ::googlesql::Type* bytes = type_factory.get_bytes();
      73            3 :   ::googlesql::ResolvedColumn out_col(
      74            3 :       /*column_id=*/300,
      75            3 :       ::googlesql::IdString::MakeGlobal("$query"),
      76            3 :       ::googlesql::IdString::MakeGlobal("v"),
      77            3 :       int64);
      78            3 :   ::googlesql::ResolvedColumn side_col(
      79            3 :       /*column_id=*/301,
      80            3 :       ::googlesql::IdString::MakeGlobal("$query"),
      81            3 :       ::googlesql::IdString::MakeGlobal("_se"),
      82            3 :       bytes);
      83              : 
      84            3 :   auto scan = MakePrivacyAggregateScan(GetParam(), out_col, side_col);
      85            3 :   ASSERT_NE(scan, nullptr);
      86              : 
      87            3 :   EvalContext ctx{.project_id = "test"};
      88            3 :   auto rows_or = MaterializeScan(scan.get(), ctx);
      89            6 :   ASSERT_TRUE(rows_or.ok()) << rows_or.status();
      90            3 :   ASSERT_EQ(rows_or->size(), 1u);
      91            3 :   const ColumnBindings& row = rows_or->at(0);
      92            3 :   auto vit = row.find(out_col.column_id());
      93            3 :   ASSERT_NE(vit, row.end());
      94            3 :   EXPECT_EQ(vit->second.int64_value(), 7);
      95            3 : }
      96              : 
      97              : INSTANTIATE_TEST_SUITE_P(
      98              :     PrivacyAggregateScanKinds,
      99              :     PrivacyAggregateStubTest,
     100              :     ::testing::Values(
     101              :         ::googlesql::RESOLVED_ANONYMIZED_AGGREGATE_SCAN,
     102              :         ::googlesql::RESOLVED_DIFFERENTIAL_PRIVACY_AGGREGATE_SCAN,
     103              :         ::googlesql::RESOLVED_AGGREGATION_THRESHOLD_AGGREGATE_SCAN));
     104              : 
     105              : }  // namespace
     106              : }  // namespace semantic
     107              : }  // namespace engine
     108              : }  // namespace backend
     109              : }  // namespace bigquery_emulator