Fully custom PushProcessor database example (drizzle)

drizzle-database.ts

import { drizzle as drizzleBuilder } from 'drizzle-orm/postgres-js'
import { sql } from 'drizzle-orm'
import postgres from 'postgres'

import type {
  Database,
  TransactionProviderHooks,
  TransactionProviderInput
} 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 DrizzleDatabase implements Database<DrizzleTransaction> {
  transaction<R>(
    callback: (
      tx: DrizzleTransaction,
      transactionHooks: TransactionProviderHooks,
    ) => Promise<R>,
    transactionInput: TransactionProviderInput,
  ): Promise<R> {
    return drizzle.transaction((drizzleTx) => callback(drizzleTx, {
        // Because we're implementing our own custom push database, we have to handle
        // updating the last mutation ID manually (the ZQLDatabase does this for you)
        async updateClientMutationID() {
          const query =
            sql<{lastMutationID: bigint}[]>`INSERT INTO ${transactionInput.upstreamSchema}.clients 
                    as current ("clientGroupID", "clientID", "lastMutationID")
                        VALUES (${transactionInput.clientGroupID}, ${transactionInput.clientID}, ${1})
                    ON CONFLICT ("clientGroupID", "clientID")
                    DO UPDATE SET "lastMutationID" = current."lastMutationID" + 1
                    RETURNING "lastMutationID"`

          const [{lastMutationID}] = await drizzleTx.execute(query)

          return {lastMutationID};
        },
      })
    )
  }
}

server.ts

import { DrizzleDatabase, type DrizzleTransaction } from './drizzle-database'
import { PushProcessor, ZQLDatabase, type CustomMutatorDefs } from '@rocicorp/zero/pg'

const processor = new PushProcessor(
  new DrizzleDatabase()
)

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

      // tx is the drizzle transaction object, with no further wrapping or abstraction on top.
      // This is NOT a zero ServerTransaction, and no ZQL methods are on the tx.
      // You can't pass this into a client mutator.
      await tx.insert(ExampleTable).values({
        id
      })
    }
  },
// This tells TS that we're passing the drizzle tx itself in as the mutator transaction
// without any additional wrapping (e.g. no ServerTransaction object around it)
} satisfies CustomMutatorDefs<DrizzleTransaction>

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