From 2bc607b06080193976940e58f1996741ad2c8c65 Mon Sep 17 00:00:00 2001
From: Michael Hunger <github@jexp.de>
Date: Wed, 10 Aug 2011 15:57:34 +0200
Subject: [PATCH] minor repository refactoring

---
 .../repository/GraphEntityInformation.java    |  1 +
 .../GraphMetamodelEntityInformation.java      | 20 ++++++++++++-------
 .../repository/GraphRepositoryFactory.java    |  9 +++------
 3 files changed, 17 insertions(+), 13 deletions(-)

diff --git a/spring-data-neo4j/src/main/java/org/springframework/data/neo4j/repository/GraphEntityInformation.java b/spring-data-neo4j/src/main/java/org/springframework/data/neo4j/repository/GraphEntityInformation.java
index 3ea397c7d..132847016 100644
--- a/spring-data-neo4j/src/main/java/org/springframework/data/neo4j/repository/GraphEntityInformation.java
+++ b/spring-data-neo4j/src/main/java/org/springframework/data/neo4j/repository/GraphEntityInformation.java
@@ -28,4 +28,5 @@ public interface GraphEntityInformation<S extends PropertyContainer, T extends G
 
     boolean isNodeEntity();
     boolean isPartialEntity();
+    boolean isRelationshipEntity();
 }
diff --git a/spring-data-neo4j/src/main/java/org/springframework/data/neo4j/repository/GraphMetamodelEntityInformation.java b/spring-data-neo4j/src/main/java/org/springframework/data/neo4j/repository/GraphMetamodelEntityInformation.java
index 7a5120369..2e0605c46 100644
--- a/spring-data-neo4j/src/main/java/org/springframework/data/neo4j/repository/GraphMetamodelEntityInformation.java
+++ b/spring-data-neo4j/src/main/java/org/springframework/data/neo4j/repository/GraphMetamodelEntityInformation.java
@@ -23,6 +23,7 @@ package org.springframework.data.neo4j.repository;
 
 import org.neo4j.graphdb.PropertyContainer;
 import org.springframework.data.neo4j.annotation.NodeEntity;
+import org.springframework.data.neo4j.annotation.RelationshipEntity;
 import org.springframework.data.neo4j.core.GraphBacked;
 import org.springframework.data.neo4j.core.NodeBacked;
 import org.springframework.data.neo4j.core.RelationshipBacked;
@@ -31,26 +32,31 @@ import org.springframework.data.repository.core.support.AbstractEntityInformatio
 
 public class GraphMetamodelEntityInformation<S extends PropertyContainer, T extends GraphBacked<S>> extends AbstractEntityInformation<T,Long> implements GraphEntityInformation<S,T> {
 
-    private final boolean isNodeEntity;
-    private final boolean isPartialEntity;
     private GraphDatabaseContext graphDatabaseContext;
+    private final RelationshipEntity relationshipEntity;
+    private final NodeEntity nodeEntity;
 
     public GraphMetamodelEntityInformation(Class domainClass, GraphDatabaseContext graphDatabaseContext) {
         super(domainClass);
         this.graphDatabaseContext = graphDatabaseContext;
-        NodeEntity nodeEntity = getJavaType().getAnnotation(NodeEntity.class);
-        isNodeEntity = nodeEntity!=null;
-        isPartialEntity = isNodeEntity && nodeEntity.partial();
+        nodeEntity = getJavaType().getAnnotation(NodeEntity.class);
+        relationshipEntity = getJavaType().getAnnotation(RelationshipEntity.class);
+
     }
 
     @Override
     public boolean isNodeEntity() {
-        return isNodeEntity;
+        return nodeEntity!=null;
     }
 
     @Override
     public boolean isPartialEntity() {
-        return isPartialEntity;
+        return nodeEntity!=null && nodeEntity.partial();
+    }
+
+    @Override
+    public boolean isRelationshipEntity() {
+        return relationshipEntity!=null;
     }
 
     @Override
diff --git a/spring-data-neo4j/src/main/java/org/springframework/data/neo4j/repository/GraphRepositoryFactory.java b/spring-data-neo4j/src/main/java/org/springframework/data/neo4j/repository/GraphRepositoryFactory.java
index 71e298436..b150c1713 100644
--- a/spring-data-neo4j/src/main/java/org/springframework/data/neo4j/repository/GraphRepositoryFactory.java
+++ b/spring-data-neo4j/src/main/java/org/springframework/data/neo4j/repository/GraphRepositoryFactory.java
@@ -92,12 +92,9 @@ public class GraphRepositoryFactory extends RepositoryFactorySupport {
     @Override
     protected Class<?> getRepositoryBaseClass(RepositoryMetadata repositoryMetadata) {
         Class<?> domainClass = repositoryMetadata.getDomainClass();
-        if (findAnnotation(domainClass, NodeEntity.class) !=null) {
-            return NodeGraphRepository.class;
-        }
-        if (findAnnotation(domainClass, RelationshipEntity.class) !=null) {
-            return RelationshipGraphRepository.class;
-        }
+        final GraphEntityInformation entityInformation = (GraphEntityInformation) getEntityInformation(domainClass);
+        if (entityInformation.isNodeEntity()) return NodeGraphRepository.class;
+        if (entityInformation.isRelationshipEntity()) return RelationshipGraphRepository.class;
         throw new IllegalArgumentException("Invalid Domain Class "+ domainClass+" neither Node- nor RelationshipEntity");
     }