Drizzle Seed

PostgreSQL
SQLite
MySQL
SingleStore
IMPORTANT

drizzle-seed 只能与 drizzle-orm@0.36.4 或更高版本一起使用。低于此版本的版本在运行时可能能工作,但可能存在类型问题和标识列问题,因为此补丁是从 drizzle-orm@0.36.4 开始引入的

drizzle-seed 是一个 TypeScript 库,帮助你生成确定性的、但又逼真的假数据来填充你的数据库。通过利用一个可设置种子的伪随机数生成器(pRNG),它保证你生成的数据在不同运行间保持一致且可复现。这对于测试、开发和调试尤为有用。

什么是确定性数据生成?

确定性数据生成意味着相同的输入总是产生相同的输出。在 drizzle-seed 的上下文中,当你用相同的种子数字初始化库时,它每次都会生成同样的假数据序列。这允许数据集是可预测和可重复的。

伪随机数生成器(pRNG)

伪随机数生成器是一种算法,用以生成近似随机数性质的数字序列。但由于它基于一个初始值称为种子,你可以控制它的随机性。使用相同的种子,pRNG 会生成相同的数字序列,使你的数据生成过程可复现。

使用 pRNG 的好处:

使用 drizzle-seed,你可以兼得生成逼真假数据的能力和随时复现这些数据的控制权。

安装

npm
yarn
pnpm
bun
npm i drizzle-seed

基本用法

这个示例中,我们将创建 10 个具有随机名称和 ID 的用户

import { pgTable, integer, text } from "drizzle-orm/pg-core";
import { drizzle } from "drizzle-orm/node-postgres";
import { seed } from "drizzle-seed";

const users = pgTable("users", {
  id: integer().primaryKey(),
  name: text().notNull(),
});

async function main() {
  const db = drizzle(process.env.DATABASE_URL!);
  await seed(db, { users });
}

main();

配置选项

count

默认情况下,seed 函数会创建 10 个实体。如果测试需要更多数据,可以在 seed 的选项对象中指定

await seed(db, schema, { count: 1000 });

seed

如果你想针对后续运行生成不同的数据集,可以在 seed 选项中定义一个不同的数字。不同的数字将生成唯一的数据集

await seed(db, schema, { seed: 12345 });

重置数据库

使用 drizzle-seed 可以轻松重置数据库并用新值填充,例如在测试套件中

// 指向你想重置的 schema 文件的路径
import * as schema from "./schema.ts";
import { reset } from "drizzle-seed";

async function main() {
  const db = drizzle(process.env.DATABASE_URL!);
  await reset(db, schema);
}

main();

不同数据库方言有不同的重置策略

PostgreSQL
MySQL
SQLite

针对 PostgreSQL,drizzle-seed 包会生成带有 CASCADE 选项的 TRUNCATE 语句, 确保运行 reset 函数后所有表都为空

TRUNCATE tableName1, tableName2, ... CASCADE;

精细化定制

如果你需要更改 drizzle-seed 默认使用的种子生成函数的行为,可以指定你自己的实现,甚至自己定义值列表用于种子过程。

.refine 是一个回调函数,会接收所有 drizzle-seed 中可用生成器函数的列表。它需要返回一个对象,键是你想要精细化定制的表名,值定义这些表的行为。 每个表可以指定多个属性来简化数据库的种子过程:

info

你还可以为要创建的被引用值数量指定加权随机分布。关于该 API 的细节可以参考 加权随机(Weighted Random)文档 部分。

API

await seed(db, schema).refine((f) => ({
  users: {
    columns: {},
    count: 10,
    with: {
        posts: 10
    }
  },
}));

下面是一些带解释的示例:

schema.ts
import { pgTable, integer, text } from "drizzle-orm/pg-core";

export const users = pgTable("users", {
  id: integer().primaryKey(),
  name: text().notNull(),
});

export const posts = pgTable("posts", {
  id: integer().primaryKey(),
  description: text(),
  userId: integer().references(() => users.id),
});

示例1:只对 users 表进行种子,生成 20 个实体,并对 name 列应用精细化的种子逻辑

index.ts
import { drizzle } from "drizzle-orm/node-postgres";
import { seed } from "drizzle-seed";
import * as schema from './schema.ts'

async function main() {
  const db = drizzle(process.env.DATABASE_URL!);

  await seed(db, { users: schema.users }).refine((f) => ({
    users: {
        columns: {
            name: f.fullName(),
        },
        count: 20
    }
  }));
}

main();

示例2:对 users 表插入 20 个实体,并且为每个用户生成 10 个 posts,通过给 posts 表种子并建立其到 users 表的引用实现关联

index.ts
import { drizzle } from "drizzle-orm/node-postgres";
import { seed } from "drizzle-seed";
import * as schema from './schema.ts'

