Zero Drizzle custom mutator example (postgres.js)

drizzleConnection.ts

import { drizzle as drizzleBuilder } from 'drizzle-orm/postgres-js'
import postgres from 'postgres'
import type { DBConnection, DBTransaction, Row } from '@rocicorp/zero/pg'

// drizzle schema
import * as schema from './schema'

// build your drizzle client instance as per usual
const drizzle = drizzleBuilder(
  postgres(DB_CONNECTION_STRING, /* options... */),
  { schema }
)

// to be able to refer to the type of the drizzle client's
// transaction object, we need to extract it out of the
// drizzle.transaction function's callback parameter
export type DrizzleTransaction = Parameters<
  Parameters<typeof drizzle.transaction>[0]
>[0]

export class Connection implements DBConnection<DrizzleTransaction> {
  query(sql: string, params: unknown[]): Promise<Row[]> {
    // while drizzle has some utilities (e.g. the sql`` template string helper)
    // that can help you work with sql strings, it doesn't have a way to
    // directly pass a parameterized SQL string and a list of parameters
    // through to the postgres client the way the Zero interface wants.
    // luckily, drizzle exposes a 'session' object which contains a reference
    // to the underlying postgres.js client instance. this allows us to
    // bypass drizzle entirely when executing sql on behalf of Zero,
    // while still using drizzle transactions + query builders in our
    // server mutators.
  
    // Unfortunately, the types for the drizzle postgres js driver don't
    // include the 'client' field that holds the postgres js instance,
    // they only include the session object. so we have to cast the session
    // object to call the client.
    return (drizzle._.session as unknown as any).client.unsafe(sql, params)
  }

  transaction<T>(
    fn: (tx: DBTransaction<DrizzleTransaction>) => Promise<T>
  ): Promise<T> {
    return drizzle.transaction((drizzleTx) => fn(new Transaction(drizzleTx)))
  }
}

class Transaction implements DBTransaction<DrizzleTransaction> {
  readonly wrappedTransaction: DrizzleTransaction

  constructor(drizzleTx: DrizzleTransaction) {
    this.wrappedTransaction = drizzleTx
  }

  query(sql: string, params: unknown[]): Promise<Row[]> {
    // the session.client object nested inside a drizzle transaction object
    // holds a reference to the postgres.js client instance for that specific
    // transaction, so this correctly bypasses drizzle while still staying inside
    // the mutation's transaction context.
    return (this.wrappedTransaction._.session as unknown as any).client.unsafe(
      sql,
      params
    )
  }
}

export const drizzleConnection = new Connection()

server.ts

import { drizzleConnection, type DrizzleTransaction } from './drizzleConnection'
import { PushProcessor, ZQLDatabase, type CustomMutatorDefs } from '@rocicorp/zero/pg'

// Import your zero schema from wherever you have it defined
import { schema, type Schema } from './zero-schema'

const processor = new PushProcessor(
  new ZQLDatabase(
    drizzleConnection,
    schema
  )
)

const mutators = {
  test: {
    example: async (
      tx,
      { id }: { id: string }
    ) => {

      // tx.dbTransaction.wrappedTransaction is the drizzle transaction object,
      // with all the usual drizzle query builder apis on it.
      // The top level `tx` object still has all the usual ZQL query and mutate functions.
      await tx.dbTransaction.wrappedTransaction.insert(ExampleTable).values({
        id
      })

      // if sharing mutator code, call your client mutators here with `tx` as the transaction...
    }
  },
// Using ServerTransaction here automatically narrows the type of the transaction
// in your server mutator code so you can safely access dbTransaction.wrappedTransaction,
// but since ServerTransaction still satisfies the ClientTransaction interface it is also
// safe to pass into client mutators!
} satisfies CustomMutatorDefs<ServerTransaction<Schema, DrizzleTransaction>>

/*  wire up your push processor to your API server endpoint per the zero mutator docs...   */