Check if aggregation is used in right clause

Reviewers: florijan, mislav.bradac

Reviewed By: mislav.bradac

Subscribers: pullbot

Differential Revision: https://phabricator.memgraph.io/D274

src/query/frontend/semantic/symbol_generator.cpp

@@ -31,15 +31,24 @@ auto SymbolGenerator::GetOrCreateSymbol(const std::string &name,
 void SymbolGenerator::Visit(Create &create) { scope_.in_create = true; }
 void SymbolGenerator::PostVisit(Create &create) { scope_.in_create = false; }
 
+void SymbolGenerator::Visit(Return &ret) {
+  scope_.in_return = true;
+}
 void SymbolGenerator::PostVisit(Return &ret) {
   for (auto &named_expr : ret.named_expressions_) {
     // Named expressions establish bindings for expressions which come after
     // return, but not for the expressions contained inside.
     symbol_table_[*named_expr] = CreateSymbol(named_expr->name_);
   }
+  scope_.in_return = false;
 }
 
-void SymbolGenerator::SetWithSymbols(With &with) {
+bool SymbolGenerator::PreVisit(With &with) {
+  scope_.in_with = true;
+  for (auto &expr : with.named_expressions_) {
+    expr->Accept(*this);
+  }
+  scope_.in_with = false;
   // WITH clause removes declarations of all the previous variables and declares
   // only those established through named expressions. New declarations must not
   // be visible inside named expressions themselves.
@@ -47,26 +56,8 @@ void SymbolGenerator::SetWithSymbols(With &with) {
   for (auto &named_expr : with.named_expressions_) {
     symbol_table_[*named_expr] = CreateSymbol(named_expr->name_);
   }
-}
-
-void SymbolGenerator::Visit(With &with) {
-  scope_.with = &with;
-}
-
-void SymbolGenerator::Visit(Where &where) {
-  if (scope_.with) {
-    // New symbols must be visible in WHERE clause, so this must be done here
-    // and not in PostVisit(With&).
-    SetWithSymbols(*scope_.with);
-  }
-}
-
-void SymbolGenerator::PostVisit(With &with) {
-  if (!with.where_) {
-    // This wasn't done when visiting Where, so do it here.
-    SetWithSymbols(with);
-  }
-  scope_.with = nullptr;
+  if (with.where_) with.where_->Accept(*this);
+  return false;  // We handled the traversal ourselves.
 }
 
 // Expressions
@@ -104,12 +95,20 @@ void SymbolGenerator::Visit(Identifier &ident) {
 }
 
 void SymbolGenerator::Visit(Aggregation &aggr) {
-  // Create a virtual symbol for aggregation result.
-  symbol_table_[aggr] = symbol_table_.CreateSymbol("");
+  // Check if the aggregation can be used in this context. This check should
+  // probably move to a separate phase, which checks if the query is well
+  // formed.
+  if (!scope_.in_return && !scope_.in_with) {
+    throw SemanticException(
+        "Aggregation functions are only allowed in WITH and RETURN");
+  }
   if (scope_.in_aggregation) {
     throw SemanticException(
-        "Using aggregate functions inside aggregate functions is not allowed");
+        "Using aggregation functions inside aggregation functions is not "
+        "allowed");
   }
+  // Create a virtual symbol for aggregation result.
+  symbol_table_[aggr] = symbol_table_.CreateSymbol("");
   scope_.in_aggregation = true;
 }
 

src/query/frontend/semantic/symbol_generator.hpp

@@ -20,16 +20,16 @@ class SymbolGenerator : public TreeVisitorBase {
  public:
   SymbolGenerator(SymbolTable &symbol_table) : symbol_table_(symbol_table) {}
 
+  using TreeVisitorBase::PreVisit;
   using TreeVisitorBase::Visit;
   using TreeVisitorBase::PostVisit;
 
   // Clauses
   void Visit(Create &) override;
   void PostVisit(Create &) override;
+  void Visit(Return &) override;
   void PostVisit(Return &) override;
-  void Visit(With &) override;
-  void PostVisit(With &) override;
-  void Visit(Where &) override;
+  bool PreVisit(With &) override;
 
   // Expressions
   void Visit(Identifier &) override;
@@ -59,8 +59,8 @@ class SymbolGenerator : public TreeVisitorBase {
     bool in_edge_atom{false};
     bool in_property_map{false};
     bool in_aggregation{false};
-    // Pointer to With clause if we are inside it, otherwise nullptr.
-    With *with{nullptr};
+    bool in_return{false};
+    bool in_with{false};
     std::map<std::string, Symbol> symbols;
   };
 
@@ -76,9 +76,6 @@ class SymbolGenerator : public TreeVisitorBase {
   auto GetOrCreateSymbol(const std::string &name,
                          Symbol::Type type = Symbol::Type::Any);
 
-  // Clear old symbol bindings and establish new from WITH clause.
-  void SetWithSymbols(With &with);
-
   SymbolTable &symbol_table_;
   Scope scope_;
 };

tests/unit/query_semantic.cpp

@@ -446,6 +446,20 @@ TEST(TestSymbolGenerator, NestedAggregation) {
   EXPECT_THROW(query->Accept(symbol_generator), SemanticException);
 }
 
+TEST(TestSymbolGenerator, WrongAggregationContext) {
+  // Test MATCH (n) WITH n.prop AS prop WHERE SUM(prop) < 42
+  Dbms dbms;
+  auto dba = dbms.active();
+  auto prop = dba->property("prop");
+  AstTreeStorage storage;
+  auto query = QUERY(MATCH(PATTERN(NODE("n"))),
+                     WITH(PROPERTY_LOOKUP("n", prop), AS("prop")),
+                     WHERE(LESS(SUM(IDENT("prop")), LITERAL(42))));
+  SymbolTable symbol_table;
+  SymbolGenerator symbol_generator(symbol_table);
+  EXPECT_THROW(query->Accept(symbol_generator), SemanticException);
+}
+
 TEST(TestSymbolGenerator, MatchPropCreateNodeProp) {
   // Test MATCH (n) CREATE (m {prop: n.prop})
   Dbms dbms;