add db seed handler

src/modules/core/config.ts

@@ -14,6 +14,7 @@ export const getDefaultAppConfig = (configure: Configure) => ({
     https: configure.env.get('APP_SSL', (v) => toBoolean(v), false),
     locale: configure.env.get('APP_LOCALE', 'zh_CN'),
     fallbackLocale: configure.env.get('APP_FALLBACK_LOCALE', 'en'),
+    timezone: configure.env.get('APP_TIMEZONE', 'Asia/Shanghai'),
 });
 
 export const createAppConfig: (

src/modules/core/helpers/time.ts

@@ -0,0 +1,30 @@
+import dayjs from 'dayjs';
+
+import advancedFormat from 'dayjs/plugin/advancedFormat';
+import customParseFormat from 'dayjs/plugin/customParseFormat';
+import dayOfYear from 'dayjs/plugin/dayOfYear';
+import localeData from 'dayjs/plugin/localeData';
+import timezone from 'dayjs/plugin/timezone';
+import utc from 'dayjs/plugin/utc';
+
+import { Configure } from '@/modules/config/configure';
+import { AppConfig, TimeOptions } from '@/modules/core/types';
+
+dayjs.extend(localeData);
+dayjs.extend(utc);
+dayjs.extend(advancedFormat);
+dayjs.extend(customParseFormat);
+dayjs.extend(dayOfYear);
+dayjs.extend(timezone);
+
+/**
+ * 获取一个dayjs时间对象
+ * @param configure
+ * @param options
+ */
+export async function getTime(configure: Configure, options?: TimeOptions) {
+    const { date, format, locale, strict, zonetime } = options ?? {};
+    const config = await configure.get<AppConfig>('app');
+    const now = dayjs(date, format, locale ?? config.locale, strict).clone();
+    return now.tz(zonetime ?? config.timezone);
+}

src/modules/core/types.ts

@@ -1,6 +1,7 @@
 import { ModuleMetadata, PipeTransform, Type } from '@nestjs/common';
 import { NestFastifyApplication } from '@nestjs/platform-fastify';
 
+import dayjs from 'dayjs';
 import { Ora } from 'ora';
 import { StartOptions } from 'pm2';
 import { CommandModule } from 'yargs';
@@ -53,10 +54,21 @@ export interface AppConfig {
 
     https: boolean;
 
+    /**
+     * 语言,默认zh-cn
+     */
     locale: string;
 
+    /**
+     * 备用语言
+     */
     fallbackLocale: string;
 
+    /**
+     * 时区,默认Asia/Shanghai
+     */
+    timezone: string;
+
     url?: string;
 
     prefix?: string;
@@ -64,6 +76,32 @@ export interface AppConfig {
     pm2?: Omit<StartOptions, 'name' | 'cwd' | 'script' | 'args' | 'interpreter' | 'watch'>;
 }
 
+/**
+ * 时间参数选项类型
+ */
+export interface TimeOptions {
+    /**
+     * 时间属性，如果不传入则获取当前属性
+     */
+    date?: dayjs.ConfigType;
+    /**
+     * 输出时间格式
+     */
+    format?: dayjs.OptionType;
+    /**
+     * 语言，如果不传入则使用app配置中设置的默认语言
+     */
+    locale?: string;
+    /**
+     * 是否开启严格模式
+     */
+    strict?: boolean;
+    /**
+     * 时区。如果不传入则使用app配置中设置的默认时区
+     */
+    zonetime?: string;
+}
+
 export interface PanicOption {
     message: string;
 

src/modules/database/base/BaseSeeder.ts

@@ -0,0 +1,85 @@
+import { isNil } from 'lodash';
+import { Ora } from 'ora';
+
+import { DataSource, EntityManager, EntityTarget, ObjectLiteral } from 'typeorm';
+
+import { Configure } from '@/modules/config/configure';
+import { panic } from '@/modules/core/helpers';
+import {
+    Seeder,
+    SeederConstructor,
+    SeederLoadParams,
+    SeederOptions,
+} from '@/modules/database/commands/types';
+import { DBOptions } from '@/modules/database/types';
+
+/**
+ * 数据填充基类
+ */
+export abstract class BaseSeeder implements Seeder {
+    protected connection: string;
+    protected dataSource: DataSource;
+    protected em: EntityManager;
+    protected configure: Configure;
+    protected ignoreLock: boolean;
+    protected truncates: EntityTarget<ObjectLiteral>[] = [];
+
+    constructor(
+        protected readonly spinner: Ora,
+        protected readonly args: SeederOptions,
+    ) {}
+
+    /**
+     * 清空原数据并重新加载数据
+     * @param params
+     */
+    async load(params: SeederLoadParams): Promise<any> {
+        const { connection, dataSource, em, configure, ignoreLock } = params;
+        this.connection = connection;
+        this.dataSource = dataSource;
+        this.em = em;
+        this.configure = configure;
+        this.ignoreLock = ignoreLock;
+
+        if (this.ignoreLock) {
+            for (const option of this.truncates) {
+                await this.em.clear(option);
+            }
+        }
+
+        return this.run(this.dataSource);
+    }
+
+    /**
+     * 运行seeder的关键方法
+     * @param dataSource
+     * @param em
+     * @protected
+     */
+    protected abstract run(dataSource: DataSource, em?: EntityManager): Promise<any>;
+
+    protected async getDBConfig() {
+        const { connections = [] }: DBOptions = await this.configure.get<DBOptions>('database');
+        const dbConfig = connections.find(({ name }) => name === this.connection);
+        if (isNil(dbConfig)) {
+            await panic(`Database connection named ${this.connection} not exists!`);
+        }
+        return dbConfig;
+    }
+
+    /**
+     * 运行子seeder
+     * @param SubSeeder
+     * @protected
+     */
+    protected async call(SubSeeder: SeederConstructor) {
+        const subSeeder: Seeder = new SubSeeder(this.spinner, this.args);
+        await subSeeder.load({
+            connection: this.connection,
+            dataSource: this.dataSource,
+            em: this.em,
+            configure: this.configure,
+            ignoreLock: this.ignoreLock,
+        });
+    }
+}

src/modules/database/commands/types.ts

@@ -1,5 +1,9 @@
+import { Ora } from 'ora';
+import { DataSource, EntityManager } from 'typeorm';
 import { Arguments } from 'yargs';
 
+import { Configure } from '@/modules/config/configure';
+
 /**
  * 基础数据库命令参数类型
  */
@@ -60,3 +64,61 @@ export interface MigrationRevertOptions {
  * 恢复迁移的命令参数
  */
 export type MigrationRevertArguments = TypeOrmArguments & MigrationRevertOptions;
+
+/**
+ * 数据填充处理器选项
+ */
+export interface SeederOptions {
+    /**
+     * 数据库连接名称
+     */
+    connection?: string;
+    /**
+     * 是否通过事务来运行填充
+     */
+    transaction?: boolean;
+    /**
+     * 是否忽略已经被执行过的填充类
+     */
+    ignorelock?: boolean;
+}
+
+/**
+ * 数据填充类接口
+ */
+export interface SeederConstructor {
+    new (spinner: Ora, args: SeederOptions): Seeder;
+}
+
+/**
+ * 数据填充类方法对象
+ */
+export interface Seeder {
+    load: (params: SeederLoadParams) => Promise<void>;
+}
+
+/**
+ * 数据填充类的load函数参数
+ */
+export interface SeederLoadParams {
+    /**
+     * 数据库连接名称
+     */
+    connection: string;
+    /**
+     * 数据库连接
+     */
+    dataSource: DataSource;
+    /**
+     * EntityManager实例
+     */
+    em: EntityManager;
+    /**
+     * 项目配置类
+     */
+    configure: Configure;
+    /**
+     * 是否忽略锁定
+     */
+    ignoreLock: boolean;
+}

src/modules/database/config.ts

@@ -1,5 +1,7 @@
 import { resolve } from 'path';
 
+import { SeederRunner } from '@/modules/database/resolver/seeder.runner';
+
 import { ConfigureFactory, ConfigureRegister } from '../config/types';
 import { createConnectionOptions } from '../config/utils';
 import { deepMerge } from '../core/helpers';
@@ -12,7 +14,7 @@ export const createDBConfig: (
     register,
     hook: (configure, value) => createDBOptions(value),
     defaultRegister: () => ({
-        common: { charset: 'utf8mb4', logging: ['error'] },
+        common: { charset: 'utf8mb4', logging: ['error'], seeders: [], seedRunner: SeederRunner },
         connections: [],
     }),
 });

src/modules/database/resolver/seeder.runner.ts

@@ -0,0 +1,37 @@
+import { resolve } from 'path';
+
+import { Type } from '@nestjs/common';
+import { ensureFileSync, readFileSync, writeFileSync } from 'fs-extra';
+import { get, isNil, set } from 'lodash';
+import { DataSource, EntityManager } from 'typeorm';
+import YAML from 'yaml';
+
+import { BaseSeeder } from '@/modules/database/base/BaseSeeder';
+
+/**
+ * 默认的Seed Runner
+ */
+export class SeederRunner extends BaseSeeder {
+    protected async run(dataSource: DataSource, em?: EntityManager): Promise<any> {
+        let seeders: Type<any>[] = ((await this.getDBConfig()) as any).seeders ?? [];
+        const seedLockFile = resolve(__dirname, '../../../..', 'seed-lock.yml');
+        ensureFileSync(seedLockFile);
+        const lockFileYml = YAML.parse(readFileSync(seedLockFile, 'utf8'));
+        const locked = isNil(lockFileYml) ? {} : lockFileYml;
+        const lockNames = get<string[]>(locked, this.connection, []);
+        if (!this.ignoreLock) {
+            seeders = seeders.filter((s) => !lockNames.includes(s.name));
+        }
+        for (const seeder of seeders) {
+            await this.call(seeder);
+        }
+        set(
+            locked,
+            this.connection,
+            this.ignoreLock
+                ? seeders.map((s) => s.name)
+                : [...lockNames, ...seeders.map((s) => s.name)],
+        );
+        writeFileSync(seedLockFile, JSON.stringify(locked, null, 4));
+    }
+}

src/modules/database/types.ts

@@ -7,6 +7,7 @@ import {
     TreeRepository,
 } from 'typeorm';
 
+import { SeederConstructor } from '@/modules/database/commands/types';
 import { OrderType, SelectTrashMode } from '@/modules/database/constants';
 
 import { BaseRepository } from './base/repository';
@@ -92,4 +93,13 @@ type DBAdditionalOption = {
      * 是否在启动应用后自动运行迁移
      */
     autoMigrate?: boolean;
+
+    /**
+     * 数据填充类列表
+     */
+    seeders?: SeederConstructor[];
+    /**
+     * 数据填充入口类
+     */
+    seedRunner?: SeederConstructor;
 };

src/modules/database/utils.ts

@@ -2,14 +2,25 @@ import { Type } from '@nestjs/common';
 import { TypeOrmModule } from '@nestjs/typeorm';
 import { EntityClassOrSchema } from '@nestjs/typeorm/dist/interfaces/entity-class-or-schema.type';
 import { isArray, isNil } from 'lodash';
-import { DataSource, ObjectLiteral, ObjectType, Repository, SelectQueryBuilder } from 'typeorm';
+import { Ora } from 'ora';
+import {
+    DataSource,
+    DataSourceOptions,
+    EntityManager,
+    ObjectLiteral,
+    ObjectType,
+    Repository,
+    SelectQueryBuilder,
+} from 'typeorm';
 
 import { Configure } from '@/modules/config/configure';
+import { Seeder, SeederConstructor, SeederOptions } from '@/modules/database/commands/types';
 import {
     DBOptions,
     OrderQueryType,
     PaginateOptions,
     PaginateReturn,
+    TypeormOption,
 } from '@/modules/database/types';
 
 import { CUSTOM_REPOSITORY_METADATA } from './constants';
@@ -151,3 +162,86 @@ export async function addSubscribers(
     configure.set('database.connections', newSubscribers);
     return subscribers;
 }
+
+/**
+ * 忽略外键
+ * @param em EntityManager实例
+ * @param type 数据库类型
+ * @param disabled 是否禁用外键
+ */
+export async function resetForeignKey(
+    em: EntityManager,
+    type = 'mysql',
+    disabled = true,
+): Promise<EntityManager> {
+    let key: string;
+    let query: string;
+    if (type === 'sqlite') {
+        key = disabled ? 'OFF' : 'ON';
+        query = `PRAGMA foreign_keys = ${key}`;
+    } else {
+        key = disabled ? '0' : '1';
+        query = `SET FOREIGN_KEY_CHECKS = ${key}`;
+    }
+    await em.query(query);
+    return em;
+}
+
+/**
+ * 数据填充函数
+ * @param Clazz 填充类
+ * @param args 填充命令参数
+ * @param spinner Ora雪碧图标
+ * @param configure 配置对象
+ * @param dbConfig 当前数据库连接池的配置
+ */
+export async function runSeeder(
+    Clazz: SeederConstructor,
+    args: SeederOptions,
+    spinner: Ora,
+    configure: Configure,
+    dbConfig: TypeormOption,
+): Promise<DataSource> {
+    const seeder: Seeder = new Clazz(spinner, args);
+    const dataSource: DataSource = new DataSource({ ...dbConfig } as DataSourceOptions);
+
+    await dataSource.initialize();
+    if (typeof args.transaction === 'boolean' && !args.transaction) {
+        const em = await resetForeignKey(dataSource.manager, dataSource.options.type);
+        await seeder.load({
+            dataSource,
+            em,
+            configure,
+            connection: args.connection ?? 'default',
+            ignoreLock: args.ignorelock,
+        });
+        await resetForeignKey(em, dataSource.options.type, false);
+    } else {
+        // 在事务中运行
+        const queryRunner = dataSource.createQueryRunner();
+        await queryRunner.connect();
+        await queryRunner.startTransaction();
+        try {
+            const em = await resetForeignKey(dataSource.manager, dataSource.options.type);
+            await seeder.load({
+                dataSource,
+                em,
+                configure,
+                connection: args.connection ?? 'default',
+                ignoreLock: args.ignorelock,
+            });
+            await resetForeignKey(em, dataSource.options.type, false);
+            await queryRunner.commitTransaction();
+        } catch (e) {
+            console.error(e);
+            await queryRunner.rollbackTransaction();
+        } finally {
+            await queryRunner.release();
+        }
+    }
+
+    if (dataSource && dataSource.isInitialized) {
+        await dataSource.destroy();
+    }
+    return dataSource;
+}