手撸一个造数脚本工具类

手撸一个造数脚本工具类

viEcho Lv5
  2024-06-18 22:55:48 2024-06-18 22:55:48 Created   2024-06-18 23:08:42 2024-06-18 23:08:42 Updated      1.6k Words  9 Mins

上篇博客讲到新建项目后我们需要根据表生成基础的代码，那基础代码写好后，我们要造点数据到表里怎么弄呢？有没有办法，我能生成对应的造数脚本呢？答案是有的，我们借助python强大的faker库即可生成；那总不能每次都要手写新的脚本吧，那必须不能啊；我们java写个生成py脚本的工具类即可，开搞!

定义一个注解

1
2
3
4
5
6
7
8
9
10
11
12
import com.local.demo.enums.FakerTypeEnum;
import java.lang.annotation.*;

@Target(ElementType.FIELD)
@Retention(RetentionPolicy.RUNTIME)
@Documented
public @interface Faker {

     FakerTypeEnum fakerType() default FakerTypeEnum.IGNORE;

     String customVal() default "";
}

实体属性加上注解

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
import com.local.demo.aop.Faker;
import com.local.demo.enums.FakerTypeEnum;
import lombok.Data;
import lombok.experimental.Accessors;

import java.util.Date;

/**
 * 用户信息
 *
 * @author echo
 * @date 2024/06/17
 */
@Data
@Accessors(chain = true)
public class UserInfo {

    /**
     * id
     */
    @Faker(fakerType = FakerTypeEnum.IGNORE)
    private Long  id;

    /**
     * 名称
     */
    @Faker(fakerType = FakerTypeEnum.NAME)
    private String name;

    /**
     * 密码
     */
    @Faker(fakerType = FakerTypeEnum.PWD)
    private String password;

    /**
     * 地址
     */
    @Faker(fakerType = FakerTypeEnum.ADDRESS)
    private String address;

    /**
     * 创建人
     */
    @Faker(customVal = "admin")
    private String createdBy;

    /**
     * 创建时间
     */
    @Faker(fakerType = FakerTypeEnum.NOW_TIME)
    private Date createdTime;

    /**
     * 更新人
     */
    @Faker(customVal = "admin")
    private String updatedBy;

    /**
     * 更新时间
     */
    @Faker(fakerType = FakerTypeEnum.NOW_TIME)
    private Date updatedTime;

}

定义对应的枚举

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
import lombok.Getter;

@Getter
public enum FakerTypeEnum {

    IGNORE(1,"","忽略此字段"),
    NAME(2,"faker.name()","随机姓名"),
    ADDRESS(3,"faker.address()","随机地址"),
    NOW_TIME(4,"datetime.now().strftime('%Y-%m-%d %H:%M:%S')","当前时间"),
    PWD(4,"faker.password()","随机密码"),

    ;

    private final int code;

    private final String fakerVal;

    private final String fakerValDesc;

    FakerTypeEnum(int code,String fakerVal,String fakerValDesc){
        this.code = code;
        this.fakerVal = fakerVal;
        this.fakerValDesc = fakerValDesc;
    }
}

对应的工具类

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
import cn.hutool.core.util.StrUtil;
import com.local.demo.aop.Faker;
import com.local.demo.entity.UserInfo;
import com.local.demo.enums.FakerTypeEnum;
import lombok.extern.slf4j.Slf4j;
import org.springframework.util.CollectionUtils;

import java.io.FileWriter;
import java.io.IOException;
import java.lang.reflect.Field;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Objects;

@Slf4j
public class FakerPyUtil {

    private final static String HOST = "127.0.0.1";
    private final static String DB = "local_demo";
    private final static int PORT = 3306;
    private final static String USER = "root";
    private final static String PWD = "123456";


    public static void main(String[] args) {
        // py地址
        String baseFilePath = "/Users/echo/PycharmProjects/pythonProject/mockSqlPy/";
		// 生成对应实体类集合
        List<Class<?>> classes = Arrays.asList(UserInfo.class);
        generatePyAccordEntity(classes, true, 10, baseFilePath);
    }

