Merge pull request #332 from CorvusYe/master

feat: entity direct query supports specifying different spaces.
nebula-contrib · Dec 18, 2024 · bb39d71 · bb39d71
2 parents e7330eb + 1aef6f2
commit bb39d71
Show file tree

Hide file tree

Showing 5 changed files with 121 additions and 27 deletions.
diff --git a/src/main/java/org/nebula/contrib/ngbatis/Env.java b/src/main/java/org/nebula/contrib/ngbatis/Env.java
@@ -4,11 +4,16 @@
 //
 // This source code is licensed under Apache 2.0 License.
 
+import static org.apache.commons.lang3.StringUtils.isBlank;
+import static org.nebula.contrib.ngbatis.proxy.MapperProxy.ENV;
+import static org.nebula.contrib.ngbatis.utils.ReflectUtil.spaceFromEntity;
+
 import com.alibaba.fastjson.parser.ParserConfig;
 import com.vesoft.nebula.client.graph.SessionPool;
 import com.vesoft.nebula.client.graph.net.Session;
 import org.nebula.contrib.ngbatis.base.GraphBaseExt;
 import org.nebula.contrib.ngbatis.config.ParseCfgProps;
+import org.nebula.contrib.ngbatis.exception.ResourceLoadException;
 import org.nebula.contrib.ngbatis.models.MapperContext;
 import org.nebula.contrib.ngbatis.proxy.MapperProxy;
 import org.slf4j.Logger;
@@ -122,6 +127,55 @@ public Session openSession() {
     }
   }
 
