PostgreSQL分库场景下Drizzle ORM单分片表与迁移分片管理咨询

一、架构层面：隔离分片专属表 #

1. 拆分ORM实例与数据库连接 #

将分片1、分片2的数据库连接和Drizzle ORM实例完全分离，同时保留原有的分片路由逻辑（针对user这类全局表）：

// db/connections.ts
import { drizzle } from 'drizzle-orm/node-postgres';
import { Pool } from 'pg';

// 分片1连接（托管country表）
const poolShard1 = new Pool({
  connectionString: process.env.DATABASE_URL_SHARD1,
});
export const dbShard1 = drizzle(poolShard1);

// 分片2连接
const poolShard2 = new Pool({
  connectionString: process.env.DATABASE_URL_SHARD2,
});
export const dbShard2 = drizzle(poolShard2);

// 分片路由函数（针对user表）
export function getShardForUser(userId: string) {
  // 原有的分片规则，比如取userId哈希值取模
  const hash = userId.split('').reduce((acc, char) => acc + char.charCodeAt(0), 0);
  return hash % 2 === 0 ? dbShard1 : dbShard2;
}

2. 表结构的分片专属注册 #

将country表的Schema仅注册到分片1的ORM实例中，避免被同步到分片2：

// schemas/country.ts
import { pgTable, varchar, serial } from 'drizzle-orm/pg-core';

export const country = pgTable('country', {
  id: serial('id').primaryKey(),
  code: varchar('code', { length: 2 }).notNull().unique(),
  name: varchar('name', { length: 100 }).notNull(),
});

// schemas/index.ts
import { user } from './user';
// 仅在分片1的Schema集合中导入country
export const shard1Schemas = { user, country };
// 分片2的Schema集合仅包含全局表
export const shard2Schemas = { user };

// 初始化时绑定Schema（可选，用于类型提示）
// dbShard1 = drizzle(poolShard1, { schema: shard1Schemas });
// dbShard2 = drizzle(poolShard2, { schema: shard2Schemas });

3. 读密集表的优化（可选） #

针对country这类读密集表，在应用层添加Redis缓存，减少分片1的数据库压力：

// services/countryService.ts
import { dbShard1 } from '../db/connections';
import { country } from '../schemas/country';
import { redisClient } from '../db/redis';
import { eq } from 'drizzle-orm';

export async function getCountryByCode(code: string) {
  const cacheKey = `country:${code}`;
  const cached = await redisClient.get(cacheKey);
  if (cached) return JSON.parse(cached);

  const result = await dbShard1.select().from(country).where(eq(country.code, code)).limit(1);
  if (result[0]) {
    await redisClient.setEx(cacheKey, 86400, JSON.stringify(result[0]));
  }
  return result[0];
}

二、迁移脚本：实现分片特定执行 #

1. 迁移文件的标记规则 #

给迁移文件添加分片标记，推荐两种方式二选一：

• 文件名后缀：全局迁移用.global.ts，分片专属用.shard1.ts/.shard2.ts
  示例：
  • 0001_create_user_table.global.ts（所有分片执行）
  • 0002_create_country_table.shard1.ts（仅分片1执行）
• 文件头部注释标记：在迁移文件顶部添加注释，方便脚本解析

  // @shard: shard1
  import { pgTable, varchar, serial } from 'drizzle-orm/pg-core';
  
  export default {
    async up(db) { /* ... */ },
    async down(db) { /* ... */ },
  };
  

2. 自定义迁移执行脚本 #

替换Drizzle默认的迁移命令，编写Node.js脚本实现分片筛选与执行：

// scripts/migrate.ts
import { readdirSync, readFileSync } from 'fs';
import { join } from 'path';
import { migrate } from 'drizzle-orm/node-postgres/migrator';
import { dbShard1, dbShard2 } from '../db/connections';

// 定义分片配置
const shards = [
  { name: 'shard1', db: dbShard1 },
  { name: 'shard2', db: dbShard2 },
];

// 读取迁移文件并筛选对应分片的文件
function getMigrationsForShard(shardName: string) {
  const migrationsDir = join(__dirname, '../drizzle');
  const files = readdirSync(migrationsDir).filter(f => f.endsWith('.ts'));
  
  return files.map(file => {
    // 方式1：通过文件名后缀筛选（注释掉方式2即可启用）
    const isGlobal = file.endsWith('.global.ts');
    const isTargetShard = file.endsWith(`.${shardName}.ts`);
    if (!isGlobal && !isTargetShard) return null;

    // 方式2：通过注释标记筛选（注释掉方式1即可启用）
    // const content = readFileSync(join(migrationsDir, file), 'utf8');
    // const shardMatch = content.match(/\/\/ @shard: (\w+)/);
    // const isTargetShard = !shardMatch || shardMatch[1] === shardName;
    // if (!isTargetShard) return null;

    return require(join(migrationsDir, file)).default;
  }).filter(Boolean);
}

// 执行分片迁移
async function runMigrations() {
  for (const shard of shards) {
    console.log(`Running migrations for ${shard.name}...`);
    const migrations = getMigrationsForShard(shard.name);
    await migrate(shard.db, {
      migrations: migrations,
      // 禁用自动读取文件夹，避免重复执行
      migrationsFolder: undefined,
    });
    console.log(`Migrations completed for ${shard.name}`);
  }
}

runMigrations().catch(err => {
  console.error('Migration failed:', err);
  process.exit(1);
});

3. 脚本使用与维护 #

• 在package.json中添加脚本命令：

  {
    "scripts": {
      "migrate": "ts-node scripts/migrate.ts"
    }
  }
  

• 新增迁移文件时，按照标记规则命名或添加注释，脚本会自动筛选对应分片执行。

内容的提问来源于stack exchange，提问作者tapan radadiya