    /**
     * 根据实体生成PY脚本
     *
     * @param clazzList           clazz列表
     * @param needTruncateOldData 是否需要清空旧数据
     * @param generateNum         生成num
     * @param baseFilePath        基本文件路径
     */
    public static void generatePyAccordEntity(List<Class<?>> clazzList, Boolean needTruncateOldData, int generateNum, String baseFilePath) {
        if (CollectionUtils.isEmpty(clazzList)) {
            return;
        }
        for (Class<?> aClass : clazzList) {
            try {
                log.info("===>> 开始生成{}实体类的造数脚本", aClass);
                String className = aClass.getName().substring(aClass.getName().lastIndexOf(".") + 1);
                String tableName = convertToUnderlineName(className);
                log.info("实体类对应的表：{}", tableName);

                String filePath = baseFilePath + tableName + ".py";
                log.info("输出位置：{}", filePath);
                FileWriter fileWriter = new FileWriter(filePath);
                fileWriter.write("# This is a Python script\n");
                fileWriter.write("import faker \n");
                fileWriter.write("import random \n");
                fileWriter.write("import pymysql \n");
                fileWriter.write("from datetime import datetime \n\n");

                fileWriter.write("connection = pymysql.connect(host='"+HOST+"', port="+ PORT +", user='"+USER +"', password='"+PWD +"', db='"+DB +"')\n\n");
                fileWriter.write("print (\"数据库连接成功\")\n\n");
                fileWriter.write("faker = faker.Faker(\"zh_CN\")\n\n\n");

                fileWriter.write("def generate_table_name():\n");
                fileWriter.write("    table_name1 = faker.pystr(min_chars=2,max_chars=3)\n");
                fileWriter.write("    table_name2 = faker.pystr(min_chars=2,max_chars=3)\n");
                fileWriter.write("    words =[table_name1, table_name2]\n");
                fileWriter.write("    return '_'. join(words).lower()\n\n\n");
                fileWriter.write("def generate_field_name():\n");
                fileWriter.write("    field_name1 = faker.pystr(min_chars=3,max_chars=5)\n");
                fileWriter.write("    field_name2 = faker.pystr(min_chars=3,max_chars=5)\n");
                fileWriter.write("    words =[field_name1, field_name2]\n");
                fileWriter.write("    return '_'. join(words).lower()\n\n\n");

                fileWriter.write("def create_data():\n");
                fileWriter.write("    try :\n");
                fileWriter.write("        connection_cursor = connection.cursor()\n");
                if (needTruncateOldData) {
                    fileWriter.write("        connection_cursor.execute(\"TRUNCATE TABLE " + tableName + "\")\n");
                    fileWriter.write("        connection.commit()\n");
                }
                fileWriter.write("        for _ in range(" + generateNum + "):\n");
                List<StringBuilder> buildSqlFieldList = buildInsertSqField(fileWriter, aClass);
                fileWriter.write("            " + "execute_sql = f\"\"\"INSERT INTO " + tableName + "(" + buildSqlFieldList.get(0) + ") values (" + buildSqlFieldList.get(1) + ");\"\"\"\n\n");
                fileWriter.write("            " + "print(execute_sql)\n\n");
                fileWriter.write("            " + "connection_cursor.execute(execute_sql)\n");
                fileWriter.write("            " + "connection.commit()\n");
                fileWriter.write("    finally:\n");
                fileWriter.write("        connection.close()\n\n\n");
                fileWriter.write("if __name__ == \"__main__\":\n");
                fileWriter.write("    create_data()\n");
                fileWriter.close();
                log.info("====>> 生成{}实体的造数脚本完毕", aClass);
            } catch (IOException e) {
                throw new RuntimeException(e);
            }

        }
    }

