View Javadoc
1   /**
2    *    Copyright 2006-2016 the original author or authors.
3    *
4    *    Licensed under the Apache License, Version 2.0 (the "License");
5    *    you may not use this file except in compliance with the License.
6    *    You may obtain a copy of the License at
7    *
8    *       http://www.apache.org/licenses/LICENSE-2.0
9    *
10   *    Unless required by applicable law or agreed to in writing, software
11   *    distributed under the License is distributed on an "AS IS" BASIS,
12   *    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13   *    See the License for the specific language governing permissions and
14   *    limitations under the License.
15   */
16  package org.mybatis.generator.codegen.mybatis3.model;
17  
18  import static org.mybatis.generator.internal.util.JavaBeansUtil.getJavaBeansField;
19  import static org.mybatis.generator.internal.util.JavaBeansUtil.getJavaBeansGetter;
20  import static org.mybatis.generator.internal.util.JavaBeansUtil.getJavaBeansSetter;
21  import static org.mybatis.generator.internal.util.messages.Messages.getString;
22  
23  import java.util.ArrayList;
24  import java.util.List;
25  
26  import org.mybatis.generator.api.CommentGenerator;
27  import org.mybatis.generator.api.FullyQualifiedTable;
28  import org.mybatis.generator.api.IntrospectedColumn;
29  import org.mybatis.generator.api.Plugin;
30  import org.mybatis.generator.api.dom.java.CompilationUnit;
31  import org.mybatis.generator.api.dom.java.Field;
32  import org.mybatis.generator.api.dom.java.FullyQualifiedJavaType;
33  import org.mybatis.generator.api.dom.java.JavaVisibility;
34  import org.mybatis.generator.api.dom.java.Method;
35  import org.mybatis.generator.api.dom.java.Parameter;
36  import org.mybatis.generator.api.dom.java.TopLevelClass;
37  import org.mybatis.generator.codegen.AbstractJavaGenerator;
38  import org.mybatis.generator.codegen.RootClassInfo;
39  
40  /**
41   * 
42   * @author Jeff Butler
43   * 
44   */
45  public class PrimaryKeyGenerator extends AbstractJavaGenerator {
46  
47      public PrimaryKeyGenerator() {
48          super();
49      }
50  
51      @Override
52      public List<CompilationUnit> getCompilationUnits() {
53          FullyQualifiedTable table = introspectedTable.getFullyQualifiedTable();
54          progressCallback.startTask(getString(
55                  "Progress.7", table.toString())); //$NON-NLS-1$
56          Plugin plugins = context.getPlugins();
57          CommentGenerator commentGenerator = context.getCommentGenerator();
58  
59          TopLevelClass topLevelClass = new TopLevelClass(introspectedTable
60                  .getPrimaryKeyType());
61          topLevelClass.setVisibility(JavaVisibility.PUBLIC);
62          commentGenerator.addJavaFileComment(topLevelClass);
63  
64          String rootClass = getRootClass();
65          if (rootClass != null) {
66              topLevelClass.setSuperClass(new FullyQualifiedJavaType(rootClass));
67              topLevelClass.addImportedType(topLevelClass.getSuperClass());
68          }
69  
70          if (introspectedTable.isConstructorBased()) {
71              addParameterizedConstructor(topLevelClass);
72              
73              if (!introspectedTable.isImmutable()) {
74                  addDefaultConstructor(topLevelClass);
75              }
76          }
77  
78          commentGenerator.addModelClassComment(topLevelClass, introspectedTable);
79  
80          for (IntrospectedColumn introspectedColumn : introspectedTable
81                  .getPrimaryKeyColumns()) {
82              if (RootClassInfo.getInstance(rootClass, warnings)
83                      .containsProperty(introspectedColumn)) {
84                  continue;
85              }
86  
87              Field field = getJavaBeansField(introspectedColumn, context, introspectedTable);
88              if (plugins.modelFieldGenerated(field, topLevelClass,
89                      introspectedColumn, introspectedTable,
90                      Plugin.ModelClassType.PRIMARY_KEY)) {
91                  topLevelClass.addField(field);
92                  topLevelClass.addImportedType(field.getType());
93              }
94  
95              Method method = getJavaBeansGetter(introspectedColumn, context, introspectedTable);
96              if (plugins.modelGetterMethodGenerated(method, topLevelClass,
97                      introspectedColumn, introspectedTable,
98                      Plugin.ModelClassType.PRIMARY_KEY)) {
99                  topLevelClass.addMethod(method);
100             }
101 
102             if (!introspectedTable.isImmutable()) {
103                 method = getJavaBeansSetter(introspectedColumn, context, introspectedTable);
104                 if (plugins.modelSetterMethodGenerated(method, topLevelClass,
105                         introspectedColumn, introspectedTable,
106                         Plugin.ModelClassType.PRIMARY_KEY)) {
107                     topLevelClass.addMethod(method);
108                 }
109             }
110         }
111 
112         List<CompilationUnit> answer = new ArrayList<CompilationUnit>();
113         if (context.getPlugins().modelPrimaryKeyClassGenerated(
114                 topLevelClass, introspectedTable)) {
115             answer.add(topLevelClass);
116         }
117         return answer;
118     }
119     
120     private void addParameterizedConstructor(TopLevelClass topLevelClass) {
121         Method method = new Method();
122         method.setVisibility(JavaVisibility.PUBLIC);
123         method.setConstructor(true);
124         method.setName(topLevelClass.getType().getShortName());
125         context.getCommentGenerator().addGeneralMethodComment(method, introspectedTable);
126         
127         StringBuilder sb = new StringBuilder();
128         for (IntrospectedColumn introspectedColumn : introspectedTable
129                 .getPrimaryKeyColumns()) {
130             method.addParameter(new Parameter(introspectedColumn.getFullyQualifiedJavaType(),
131                     introspectedColumn.getJavaProperty()));
132             sb.setLength(0);
133             sb.append("this."); //$NON-NLS-1$
134             sb.append(introspectedColumn.getJavaProperty());
135             sb.append(" = "); //$NON-NLS-1$
136             sb.append(introspectedColumn.getJavaProperty());
137             sb.append(';');
138             method.addBodyLine(sb.toString());
139         }
140         
141         topLevelClass.addMethod(method);
142     }
143 }