feat(drizzle): support schema parameter (#183)

Author: onmax

examples/drizzle/index.ts

@@ -10,18 +10,20 @@ export const users = sqliteTable("users", {
   name: text("full_name"),
 });
 
+const schema = { users };
+
 async function main() {
-  const db = createDatabase(sqlite({}));
-  const drizzleDb = drizzle(db);
+  const db0 = createDatabase(sqlite({}));
+  const db = drizzle(db0, { schema });
 
-  await db.sql`create table if not exists users (
+  await db0.sql`create table if not exists users (
     id integer primary key autoincrement,
     full_name text
   )`;
 
-  await db.sql`insert into users (full_name) values ('John Doe')`;
+  await db0.sql`insert into users (full_name) values ('John Doe')`;
 
-  const res = await drizzleDb.select().from(users).all();
+  const res = await db.select().from(users).all();
   console.log({ res });
 }
 

src/integrations/drizzle/index.ts

@@ -1,23 +1,59 @@
+import type { Database } from "db0";
+import { DB0Session, type DB0SessionOptions } from "./_session.ts";
+
+import { DefaultLogger } from "drizzle-orm/logger";
+
 import {
   BaseSQLiteDatabase,
   SQLiteAsyncDialect,
 } from "drizzle-orm/sqlite-core";
-import type { Database } from "db0";
-import { DB0Session } from "./_session.ts";
+
+import {
+  type DrizzleConfig as DrizzleBaseConfig,
+  type RelationalSchemaConfig,
+  type TablesRelationalConfig,
+  createTableRelationsHelpers,
+  extractTablesRelationalConfig,
+} from "drizzle-orm";
 
 export type DrizzleDatabase<
   TSchema extends Record<string, unknown> = Record<string, never>,
 > = BaseSQLiteDatabase<"async", any, TSchema>;
 
+export type DrizzleConfig<
+  TSchema extends Record<string, unknown> = Record<string, never>,
+> = DrizzleBaseConfig<TSchema>;
+
 export function drizzle<
   TSchema extends Record<string, unknown> = Record<string, never>,
->(db: Database): DrizzleDatabase<TSchema> {
-  // TODO: Support schema
-  const schema = undefined;
+>(db: Database, config?: DrizzleConfig<TSchema>): DrizzleDatabase<TSchema> {
+  const dialect = new SQLiteAsyncDialect({ casing: config?.casing });
+
+  let logger: DB0SessionOptions["logger"];
+  if (config?.logger === true) {
+    logger = new DefaultLogger();
+  } else if (config?.logger !== false && config?.logger !== undefined) {
+    logger = config.logger;
+  }
 
-  const dialect = new SQLiteAsyncDialect();
+  // Transform user schema to RelationalSchemaConfig
+  // Reference: https://github.com/drizzle-team/drizzle-orm/blob/main/drizzle-orm/src/d1/driver.ts
+  let schema: RelationalSchemaConfig<TablesRelationalConfig> | undefined;
+  if (config?.schema) {
+    const tablesConfig = extractTablesRelationalConfig(
+      config.schema,
+      createTableRelationsHelpers,
+    );
+    schema = {
+      fullSchema: config.schema,
+      schema: tablesConfig.tables,
+      tableNamesMap: tablesConfig.tableNamesMap,
+    };
+  }
 
-  const session = new DB0Session(db, dialect, schema);
+  const session = new DB0Session(db, dialect, schema, {
+    logger,
+  });
 
   return new BaseSQLiteDatabase(
     "async",

test/integrations/drizzle.test.ts

@@ -47,11 +47,65 @@ describe("integrations: drizzle: better-sqlite3", () => {
     expect(res[0].name).toBe("John Doe");
   });
 
+  it("accepts boolean logger config", () => {
+    expect(() => drizzle(db, { logger: true })).not.toThrow();
+  });
+
   afterAll(async () => {
     await db.sql`DROP TABLE IF EXISTS users`;
   });
 });
 
+describe("integrations: drizzle: with schema parameter", () => {
+  const users = dSqlite.sqliteTable("users_schema", {
+    id: dSqlite.numeric("id"),
+    name: dSqlite.text("name"),
+    email: dSqlite.text("email"),
+  });
+
+  const schema = { users };
+
+  let drizzleDb: DrizzleDatabase<typeof schema>;
+  let db: Database;
+
+  beforeAll(async () => {
+    db = createDatabase(sqliteConnector({}));
+    drizzleDb = drizzle(db, { schema });
+    await db.sql`DROP TABLE IF EXISTS users_schema`;
+    await db.sql`create table if not exists users_schema (
+      id integer primary key autoincrement,
+      name text,
+      email text
+    )`;
+  });
+
+  it("insert with schema", async () => {
+    const res = await drizzleDb
+      .insert(users)
+      .values({
+        name: "Jane Doe",
+        email: "[email protected]",
+      })
+      .returning();
+
+    expect(res.length).toBe(1);
+    expect(res[0].name).toBe("Jane Doe");
+    expect(res[0].email).toBe("[email protected]");
+  });
+
+  it("select with schema", async () => {
+    const res = await drizzleDb.select().from(users).all();
+
+    expect(res.length).toBe(1);
+    expect(res[0].name).toBe("Jane Doe");
+    expect(res[0].email).toBe("[email protected]");
+  });
+
+  afterAll(async () => {
+    await db.sql`DROP TABLE IF EXISTS users_schema`;
+  });
+});
+
 describe.runIf(process.env.POSTGRESQL_URL)(
   "integrations: drizzle: postgres",
   () => {