    /**
     * 生成插入sql字段
     *
     * @param fileWriter 文件写入程序
     * @param aClass     类
     * @return {@link List}<{@link StringBuilder}>
     * @throws IOException IOException
     */
    private static List<StringBuilder> buildInsertSqField(FileWriter fileWriter, Class<?> aClass) throws IOException {
        List<StringBuilder> returnList = new ArrayList<>();
        List<Field> allFields = getAllFields(aClass);
        StringBuilder beforeSqlStrBuilder = new StringBuilder();
        StringBuilder lastSqlStrBuilder = new StringBuilder();
        long ignoreCount = allFields.stream().filter(f -> Objects.equals(f.getAnnotation(Faker.class).fakerType(), FakerTypeEnum.IGNORE) &&
                StrUtil.isBlank(f.getAnnotation(Faker.class).customVal())).count();
        int executedIndex = 0;
        for (int i = 0; i < allFields.size(); i++) {
            Field field = allFields.get(i);
            if (field.isAnnotationPresent(Faker.class)) {
                Faker annotation = field.getAnnotation(Faker.class);
                //自增和有默认值的字段不处理
                if (StrUtil.isBlank(annotation.customVal()) && Objects.equals(field.getAnnotation(Faker.class).fakerType(), FakerTypeEnum.IGNORE)) {
                    continue;
                }
                String fieldName = field.getName();
                String tableFieldName = convertToUnderlineName(fieldName);
                FakerTypeEnum fakerTypeEnum = annotation.fakerType();
                if (StrUtil.isNotBlank(annotation.customVal())) {
                    fileWriter.write("            " + tableFieldName + " = '" + annotation.customVal() + "'\n");
                } else {
                    fileWriter.write("            " + tableFieldName + " = " + fakerTypeEnum.getFakerVal() + "\n");
                }
                if (executedIndex == allFields.size() - ignoreCount - 1) {
                    beforeSqlStrBuilder.append(tableFieldName);
                    lastSqlStrBuilder.append("'{").append(tableFieldName).append("}'");
                } else {
                    beforeSqlStrBuilder.append(tableFieldName).append(",");
                    lastSqlStrBuilder.append("'{").append(tableFieldName).append("}',");
                }
                executedIndex++;
            }
        }
        returnList.add(beforeSqlStrBuilder);
        returnList.add(lastSqlStrBuilder);
        return returnList;
    }


    /**
     * 获取所有字段
     *
     * @param type 类型
     * @return {@link List}<{@link Field}>
     */
    public static List<Field> getAllFields(Class<?> type) {
        List<Field> fields = new ArrayList<>();
        for (Class<?> c = type; c != null; c = c.getSuperclass()) {
            Field[] declaredFields = c.getDeclaredFields();
            for (Field declaredField : declaredFields) {
                fields.add(declaredField);
            }
        }
        return fields;
    }


    /**
     * 转换为下划线名称
     *
     * @param className 类名
     * @return {@link String}
     */
    private static String convertToUnderlineName(String className) {
        StringBuilder tableName = new StringBuilder();
        char[] charArray = className.toCharArray();
        for (int i = 0; i < charArray.length; i++) {
            if (i == 0) {
                tableName.append(Character.toLowerCase(charArray[0]));
                continue;
            }
            if (Character.isUpperCase(charArray[i])) {
                tableName.append('_').append(Character.toLowerCase(charArray[i]));
            } else {
                tableName.append(Character.toLowerCase(charArray[i]));
            }
        }
        return tableName.toString();
    }

}

生成完毕截图



至此，一个简单的生成数据的脚本我们就创建好了；可以思考下如果要基于前置数据生成有关联关系的数据脚本要怎么弄，一对多一对一要怎么做？以及有没有办法改造mybatis生成基础代码配置，让其生成对应实体时自动加上对应的注解及对应的枚举值；

  • Title: 手撸一个造数脚本工具类
  • Author: viEcho
  • Created at : 2024-06-18 22:55:48
  • Updated at : 2024-06-18 23:08:42
  • Link: https://viecho.github.io/2024/0618/py-shell-for-generate-db-data.html
  • License: This work is licensed under CC BY-NC-SA 4.0.
推荐阅读
时间复杂度&空间复杂度 时间复杂度&空间复杂度 手撸两个注解 手撸两个注解 统一返回处理&统一异常封装 统一返回处理&统一异常封装
推荐阅读
时间复杂度&空间复杂度 时间复杂度&空间复杂度 手撸两个注解 手撸两个注解
Comments
On this page
手撸一个造数脚本工具类
  1. 定义一个注解
  2. 实体属性加上注解
  3. 定义对应的枚举
  4. 对应的工具类
  5. 生成完毕截图