From 86a15331d164487291404bd4f2c0a964bd9b117e Mon Sep 17 00:00:00 2001
From: josipmrden <josip.mrden@external-basf.com>
Date: Mon, 4 Jul 2022 16:49:23 +0200
Subject: [PATCH] Added saving of labels to AuthQuery

---
 src/query/frontend/ast/ast.lcp                |  6 +++--
 .../frontend/ast/cypher_main_visitor.cpp      | 27 +++++++++++++++++--
 .../frontend/ast/cypher_main_visitor.hpp      |  5 ++++
 .../opencypher/grammar/MemgraphCypher.g4      |  2 +-
 4 files changed, 35 insertions(+), 5 deletions(-)

diff --git a/src/query/frontend/ast/ast.lcp b/src/query/frontend/ast/ast.lcp
index 836857a4e..a618adf69 100644
--- a/src/query/frontend/ast/ast.lcp
+++ b/src/query/frontend/ast/ast.lcp
@@ -2239,6 +2239,7 @@ cpp<#
    (user "std::string" :scope :public)
    (role "std::string" :scope :public)
    (user-or-role "std::string" :scope :public)
+   (labels "std::vector<std::string>" :scope :public)
    (password "Expression *" :initval "nullptr" :scope :public
              :slk-save #'slk-save-ast-pointer
              :slk-load (slk-load-ast-pointer "Expression"))
@@ -2264,13 +2265,14 @@ cpp<#
     #>cpp
     AuthQuery(Action action, std::string user, std::string role,
               std::string user_or_role, Expression *password,
-              std::vector<Privilege> privileges)
+              std::vector<Privilege> privileges, std::vector<std::string> labels)
         : action_(action),
           user_(user),
           role_(role),
           user_or_role_(user_or_role),
           password_(password),
-          privileges_(privileges) {}
+          privileges_(privileges),
+          labels_(labels) {}
     cpp<#)
   (:private
     #>cpp
diff --git a/src/query/frontend/ast/cypher_main_visitor.cpp b/src/query/frontend/ast/cypher_main_visitor.cpp
index 62a7aa338..162552b0f 100644
--- a/src/query/frontend/ast/cypher_main_visitor.cpp
+++ b/src/query/frontend/ast/cypher_main_visitor.cpp
@@ -1285,7 +1285,11 @@ antlrcpp::Any CypherMainVisitor::visitGrantPrivilege(MemgraphCypher::GrantPrivil
   auth->user_or_role_ = ctx->userOrRole->accept(this).as<std::string>();
   if (ctx->privilegeList()) {
     for (auto *privilege : ctx->privilegeList()->privilege()) {
-      auth->privileges_.push_back(privilege->accept(this));
+      if (privilege->LABELS()) {
+        auth->labels_ = privilege->labelList()->accept(this).as<std::vector<std::string>>();
+      } else {
+        auth->privileges_.push_back(privilege->accept(this));
+      }
     }
   } else {
     /* grant all privileges */
@@ -1330,6 +1334,22 @@ antlrcpp::Any CypherMainVisitor::visitRevokePrivilege(MemgraphCypher::RevokePriv
   return auth;
 }
 
+/**
+ * @return AuthQuery*
+ */
+antlrcpp::Any CypherMainVisitor::visitLabelList(MemgraphCypher::LabelListContext *ctx) {
+  std::vector<std::string> labels;
+  for (auto *label : ctx->label()) {
+    if (label->ASTERISK()) {
+      labels.push_back("*");
+    } else {
+      labels.push_back(label->symbolicName()->accept(this).as<std::string>());
+    }
+  }
+
+  return labels;
+}
+
 /**
  * @return AuthQuery::Privilege
  */
@@ -1355,7 +1375,10 @@ antlrcpp::Any CypherMainVisitor::visitPrivilege(MemgraphCypher::PrivilegeContext
   if (ctx->MODULE_READ()) return AuthQuery::Privilege::MODULE_READ;
   if (ctx->MODULE_WRITE()) return AuthQuery::Privilege::MODULE_WRITE;
   if (ctx->WEBSOCKET()) return AuthQuery::Privilege::WEBSOCKET;
-  if (ctx->LABELS()) return AuthQuery::Privilege::LABELS;
+  if (ctx->LABELS()) {
+    // fill labels in authquery
+    return AuthQuery::Privilege::LABELS;
+  }
   LOG_FATAL("Should not get here - unknown privilege!");
 }
 
diff --git a/src/query/frontend/ast/cypher_main_visitor.hpp b/src/query/frontend/ast/cypher_main_visitor.hpp
index 2a6b8ff5e..2fdd22b88 100644
--- a/src/query/frontend/ast/cypher_main_visitor.hpp
+++ b/src/query/frontend/ast/cypher_main_visitor.hpp
@@ -473,6 +473,11 @@ class CypherMainVisitor : public antlropencypher::MemgraphCypherBaseVisitor {
    */
   antlrcpp::Any visitPrivilege(MemgraphCypher::PrivilegeContext *ctx) override;
 
+  /**
+   * @return AuthQuery::LabelList
+   */
+  antlrcpp::Any visitLabelList(MemgraphCypher::LabelListContext *ctx) override;
+
   /**
    * @return AuthQuery*
    */
diff --git a/src/query/frontend/opencypher/grammar/MemgraphCypher.g4 b/src/query/frontend/opencypher/grammar/MemgraphCypher.g4
index 551b165e8..12d54a916 100644
--- a/src/query/frontend/opencypher/grammar/MemgraphCypher.g4
+++ b/src/query/frontend/opencypher/grammar/MemgraphCypher.g4
@@ -262,7 +262,7 @@ privilegeList : privilege ( ',' privilege )* ;
 
 labelList : label ( ',' label )* ;
 
-label : ( '*' | StringLiteral ) ;
+label : ( '*' | symbolicName ) ;
 
 showPrivileges : SHOW PRIVILEGES FOR userOrRole=userOrRoleName ;