Batch Updates / Update Many

[ryanking1809]

I've written a batch update function using drizzle. It feels a little messy and I was curious if anyone knew a simpler way to do this?

https://orm.drizzle.team/docs/update#with-update-clause looks promising but you can't seem to pass it multiple rows.

Desired SQL:

WITH update_values (id, username, email) AS (
  VALUES 
    ($1::UUID, $2::VARCHAR, $3::VARCHAR),
    ($4::UUID, $5::VARCHAR, $6::VARCHAR)
)
UPDATE users t
SET username = update_values.username, 
    email = update_values.email
FROM update_values
WHERE t.id = update_values.id

params = [
  '123e4567-e89b-12d3-a456-426614174000',
  'john_doe',
  '[email protected]',
  '98765432-e89b-12d3-a456-426614174000',
  'jane_smith',
  '[email protected]',
]

The Code:

update(db, users, [
  {
    id: "123e4567-e89b-12d3-a456-426614174000",
    username: "john_doe",
    email: "[email protected]",
  },
  {
    id: "98765432-e89b-12d3-a456-426614174000",
    username: "jane_smith",
    email: "[email protected]",
  },
]);

import { sql, SQL } from "drizzle-orm";
import { PgTable, PgColumn, getTableConfig, PgDialect } from "drizzle-orm/pg-core";
const pgDialect = new PgDialect();

type TableWithModel<T> = T & {
  id: PgColumn;
};

type UpdateRow = {
  id: string | number;
  [key: string]: any;
};

const generateUpdateStatements = <T extends TableWithModel<PgTable>>(table: T, columns: string[]): string =>
  columns.map((col) => `"${table[col].name}" = update_values."${table[col].name}"`).join(", ");

const generateValues = <T extends PgTable>(table: T, updates: UpdateRow[]): SQL => {
  const tableConfig = getTableConfig(table);

  const rows = updates.map(
    (row) =>
      sql`(${sql.join(
        Object.entries(row).map(([key, value]) => {
          const column = table[key] as PgColumn;
          if (!column) {
            throw new Error(`Column ${key} not found in table ${tableConfig.name}`);
          }
          return sql`${value}::${sql.raw(column.columnType.replace("Pg", ""))}`;
        }),
        sql`, `,
      )})`,
  );
  return sql.join(rows, sql`, `);
};
export async function update<T extends TableWithModel<PgTable>>(db: any, table: T, updates: UpdateRow[]) {
  if (updates.length === 0) return;
  const updateColumns = Object.keys(updates[0]!).filter((key) => key !== "id");

  const values = generateValues(table, updates);
  const updateStatements = generateUpdateStatements(table, updateColumns);

  const tableName = getTableConfig(table).name;

  const query = sql`
        WITH update_values (id, ${sql.raw(updateColumns.map((col) => `"${table[col].name}"`).join(", "))}) AS (
        VALUES ${values}
        )
        UPDATE ${sql.raw(tableName)} t
        SET ${sql.raw(updateStatements)}
        FROM update_values
        WHERE t.id = update_values.id
    `;
  return db.execute(query);
}

export async function batchUpdate<T extends TableWithModel<PgTable>>(
  db: any,
  table: T,
  updates: UpdateRow[],
  batchSize = 500,
) {
  for (let i = 0; i < updates.length; i += batchSize) {
    const batch = updates.slice(i, i + batchSize);
    await update(db, table, batch);
  }
}

[wilfredjonathanjames]

Seems simpler to run upsert, as insert supports multirow. This is my solution:

...
  await db
    .insert(Table)
    .values(values)
    .onConflictDoUpdate({
      target: Table.id,
      set: conflictUpdateSetAllColumns(Table, ["id"]),
    })
    .returning()

// src/db/helpers.ts
import { sql } from "drizzle-orm"
import { PgUpdateSetSource, PgTable } from "drizzle-orm/pg-core"
import { getTableColumns } from "drizzle-orm"
import { getTableConfig } from "drizzle-orm/pg-core"

export function conflictUpdateSetAllColumns<
  T extends PgTable,
  E extends (keyof T["$inferInsert"])[],
>(table: T, except?: E): PgUpdateSetSource<T> {
  const columns = getTableColumns(table)
  const config = getTableConfig(table)
  const { name: tableName } = config
  const conflictUpdateSet = Object.entries(columns).reduce(
    (acc, [columnName, columnInfo]) => {
      if (except && except.includes(columnName as E[number])) {
        return acc
      }
      if (!columnInfo.default) {
        // @ts-ignore
        acc[columnName] = sql.raw(
          `COALESCE("excluded"."${columnInfo.name}", "${tableName}"."${columnInfo.name}")`,
        )
      }
      return acc
    },
    {},
  ) as PgUpdateSetSource<T>
  return conflictUpdateSet
}

[ryanking1809]

Amazing, thank you!