zero `/push` next.js route with node-postgres

route.ts

import { hmac } from "@oslojs/crypto/hmac";
import { SHA256 } from "@oslojs/crypto/sha2";
import { constantTimeEqual } from "@oslojs/crypto/subtle";
import {
  joseAlgorithmHS256,
  JWSRegisteredHeaders,
  JWTRegisteredClaims,
  parseJWT,
} from "@oslojs/jwt";
import type { JSONValue } from "@rocicorp/zero";
import {
  PushProcessor,
  ZQLDatabase,
  type DBConnection,
  type DBTransaction,
  type Row,
} from "@rocicorp/zero/pg";
import { NextResponse, type NextRequest } from "next/server";
import { Client, type ClientBase } from "pg";
import { z } from "zod";

import { createMutators, schema } from "@trip/zero";

import { env } from "~/env";

const payloadSchema = z
  .object({
    sub: z.string(),
  })
  .passthrough();

async function handler(request: NextRequest) {
  // JSON parsing moved before client connection
  const json = await request
    .json()
    .then((data) => data as JSONValue)
    .catch(() => {
      return {};
    });

  const [header, payload, signature, signatureMessage] = parseJWT(
    request.headers.get("Authorization")?.replace("Bearer ", "") ?? "",
  );

  // Check if the JWT algorithm is valid
  const headerParameters = new JWSRegisteredHeaders(header);
  if (headerParameters.algorithm() !== joseAlgorithmHS256) {
    return new NextResponse("Unsupported algorithm", { status: 401 });
  }

  // Check expiration
  const claims = new JWTRegisteredClaims(payload);
  if (!claims.verifyExpiration()) {
    return new NextResponse("Token expired", { status: 401 });
  }

  // Verify signature
  const secretKey = new TextEncoder().encode(env.SECRET);
  const expectedSignature = hmac(SHA256, secretKey, signatureMessage);

  if (
    expectedSignature.length !== signature.length ||
    !constantTimeEqual(expectedSignature, signature)
  ) {
    return new NextResponse("Invalid signature", { status: 401 });
  }

  const payloadResult = payloadSchema.safeParse(payload);
  if (!payloadResult.success) {
    return new NextResponse("Invalid payload", { status: 401 });
  }
  const { sub } = payloadResult.data;

  const mutators = createMutators({
    sub,
  });

  let client: Client | undefined;

  try {
    client = new Client({
      connectionString: env.DATABASE_URL,
    });
    await client.connect();

    const processor = new PushProcessor(
      new ZQLDatabase(new PgConnection(client), schema),
    );
    const searchParams = request.nextUrl.searchParams;

    const response = await processor.process(mutators, searchParams, json);

    return NextResponse.json(response);
  } catch {
    return NextResponse.json(
      {
        error: "Error processing request",
      },
      {
        status: 500,
      },
    );
  } finally {
    if (client) {
      await client.end();
    }
  }
}

export class PgConnection implements DBConnection<ClientBase> {
  readonly #client: ClientBase;

  constructor(client: ClientBase) {
    this.#client = client;
  }

  async query(sql: string, params: unknown[]): Promise<Row[]> {
    const result = await this.#client.query<Row>(sql, params as JSONValue[]);
    return result.rows;
  }

  async transaction<T>(
    fn: (tx: DBTransaction<ClientBase>) => Promise<T>,
  ): Promise<T> {
    if (!(this.#client instanceof Client)) {
      throw new Error("Transactions require a non-pooled Client instance");
    }

    const tx = new PgTransaction(this.#client);

    try {
      await this.#client.query("BEGIN");
      const result = await fn(tx);
      await this.#client.query("COMMIT");
      return result;
    } catch (error) {
      await this.#client.query("ROLLBACK");
      throw error;
    }
  }
}

class PgTransaction implements DBTransaction<ClientBase> {
  readonly wrappedTransaction: ClientBase;

  constructor(client: ClientBase) {
    this.wrappedTransaction = client;
  }

  async query(sql: string, params: unknown[]): Promise<Row[]> {
    const result = await this.wrappedTransaction.query<Row>(
      sql,
      params as JSONValue[],
    );
    return result.rows;
  }
}

export function pgConnectionProvider(client: Client): () => PgConnection {
  return () => new PgConnection(client);
}

export { handler as GET, handler as POST };