async function main() {
  const db = drizzle(process.env.DATABASE_URL!);

  await seed(db, schema).refine((f) => ({
    users: {
        count: 20,
        with: {
            posts: 10
        }
    }
  }));
}

main();

示例3:对 users 表种子 5 个实体,且向数据库插入 100 个 posts,但不与 users 实体关联。针对 usersid 精细化为给出 1000020000 之间且唯一的整数,对 postsdescription 精细化为从自定义数组里获取值

index.ts
import { drizzle } from "drizzle-orm/node-postgres";
import { seed } from "drizzle-seed";
import * as schema from './schema.ts'

async function main() {
  const db = drizzle(process.env.DATABASE_URL!);

  await seed(db, schema).refine((f) => ({
    users: {
        count: 5,
        columns: {
            id: f.int({
              minValue: 10000,
              maxValue: 20000,
              isUnique: true,
            }),
        }
    },
    posts: {
        count: 100,
        columns: {
            description: f.valuesFromArray({
            values: [
                "太阳落山,山后染上橙紫色的天空", 
                "我真不敢相信这自制披萨做得这么好!", 
                "有时候,你只需要一本好书和一个安静的角落。", 
                "谁认为雨天是刷老电影的绝佳日子?", 
                "今天试了条新徒步路线,发现了最棒的瀑布!",
                // ...
            ],
          })
        }
    }
  }));
}

main();
IMPORTANT

此文档中将定义更多可能性,目前你可以先浏览这些部分。查看 生成器(Generators) 部分,熟悉所有可用的生成器函数。

一个特别棒的功能是可以使用加权随机,无论是针对列生成的值,还是 drizzle-seed 生成相关实体数量时。

请参考 加权随机(Weighted Random)文档 获取更多信息。

加权随机(Weighted Random)

有时你可能需要使用多个数据集以不同优先级在种子阶段插入数据库。针对这种情况,drizzle-seed 提供了一个称作加权随机的 API。

Drizzle Seed 包在以下场景支持加权随机:

下面给出两个示例:

schema.ts
import { pgTable, integer, text, varchar, doublePrecision } from "drizzle-orm/pg-core";

export const orders = pgTable(
  "orders",
  {
    id: integer().primaryKey(),
    name: text().notNull(),
    quantityPerUnit: varchar().notNull(),
    unitPrice: doublePrecision().notNull(),
    unitsInStock: integer().notNull(),
    unitsOnOrder: integer().notNull(),
    reorderLevel: integer().notNull(),
    discontinued: integer().notNull(),
  }
);

export const details = pgTable(
  "details",
  {
    unitPrice: doublePrecision().notNull(),
    quantity: integer().notNull(),
    discount: doublePrecision().notNull(),

    orderId: integer()
      .notNull()
      .references(() => orders.id, { onDelete: "cascade" }),
  }
);

示例1:精细化 unitPrice 的生成逻辑,生成 5000 条随机价格数据,其中 30% 是 10-100 之间的价格,70% 是 100-300 之间的价格

index.ts
import { drizzle } from "drizzle-orm/node-postgres";
import { seed } from "drizzle-seed";
import * as schema from './schema.ts'

async function main() {
  const db = drizzle(process.env.DATABASE_URL!);

  await seed(db, schema).refine((f) => ({
    orders: {
       count: 5000,
       columns: {
           unitPrice: f.weightedRandom(
               [
                   {
                       weight: 0.3,
                       value: funcs.int({ minValue: 10, maxValue: 100 })
                   },
                   {
                       weight: 0.7,
                       value: funcs.number({ minValue: 100, maxValue: 300, precision: 100 })
                   }
               ]
           ),
       }
    }
  }));
}

main();

示例2:为每个订单生成 1 到 3 条详情的概率为 60%,生成 5 到 7 条详情概率为 30%,生成 8 到 10 条详情概率为 10%

index.ts
import { drizzle } from "drizzle-orm/node-postgres";
import { seed } from "drizzle-seed";
import * as schema from './schema.ts'

async function main() {
  const db = drizzle(process.env.DATABASE_URL!);

  await seed(db, schema).refine((f) => ({
    orders: {
       with: {
           details:
               [
                   { weight: 0.6, count: [1, 2, 3] },
                   { weight: 0.3, count: [5, 6, 7] },
                   { weight: 0.1, count: [8, 9, 10] },
               ]
       }
    }
  }));
}

main();

复杂示例

main.ts
schema.ts
import { seed } from "drizzle-seed";
import * as schema from "./schema.ts";

