From 8dc3153fdefbb756fa184f673f646ac842706fc8 Mon Sep 17 00:00:00 2001
From: antonio2368 <antonio2368@users.noreply.github.com>
Date: Fri, 24 Sep 2021 13:14:05 +0200
Subject: [PATCH] Don't strip queries after execute for CREATE TRIGGER only
 (#245)

---
 src/query/frontend/stripped.cpp | 23 ++++++++++++++++----
 tests/unit/stripped.cpp         | 38 +++++++++++++++++++++++++++++++++
 2 files changed, 57 insertions(+), 4 deletions(-)

diff --git a/src/query/frontend/stripped.cpp b/src/query/frontend/stripped.cpp
index 9a6c31959..547bdb3c8 100644
--- a/src/query/frontend/stripped.cpp
+++ b/src/query/frontend/stripped.cpp
@@ -3,6 +3,7 @@
 #include <cctype>
 #include <cstdint>
 #include <iostream>
+#include <span>
 #include <string>
 #include <vector>
 
@@ -60,11 +61,25 @@ StrippedQuery::StrippedQuery(const std::string &query) : original_(query) {
     tokens.emplace_back(token, original_.substr(i, len));
     i += len;
 
-    // if we notice execute, we create a trigger which has defined statements
-    // the statements will be parsed separately later on so we skip it for now
+    // If we notice execute, we possibly create a trigger which has defined statements.
+    // The statements will be parsed separately later on so we skip it for now.
     if (utils::IEquals(tokens.back().second, "execute")) {
-      unstripped_chunk = original_.substr(i);
-      break;
+      // check if it's CREATE TRIGGER query
+      std::span token_span{tokens};
+
+      // query could start with spaces and/or comments
+      if (token_span.front().first == Token::SPACE) {
+        token_span = token_span.subspan(1);
+      }
+
+      // we need to check that first and third elements are correct keywords
+      // CREATE<SPACE>TRIGGER<SPACE>trigger-name...EXECUTE
+      // trigger-name (5th element) can also be "execute" so we verify that the size is larger than 5
+      if (token_span.size() > 5 && utils::IEquals(token_span[0].second, "create") &&
+          utils::IEquals(token_span[2].second, "trigger")) {
+        unstripped_chunk = original_.substr(i);
+        break;
+      }
     }
   }
 
diff --git a/tests/unit/stripped.cpp b/tests/unit/stripped.cpp
index 3c2883016..7b54a3c39 100644
--- a/tests/unit/stripped.cpp
+++ b/tests/unit/stripped.cpp
@@ -353,4 +353,42 @@ TEST(QueryStripper, QuerySemicolonEndingQuery2) {
   StrippedQuery stripped("RETURN 42   ;");
   EXPECT_THAT(stripped.named_expressions(), UnorderedElementsAre(Pair(1, "42")));
 }
+
+TEST(QueryStripper, CreateTriggerQuery) {
+  constexpr std::string_view execute_query{
+      "      MATCH  (execute:Node)    RETURN / *test comment */    execute \"test\""};
+  {
+    SCOPED_TRACE("Everything after EXECUTE keyword in CREATE TRIGGER should not be stripped");
+
+    {
+      SCOPED_TRACE("Query starting with CREATE keyword");
+      StrippedQuery stripped(
+          fmt::format("CREATE TRIGGER execute  /*test*/ ON CREATE BEFORE COMMIT EXECUTE{}", execute_query));
+      EXPECT_EQ(stripped.query(),
+                fmt::format("CREATE TRIGGER execute ON CREATE BEFORE COMMIT EXECUTE {}", execute_query));
+    }
+
+    {
+      SCOPED_TRACE("Query starting with comments and spaces");
+      StrippedQuery stripped(fmt::format(
+          "/*comment*/    \n\n //other comment\nCREATE TRIGGER execute AFTER COMMIT EXECUTE{}", execute_query));
+      EXPECT_EQ(stripped.query(), fmt::format("CREATE TRIGGER execute AFTER COMMIT EXECUTE {}", execute_query));
+    }
+
+    {
+      SCOPED_TRACE("Query with comments and spaces between CREATE and TRIGGER");
+      StrippedQuery stripped(fmt::format(
+          "/*comment*/    \n\n //other comment\nCREATE //some comment \n   TRIGGER execute AFTER COMMIT EXECUTE{}",
+          execute_query));
+      EXPECT_EQ(stripped.query(), fmt::format("CREATE TRIGGER execute AFTER COMMIT EXECUTE {}", execute_query));
+    }
+  }
+  {
+    SCOPED_TRACE("Execute keyword should still be allowed in other queries");
+    StrippedQuery stripped("MATCH (execute:Node)   //comment \n  RETURN  /* test comment */  execute");
+    EXPECT_EQ(stripped.query(), "MATCH ( execute : Node ) RETURN execute");
+    EXPECT_THAT(stripped.named_expressions(), UnorderedElementsAre(Pair(7, "execute")));
+  }
+}
+
 }  // namespace