+  /**
+   * 通过实体类获取对应的space。
+   * 支持占位符，并从配置信息中读取。
+   * 
+   * @param entityType 实体类
+   * @return space
+   */
+  public static String spaceFromConfig(Class<?> entityType) {
+    return spaceFromConfig(entityType, ENV.getContext());
+  }
+
+  public static String spaceFromConfig(Class<?> entityType, ApplicationContext context) {
+    String space = spaceFromEntity(entityType);
+    if (isBlank(space)) return null;
+    space = tryResolvePlaceholder(space, context);
+    return space;
+  }
+
+  /**
+   * 利用Spring Environment 解析注解的值，用于 @Space 的 name 属性解析
+   * @author <a href="https://github.com/charle004">Charle004</a>
+   * @param value 需要解析的值，可能是带占位符的 ${xx.xx} ，也可以是固定的字符串
+   * @return resolveResult 解析结果
+   * @throws IllegalArgumentException 当配置了 ${xx.xx} 占位符，且spring配置文件中未指定该配置时抛出
+   */
+  public static String tryResolvePlaceholder(String value) {
+    return tryResolvePlaceholder(value, ENV.getContext());
+  }
+
+  /**
+   * 利用Spring Environment 解析注解的值，用于 @Space 的 name 属性解析
+   * @author <a href="https://github.com/charle004">Charle004</a>
+   * @param configKey 需要解析的值，可能是带占位符的 ${xx.xx} ，也可以是固定的字符串
+   * @return resolveResult 解析结果
+   * @throws IllegalArgumentException 当配置了 ${xx.xx} 占位符，且spring配置文件中未指定该配置时抛出
+   */
+  public static String tryResolvePlaceholder(String configKey, ApplicationContext context) {
+    String resolveResult = configKey;
+    if (null != context) {
+      try {
+        resolveResult = context.getEnvironment().resolveRequiredPlaceholders(configKey);
+      } catch (IllegalArgumentException e) {
+        throw new ResourceLoadException(
+          "name ( " + configKey + " ) of @Space missing configurable value"
+        );
+      }
+    }
+    return resolveResult;
+  }
 
   public String getUsername() {
     return username;

diff --git a/src/main/java/org/nebula/contrib/ngbatis/base/GraphBase.java b/src/main/java/org/nebula/contrib/ngbatis/base/GraphBase.java
@@ -1,6 +1,7 @@
 package org.nebula.contrib.ngbatis.base;
 
 import java.util.Map;
+import org.nebula.contrib.ngbatis.models.MapperContext;
 
 /**
  * 实体基类
@@ -9,6 +10,10 @@
  */
 
 public abstract class GraphBase {
+  public GraphBase() {
+    Map<String, Class<?>> tagTypeMapping = MapperContext.newInstance().getTagTypeMapping();
+    tagTypeMapping.put(getEntityTypeName(), this.getClass());
+  }
 
   protected abstract String getEntityTypeName();
 

diff --git a/src/main/java/org/nebula/contrib/ngbatis/base/GraphBaseExt.java b/src/main/java/org/nebula/contrib/ngbatis/base/GraphBaseExt.java
@@ -1,5 +1,6 @@
 package org.nebula.contrib.ngbatis.base;
 
+import static org.nebula.contrib.ngbatis.Env.spaceFromConfig;
 import static org.nebula.contrib.ngbatis.utils.ReflectUtil.getNameByColumn;
 import static org.nebula.contrib.ngbatis.utils.ReflectUtil.getValue;
 
@@ -33,7 +34,7 @@ public static ResultSet executeGql(String textTpl,
     //从env中获取本地会话调度器
     SessionDispatcher dispatcher = ENV.getDispatcher();
     //从env中获取space
-    String currentSpace = ENV.getSpace();
+    String currentSpace = getSpace(m1);
 
     ArgsResolver argsResolver = ENV.getArgsResolver();
 
@@ -55,6 +56,38 @@ public static ResultSet executeGql(String textTpl,
     }
     return result;
   }
+
+  /**
+   * 获取当前space，支持 @Space 跟 @Table 注解
+   *     如果均未指定，则返回默认space
+   *     
+   * @return 当前space
+   */
+  public static String getSpace(Map<String, Object> m1) {
+    Object edgeType = m1.get("edgeType");
+    Object tag = m1.get("tag");
+    String entityTypeName = edgeType != null ? edgeType.toString()
+            : tag != null ? tag.toString()
+              : null;
+
+    String defaultSpace = ENV.getSpace();
+    if (entityTypeName == null) {
+      return defaultSpace;
+    }
+
+    Map<String, Class<?>> entityTypeMapping = ENV.getMapperContext().getTagTypeMapping();
+    Class<?> entityType = entityTypeMapping.get(entityTypeName);
+    boolean isGraphBase = GraphBase.class.isAssignableFrom(entityType);
+    if (!isGraphBase) {
+      return defaultSpace;
+    }
+
+    String space = spaceFromConfig(entityType);
+    if (space != null) {
+      return space;
+    }
+    return ENV.getSpace();
+  }
 
   /**
    * 处理查询结果

diff --git a/src/main/java/org/nebula/contrib/ngbatis/io/MapperResourceLoader.java b/src/main/java/org/nebula/contrib/ngbatis/io/MapperResourceLoader.java
@@ -31,6 +31,7 @@
 import org.jsoup.nodes.Node;
 import org.jsoup.nodes.TextNode;
 import org.jsoup.select.Elements;
+import org.nebula.contrib.ngbatis.Env;
 import org.nebula.contrib.ngbatis.annotations.Space;
 import org.nebula.contrib.ngbatis.annotations.TimeLog;
 import org.nebula.contrib.ngbatis.config.ParseCfgProps;
@@ -154,38 +155,17 @@ private void setClassModelBySpaceAnnotation(ClassModel cm) {
         return;
       }
       String spaceClassName = genericTypes[0].getTypeName();
-      Space annotation = Class.forName(spaceClassName).getAnnotation(Space.class);
-      if (null != annotation) {
-        String spaceName = tryResolvePlaceholder(annotation.name());
-        if (!spaceName.equals("")) {
-          cm.setSpace(spaceName);
-        }
+      Class<?> entityType = Class.forName(spaceClassName);
+      String space = ReflectUtil.spaceFromEntity(entityType);
+      if (isNotBlank(space)) {
+        space = Env.tryResolvePlaceholder(space, applicationContext);
+        cm.setSpace(space);
       }
     } catch (ClassNotFoundException e) {
       e.printStackTrace();
     }
   }
 
-  /**
-   * 利用Spring Environment 解析注解的值，用于 @Space 的 name 属性解析
-   * @param value 需要解析的值，可能是带占位符的 ${xx.xx} ，也可以是固定的字符串
-   * @return resolveResult 解析结果
-   * @throws IllegalArgumentException 当配置了 ${xx.xx} 占位符，且spring配置文件中未指定该配置时抛出
-   */
-  private String tryResolvePlaceholder(String value) {
-    String resolveResult = value;
-    if (null != applicationContext) {
-      try {
-        resolveResult = applicationContext.getEnvironment().resolveRequiredPlaceholders(value);
-      } catch (IllegalArgumentException e) {
-        throw new ResourceLoadException(
-          "name ( " + value + " ) of @Space missing configurable value"
-        );
-      }
-    }
-    return resolveResult;
-  }
-
   /**
    * 解析 一个 XXXDao 的多个方法。
    *

diff --git a/src/main/java/org/nebula/contrib/ngbatis/utils/ReflectUtil.java b/src/main/java/org/nebula/contrib/ngbatis/utils/ReflectUtil.java
@@ -24,6 +24,7 @@
 import javax.persistence.Id;
 import javax.persistence.Table;
 import javax.persistence.Transient;
+import org.nebula.contrib.ngbatis.annotations.Space;
 import org.nebula.contrib.ngbatis.annotations.base.EdgeType;
 import org.nebula.contrib.ngbatis.annotations.base.GraphId;
 import org.nebula.contrib.ngbatis.annotations.base.Tag;
@@ -623,4 +624,25 @@ public static boolean isGraphId(Field field) {
     }
     return false;
   }
+
+  /**
+   * 从实体类获取 space，并解析占位符
+    */
+  public static String spaceFromEntity(Class<?> entityType) {
+    boolean hasSpace = entityType.isAnnotationPresent(Space.class);
+    String space = null;
+    if (hasSpace) {
+      Space spaceAnnotation = entityType.getAnnotation(Space.class);
+      space = spaceAnnotation.name();
+    }
+
+    boolean hasTable = entityType.isAnnotationPresent(Table.class);
+    if (hasTable) {
+      Table tableAnnotation = entityType.getAnnotation(Table.class);
+      space = tableAnnotation.schema();
+    }
+
+    return space;
+  }
+
 }