使用 Nile 数据库的 Drizzle

import { drizzle } from 'drizzle-orm/node-postgres';
import dotenv from "dotenv/config";
import { sql } from "drizzle-orm";
import { AsyncLocalStorage } from "async_hooks";

export const db = drizzle(process.env.NILEDB_URL);
export const tenantContext = new AsyncLocalStorage<string | undefined>();

export function tenantDB<T>(cb: (tx: any) => T | Promise<T>): Promise<T> {
  return db.transaction(async (tx) => {
    const tenantId = tenantContext.getStore();
    console.log("执行带租户的查询: " + tenantId);
    // 如果有租户 ID，在事务上下文中设置它
    if (tenantId) {
      await tx.execute(sql`set local nile.tenant_id = '${sql.raw(tenantId)}'`);
    }

    return cb(tx);
  }) as Promise<T>;
}

import 'dotenv/config';
import { defineConfig } from 'drizzle-kit';
export default defineConfig({
  out: './drizzle',
  schema: './src/db/schema.ts',
  dialect: 'postgresql',
  dbCredentials: {
    url: process.env.NILEDB_URL!,
  },
});

// 通过解析生成的表模式
import { pgTable, uuid, text, timestamp, varchar, vector, boolean } from "drizzle-orm/pg-core"
import { sql } from "drizzle-orm"

export const tenants = pgTable("tenants", {
	id: uuid().default(sql`public.uuid_generate_v7()`).primaryKey().notNull(),
	name: text(),
	created: timestamp({ mode: 'string' }).default(sql`LOCALTIMESTAMP`).notNull(),
	updated: timestamp({ mode: 'string' }).default(sql`LOCALTIMESTAMP`).notNull(),
	deleted: timestamp({ mode: 'string' }),
});

// 通过解析生成的表模式
import { pgTable, uuid, text, timestamp, varchar, vector, boolean } from "drizzle-orm/pg-core"
import { sql } from "drizzle-orm"

export const tenants = pgTable("tenants", {
	id: uuid().default(sql`public.uuid_generate_v7()`).primaryKey().notNull(),
	name: text(),
	created: timestamp({ mode: 'string' }).default(sql`LOCALTIMESTAMP`).notNull(),
	updated: timestamp({ mode: 'string' }).default(sql`LOCALTIMESTAMP`).notNull(),
	deleted: timestamp({ mode: 'string' }),
});

export const todos = pgTable("todos", {
	id: uuid().defaultRandom(),
	tenantId: uuid("tenant_id"),
	title: varchar({ length: 256 }),
	estimate: varchar({ length: 256 }),
	embedding: vector({ dimensions: 3 }),
	complete: boolean(),
});

import express from "express";
import { tenantDB, tenantContext, db } from "./db/db";
import {
  tenants as tenantSchema,
  todos as todoSchema,
} from "./db/schema";
import { eq } from "drizzle-orm";

const PORT = process.env.PORT || 3001;

const app = express();
app.listen(PORT, () => console.log(`服务器正在端口 ${PORT} 上运行`));
app.use(express.json());

// 根据 URL 参数在上下文中设置租户 ID
app.use('/api/tenants/:tenantId/*', (req, res, next) => {
  const tenantId = req.params.tenantId;
  console.log("设置上下文为租户: " + tenantId);
  tenantContext.run(tenantId, next);
});

// 创建新租户
app.post("/api/tenants", async (req, res) => {
  try {
    const name = req.body.name;
    var tenants: any = null;
    tenants = await tenantDB(async (tx) => {
        return await tx.insert(tenantSchema).values({ name }).returning();
    });
    res.json(tenants);
  } catch (error: any) {
    console.log("创建租户时出错: " + error.message);
    res.status(500).json({message: "内部服务器错误",});
  }
});

// 返回租户列表
app.get("/api/tenants", async (req, res) => {
  let tenants: any = [];
  try {
      tenants = await tenantDB(async (tx) => {
        return await tx.select().from(tenantSchema);
      });
    res.json(tenants);
  } catch (error: any) {
    console.log("列出租户时出错: " + error.message);
    res.status(500).json({message: "内部服务器错误",});
  }
});

// 为租户添加新任务
app.post("/api/tenants/:tenantId/todos", async (req, res) => {
  try {
    const { title, complete } = req.body;
    if (!title) {
      res.status(400).json({message: "未提供任务标题",});
    }
    const tenantId = req.params.tenantId;

    const newTodo = await tenantDB(async (tx) => {
      return await tx
        .insert(todoSchema)
        .values({ tenantId, title, complete })
        .returning();
    });
    // 返回时不包括 embedding 向量，因为它很大且无用
    res.json(newTodo);
  } catch (error: any) {
    console.log("添加任务时出错: " + error.message);
    res.status(500).json({message: "内部服务器错误",});
  }
});

// 更新租户的任务
// 因为我们在上下文中有租户，所以不需要 where 子句
app.put("/api/tenants/:tenantId/todos", async (req, res) => {
  try {
    const { id, complete } = req.body;
    await tenantDB(async (tx) => {
      return await tx
        .update(todoSchema)
        .set({ complete })
        .where(eq(todoSchema.id, id));
    });
    res.sendStatus(200);
  } catch (error: any) {
    console.log("更新任务时出错: " + error.message);
    res.status(500).json({message: "内部服务器错误",});
  }
});

// 获取租户的所有任务
app.get("/api/tenants/:tenantId/todos", async (req, res) => {
  try {
    // 这里不需要“where”子句，因为我们在上下文中设置了租户 ID
    const todos = await tenantDB(async (tx) => {
      return await tx
        .select({
          id: todoSchema.id,
          tenant_id: todoSchema.tenantId,
          title: todoSchema.title,
          estimate: todoSchema.estimate,
        })
        .from(todoSchema);
    });
    res.json(todos);
  } catch (error: any) {
    console.log("列出任务时出错: " + error.message);
    res.status(500).json({message: error.message,});
  }
});