const main = async () => {
    const titlesOfCourtesy = ["Ms.", "Mrs.", "Dr."];
    const unitsOnOrders = [0, 10, 20, 30, 50, 60, 70, 80, 100];
    const reorderLevels = [0, 5, 10, 15, 20, 25, 30];
    const quantityPerUnit = [
        "100 - 100 g pieces",
        "100 - 250 g bags",
        "10 - 200 g glasses",
        "10 - 4 oz boxes",
        "10 - 500 g pkgs.",
        "10 - 500 g pkgs."
    ];
    const discounts = [0.05, 0.15, 0.2, 0.25];

    await seed(db, schema).refine((funcs) => ({
        customers: {
            count: 10000,
            columns: {
                companyName: funcs.companyName(),
                contactName: funcs.fullName(),
                contactTitle: funcs.jobTitle(),
                address: funcs.streetAddress(),
                city: funcs.city(),
                postalCode: funcs.postcode(),
                region: funcs.state(),
                country: funcs.country(),
                phone: funcs.phoneNumber({ template: "(###) ###-####" }),
                fax: funcs.phoneNumber({ template: "(###) ###-####" })
            }
        },
        employees: {
            count: 200,
            columns: {
                firstName: funcs.firstName(),
                lastName: funcs.lastName(),
                title: funcs.jobTitle(),
                titleOfCourtesy: funcs.valuesFromArray({ values: titlesOfCourtesy }),
                birthDate: funcs.date({ minDate: "2010-12-31", maxDate: "2010-12-31" }),
                hireDate: funcs.date({ minDate: "2010-12-31", maxDate: "2024-08-26" }),
                address: funcs.streetAddress(),
                city: funcs.city(),
                postalCode: funcs.postcode(),
                country: funcs.country(),
                homePhone: funcs.phoneNumber({ template: "(###) ###-####" }),
                extension: funcs.int({ minValue: 428, maxValue: 5467 }),
                notes: funcs.loremIpsum()
            }
        },
        orders: {
            count: 50000,
            columns: {
                shipVia: funcs.int({ minValue: 1, maxValue: 3 }),
                freight: funcs.number({ minValue: 0, maxValue: 1000, precision: 100 }),
                shipName: funcs.streetAddress(),
                shipCity: funcs.city(),
                shipRegion: funcs.state(),
                shipPostalCode: funcs.postcode(),
                shipCountry: funcs.country()
            },
            with: {
                details:
                    [
                        { weight: 0.6, count: [1, 2, 3, 4] },
                        { weight: 0.2, count: [5, 6, 7, 8, 9, 10] },
                        { weight: 0.15, count: [11, 12, 13, 14, 15, 16, 17] },
                        { weight: 0.05, count: [18, 19, 20, 21, 22, 23, 24, 25] },
                    ]
            }
        },
        suppliers: {
            count: 1000,
            columns: {
                companyName: funcs.companyName(),
                contactName: funcs.fullName(),
                contactTitle: funcs.jobTitle(),
                address: funcs.streetAddress(),
                city: funcs.city(),
                postalCode: funcs.postcode(),
                region: funcs.state(),
                country: funcs.country(),
                phone: funcs.phoneNumber({ template: "(###) ###-####" })
            }
        },
        products: {
            count: 5000,
            columns: {
                name: funcs.companyName(),
                quantityPerUnit: funcs.valuesFromArray({ values: quantityPerUnit }),
                unitPrice: funcs.weightedRandom(
                    [
                        {
                            weight: 0.5,
                            value: funcs.int({ minValue: 3, maxValue: 300 })
                        },
                        {
                            weight: 0.5,
                            value: funcs.number({ minValue: 3, maxValue: 300, precision: 100 })
                        }
                    ]
                ),
                unitsInStock: funcs.int({ minValue: 0, maxValue: 125 }),
                unitsOnOrder: funcs.valuesFromArray({ values: unitsOnOrders }),
                reorderLevel: funcs.valuesFromArray({ values: reorderLevels }),
                discontinued: funcs.int({ minValue: 0, maxValue: 1 })
            }
        },
        details: {
            columns: {
                unitPrice: funcs.number({ minValue: 10, maxValue: 130 }),
                quantity: funcs.int({ minValue: 1, maxValue: 130 }),
                discount: funcs.weightedRandom(
                    [
                        { weight: 0.5, value: funcs.valuesFromArray({ values: discounts }) },
                        { weight: 0.5, value: funcs.default({ defaultValue: 0 }) }
                    ]
                )
            }
        }
    }));
}

main();

限制

with 的类型限制

由于 TypeScript 的某些限制以及 Drizzle 当前的 API,无法正确推断表之间的引用,尤其是在表之间存在循环依赖的情况下。

这意味着 with 选项会显示 schema 中的所有表,你需要手动选择具有一对多关系的表。

warning

with 选项仅针对一对多关系有效。例如,如果有一个 user 其对应多个 posts,可以用 users with posts,但不能用 posts with users

Drizzle 表第三个参数的类型限制:

目前,我们没有针对 Drizzle 表的第三个参数的类型支持。虽然运行时能工作,但类型层面不能正确识别。