001    package org.trails.hibernate;
002    
003    import com.sun.mirror.apt.AnnotationProcessor;
004    import com.sun.mirror.apt.AnnotationProcessorEnvironment;
005    import com.sun.mirror.apt.AnnotationProcessorFactory;
006    import com.sun.mirror.declaration.AnnotationTypeDeclaration;
007    import com.sun.mirror.declaration.ClassDeclaration;
008    import com.sun.mirror.declaration.Declaration;
009    import com.sun.mirror.util.SimpleDeclarationVisitor;
010    import org.apache.commons.logging.Log;
011    import org.apache.commons.logging.LogFactory;
012    import org.dom4j.Document;
013    import org.dom4j.Element;
014    import org.dom4j.io.OutputFormat;
015    import org.dom4j.io.SAXReader;
016    import org.dom4j.io.XMLWriter;
017    import org.hibernate.util.DTDEntityResolver;
018    
019    import java.io.File;
020    import java.io.PrintWriter;
021    import java.util.*;
022    
023    public class HibernateAnnotationProcessorFactory implements AnnotationProcessorFactory
024    {
025            private static final Log LOG = LogFactory.getLog(HibernateAnnotationProcessorFactory.class);
026    
027            public static final String TRAILS_PACKAGE = "org.trails";
028    
029            public static final String configFileOption = "configFile";
030    
031            public static final String destFileOption = "destFile";
032    
033            public HibernateAnnotationProcessorFactory()
034            {
035                    super();
036                    // TODO Auto-generated constructor stub
037            }
038    
039            public Collection<String> supportedOptions()
040            {
041                    return null;
042            }
043    
044            public Collection<String> supportedAnnotationTypes()
045            {
046                    return Arrays.asList("javax.persistence.Entity");
047            }
048    
049            public AnnotationProcessor getProcessorFor(
050                    Set<AnnotationTypeDeclaration> decls,
051                    AnnotationProcessorEnvironment env)
052            {
053                    return new HibernateAnnotationProcessor(env, decls);
054            }
055    
056            public class HibernateAnnotationProcessor implements AnnotationProcessor
057            {
058                    private AnnotationProcessorEnvironment env;
059    
060                    private Set<AnnotationTypeDeclaration> annotationTypeDeclarations;
061    
062                    public HibernateAnnotationProcessor(AnnotationProcessorEnvironment env,
063                                                                                            Set<AnnotationTypeDeclaration> annTypeDecls)
064                    {
065                            this.env = env;
066                            this.annotationTypeDeclarations = annTypeDecls;
067                    }
068    
069                    Element sessionFactoryElement;
070    
071                    public void process()
072                    {
073                            String configTemplateFilePath = getOptionValue(configFileOption);
074    
075                            LOG.info(configTemplateFilePath);
076                            try
077                            {
078                                    SAXReader reader = new SAXReader();
079                                    reader.setValidation(false);
080                                    reader.setEntityResolver(new DTDEntityResolver());
081                                    reader.setIncludeExternalDTDDeclarations(false);
082                                    reader.setIncludeInternalDTDDeclarations(false);
083                                    // Create file first, this is more reliable if there are spaces
084                                    // in the path
085                                    File configTemplateFile = new File(configTemplateFilePath);
086                                    Document doc = reader.read(configTemplateFile);
087    
088                                    sessionFactoryElement = (Element) doc.getRootElement()
089                                            .selectSingleNode("//session-factory");
090    
091                                    String trailsXPath = "mapping[not(starts-with(@class, '"
092                                            + TRAILS_PACKAGE + "'))]";
093                                    for (Iterator iter = sessionFactoryElement.selectNodes(
094                                            trailsXPath).iterator(); iter.hasNext();)
095                                    {
096                                            Element element = (Element) iter.next();
097    
098                                            sessionFactoryElement.remove(element);
099                                    }
100    
101                                    List listenerElements = sessionFactoryElement
102                                            .elements("listener");
103                                    sessionFactoryElement.elements().removeAll(listenerElements);
104    
105                                    for (AnnotationTypeDeclaration annotationTypeDecl : annotationTypeDeclarations)
106                                    {
107                                            LOG.info(annotationTypeDecl);
108                                            for (Declaration declaration : env
109                                                    .getDeclarationsAnnotatedWith(annotationTypeDecl))
110                                            {
111                                                    declaration.accept(new SimpleDeclarationVisitor()
112                                                    {
113                                                            public void visitClassDeclaration(
114                                                                    ClassDeclaration classDeclaration)
115                                                            {
116                                                                    LOG.info(classDeclaration
117                                                                            .getQualifiedName());
118                                                                    sessionFactoryElement.addElement("mapping")
119                                                                            .setAttributeValue(
120                                                                                    "class",
121                                                                                    classDeclaration
122                                                                                            .getQualifiedName());
123                                                            }
124                                                    });
125                                            }
126                                            sessionFactoryElement.elements().addAll(listenerElements);
127                                    }
128                                    String filename = getOptionValue(destFileOption);
129                                    LOG.info("Creating destFile: " + filename);
130                                    File f = new File(filename);
131                                    f.getParentFile().mkdirs();
132    
133                                    OutputFormat format = OutputFormat.createPrettyPrint();
134                                    XMLWriter writer = new XMLWriter(new PrintWriter(f), format);
135                                    writer.write(doc);
136                                    writer.close();
137    
138                            } catch (Exception ex)
139                            {
140                                    ex.printStackTrace();
141                            }
142    
143                    }
144    
145                    private String getOptionValue(String option)
146                    {
147                            // This is a hack due to a bug in apt
148                            for (String key : env.getOptions().keySet())
149                            {
150                                    if (key.startsWith("-A" + option))
151                                    {
152                                            return key.split("=")[1];
153                                    }
154                            }
155                            return null;
156                    }
157    
158            }
159    }