Forbid nonaliased expressions in WITH

Reviewers: teon.banek, florijan

Reviewed By: teon.banek

Subscribers: pullbot

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

src/query/frontend/ast/cypher_main_visitor.cpp

@@ -187,14 +187,17 @@ antlrcpp::Any CypherMainVisitor::visitReturnItems(
 antlrcpp::Any CypherMainVisitor::visitReturnItem(
     CypherParser::ReturnItemContext *ctx) {
   auto *named_expr = storage_.Create<NamedExpression>();
+  named_expr->expression_ = ctx->expression()->accept(this);
   if (ctx->variable()) {
     named_expr->name_ =
         std::string(ctx->variable()->accept(this).as<std::string>());
   } else {
+    if (in_with_ && !dynamic_cast<Identifier *>(named_expr->expression_)) {
+      throw SemanticException("Only variables can be non aliased in with");
+    }
     // TODO: Should we get this by text or some escaping is needed?
     named_expr->name_ = std::string(ctx->getText());
   }
-  named_expr->expression_ = ctx->expression()->accept(this);
   return named_expr;
 }
 
@@ -929,7 +932,9 @@ antlrcpp::Any CypherMainVisitor::visitPropertyExpression(
 
 antlrcpp::Any CypherMainVisitor::visitWith(CypherParser::WithContext *ctx) {
   auto *with = storage_.Create<With>();
+  in_with_ = true;
   with->body_ = ctx->returnBody()->accept(this);
+  in_with_ = false;
   if (ctx->DISTINCT()) {
     with->body_.distinct = true;
   }

src/query/frontend/ast/cypher_main_visitor.hpp

@@ -482,6 +482,10 @@ class CypherMainVisitor : public antlropencypher::CypherBaseVisitor {
   std::vector<Identifier **> anonymous_identifiers;
   AstTreeStorage storage_;
   Query *query_ = nullptr;
+  // All return items which are not variables must be aliased in with.
+  // We use this variable in visitReturnItem to check if we are in with or
+  // return.
+  bool in_with_ = false;
 };
 }
 }

tests/unit/cypher_main_visitor.cpp

@@ -857,6 +857,23 @@ TEST(CypherMainVisitorTest, With) {
   ASSERT_EQ(identifier->name_, "n");
 }
 
+TEST(CypherMainVisitorTest, WithNonAliasedExpression) {
+  ASSERT_THROW(AstGenerator("WITH n.x RETURN 1"), SemanticException);
+}
+
+TEST(CypherMainVisitorTest, WithNonAliasedVariable) {
+  AstGenerator ast_generator("WITH n RETURN 1");
+  auto *query = ast_generator.query_;
+  ASSERT_EQ(query->clauses_.size(), 2U);
+  auto *with = dynamic_cast<With *>(query->clauses_[0]);
+  ASSERT_TRUE(with);
+  ASSERT_EQ(with->body_.named_expressions.size(), 1U);
+  auto *named_expr = with->body_.named_expressions[0];
+  ASSERT_EQ(named_expr->name_, "n");
+  auto *identifier = dynamic_cast<Identifier *>(named_expr->expression_);
+  ASSERT_EQ(identifier->name_, "n");
+}
+
 TEST(CypherMainVisitorTest, WithDistinct) {
   AstGenerator ast_generator("WITH DISTINCT n AS m RETURN 1");
   auto *query = ast_generator.query_;