fix: #64 by update current space from custom nGQL

nebula-contrib · Oct 25, 2022 · f5e4134 · f5e4134
1 parent 806b9f4
commit f5e4134
Show file tree

Hide file tree

Showing 4 changed files with 59 additions and 4 deletions.
diff --git a/ngbatis-demo/src/main/java/ye/weicheng/ngbatis/demo/repository/TestRepository.java b/ngbatis-demo/src/main/java/ye/weicheng/ngbatis/demo/repository/TestRepository.java
@@ -58,5 +58,9 @@ public interface TestRepository extends NebulaDaoBasic<Person, String> {
   ResultSet testMulti();
 
   Map<String, Object> selectMapWhenNull();
+
+  void testSpaceSwitchStep1();
+
+  Integer testSpaceSwitchStep2();
 
 }
diff --git a/ngbatis-demo/src/main/resources/mapper/TestRepository.xml b/ngbatis-demo/src/main/resources/mapper/TestRepository.xml
@@ -93,6 +93,17 @@
         WHERE 1 == 2 
         RETURN n
     </select>
+
+    <select id="testSpaceSwitchStep1">
+        use cmqa;
+        INSERT VERTEX paragraph( name ) VALUES -31415926:('spaceSwitch');
+    </select>
+
+    <select id="testSpaceSwitchStep2" space="cmqa">
+        MATCH (n: paragraph)
+        WHERE id(n) == -31415926
+        RETURN count(n)
+    </select>
 
 </mapper>
 
diff --git a/ngbatis-demo/src/test/java/ye/weicheng/ngbatis/demo/NgbatisDemoApplicationTests.java b/ngbatis-demo/src/test/java/ye/weicheng/ngbatis/demo/NgbatisDemoApplicationTests.java
@@ -161,4 +161,20 @@ public void selectMapWhenNull() {
     Assert.isTrue(result == null);;
   }
 
+  @Test
+  public void testSpaceSwitch() {
+    for (int i = 0; i < 30; i ++) {
+      long l = System.currentTimeMillis();
+      int mod = (int) (l % 3);
+      switch (mod) {
+        case 0: repository.testSpaceSwitchStep1();
+          break;
+        case 1: repository.testSpaceSwitchStep2();
+          break;
+        case 2: repository.selectMapWhenNull();
+          break;
+      }
+    }
+  }
+
 }
diff --git a/src/main/java/org/nebula/contrib/ngbatis/proxy/MapperProxy.java b/src/main/java/org/nebula/contrib/ngbatis/proxy/MapperProxy.java
@@ -15,6 +15,8 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
 import org.nebula.contrib.ngbatis.ArgNameFormatter;
 import org.nebula.contrib.ngbatis.Env;
 import org.nebula.contrib.ngbatis.ResultResolver;
@@ -238,11 +240,33 @@ public static ResultSet executeWithParameter(ClassModel cm, MethodModel mm, Stri
   private static String qlWithSpace(LocalSession localSession, String gql, String currentSpace) {
     gql = gql.trim();
     String sessionSpace = localSession.getCurrentSpace();
-    if (Objects.equals(sessionSpace, currentSpace)) {
-      return String.format("\n\t\t%s", gql);
+    String qlWithSpace = Objects.equals(sessionSpace, currentSpace)
+        ? String.format("\n\t\t%s", gql)
+        : String.format("USE %s;\n\t\t%s", currentSpace, gql);
+
+    setNewSpace(localSession, gql, currentSpace);
+    return qlWithSpace;
+  }
+
+  private static void setNewSpace(LocalSession localSession, String gql, String currentSpace) {
+    String lastSpace = findLastSpaceFromGql(gql);
+    if (lastSpace == null) {
+      localSession.setCurrentSpace(currentSpace);
+    } else {
+      localSession.setCurrentSpace(lastSpace);
+    }
+  }
+
+  private static final String PATTERN = "(?s).*(\\bUSE\\s+(\\S+);).*";
+
+  private static String findLastSpaceFromGql(String gql) {
+    Pattern r = Pattern.compile(PATTERN, Pattern.MULTILINE | Pattern.CASE_INSENSITIVE);
+    Matcher m = r.matcher(gql);
+    if (!m.matches()) {
+      return null;
+    } else {
+      return m.group(2);
     }
-    localSession.setCurrentSpace(currentSpace);
-    return String.format("USE %s;\n\t\t%s", currentSpace, gql);
   }
 
   /**