From 06461b76c7d73fd87379c8986e1c9020e5a5c81f Mon Sep 17 00:00:00 2001
From: collinsmith <collinsmith70@gmail.com>
Date: Tue, 15 Dec 2020 13:23:49 -0800
Subject: [PATCH] Added PrimaryKey AnnotationMirror for PrimaryKeyElement

Added PrimaryKey AnnotationMirror for PrimaryKeyElement
Created Context#getAnnotationMirror
Created Constants#PRIMARY_KEY
---
 .../com/riiablo/table/annotation/Constants.java     |  1 +
 .../java/com/riiablo/table/annotation/Context.java  | 11 +++++++++++
 .../riiablo/table/annotation/PrimaryKeyElement.java | 13 ++++++++++---
 3 files changed, 22 insertions(+), 3 deletions(-)

diff --git a/table/annotation-processor/src/main/java/com/riiablo/table/annotation/Constants.java b/table/annotation-processor/src/main/java/com/riiablo/table/annotation/Constants.java
index f693ddab..c93817aa 100644
--- a/table/annotation-processor/src/main/java/com/riiablo/table/annotation/Constants.java
+++ b/table/annotation-processor/src/main/java/com/riiablo/table/annotation/Constants.java
@@ -7,6 +7,7 @@ final class Constants {
   private Constants() {}
 
   static final ClassName STRING = ClassName.get(String.class);
+  static final ClassName PRIMARY_KEY = ClassName.get(PrimaryKey.class);
 
   static final TypeName[] PRIMARY_KEY_TYPES = { TypeName.INT, STRING };
 }
diff --git a/table/annotation-processor/src/main/java/com/riiablo/table/annotation/Context.java b/table/annotation-processor/src/main/java/com/riiablo/table/annotation/Context.java
index 30b5f51f..b552294d 100644
--- a/table/annotation-processor/src/main/java/com/riiablo/table/annotation/Context.java
+++ b/table/annotation-processor/src/main/java/com/riiablo/table/annotation/Context.java
@@ -1,5 +1,6 @@
 package com.riiablo.table.annotation;
 
+import com.squareup.javapoet.ClassName;
 import java.lang.annotation.Annotation;
 import java.util.Objects;
 import javax.annotation.processing.Messager;
@@ -26,6 +27,16 @@ class Context {
     elementUtils = processingEnvironment.getElementUtils();
   }
 
+  AnnotationMirror getAnnotationMirror(Element element, Class<? extends Annotation> annotationClass) {
+    for (AnnotationMirror annotationMirror : element.getAnnotationMirrors()) {
+      if (Constants.PRIMARY_KEY.equals(ClassName.get(annotationMirror.getAnnotationType()))) {
+        return annotationMirror;
+      }
+    }
+
+    return null;
+  }
+
   void log(
       Diagnostic.Kind kind,
       Element element,
diff --git a/table/annotation-processor/src/main/java/com/riiablo/table/annotation/PrimaryKeyElement.java b/table/annotation-processor/src/main/java/com/riiablo/table/annotation/PrimaryKeyElement.java
index 62ffee0f..24424c91 100644
--- a/table/annotation-processor/src/main/java/com/riiablo/table/annotation/PrimaryKeyElement.java
+++ b/table/annotation-processor/src/main/java/com/riiablo/table/annotation/PrimaryKeyElement.java
@@ -2,6 +2,7 @@ package com.riiablo.table.annotation;
 
 import com.squareup.javapoet.ClassName;
 import java.util.Collection;
+import javax.lang.model.element.AnnotationMirror;
 import javax.lang.model.element.TypeElement;
 import javax.lang.model.element.VariableElement;
 import org.apache.commons.lang3.ArrayUtils;
@@ -25,7 +26,9 @@ class PrimaryKeyElement {
         if (primaryKeyElement == null) {
           primaryKeyElement = e;
         } else {
-          context.error(e, "{} already declared as {}", primaryKeyElement, PrimaryKey.class);
+          context.error(
+              e, context.getAnnotationMirror(e, PrimaryKey.class),
+              "{} already declared as {}", primaryKeyElement, PrimaryKey.class);
         }
       }
     }
@@ -41,19 +44,23 @@ class PrimaryKeyElement {
       primaryKeyElement = firstAcceptableElement;
     }
 
-    return new PrimaryKeyElement(primaryKeyElement);
+    AnnotationMirror primaryKeyMirror = context.getAnnotationMirror(primaryKeyElement, PrimaryKey.class);
+    return new PrimaryKeyElement(primaryKeyElement, primaryKeyMirror);
   }
 
   final VariableElement element;
+  final AnnotationMirror mirror;
 
-  PrimaryKeyElement(VariableElement element) {
+  PrimaryKeyElement(VariableElement element, AnnotationMirror mirror) {
     this.element = element;
+    this.mirror = mirror;
   }
 
   @Override
   public String toString() {
     return new ToStringBuilder(this)
         .append("element", element)
+        .append("mirror", mirror)
         .toString();
   }
 }