Drizzle snippets

db.ts

import { sql } from "drizzle-orm";
import { type Logger } from "drizzle-orm/logger";
import { drizzle, type PostgresJsDatabase } from "drizzle-orm/postgres-js";
import postgres from "postgres";

class QueryLogger implements Logger {
  logQuery(query: string, params: unknown[]): void {
    console.debug("___QUERY___");
    console.debug(query);
    console.debug(params);
    console.debug("___END_QUERY___");
  }
}

const client = postgres(process.env.DATABASE_URL!);

export const db = drizzle(client, { logger: new QueryLogger() });

/**
 * Run a transaction.
 *
 * @example
 * await transaction(async (tx) => {
 * const query1 = await tx.insert(...)
 * const query2 = await tx.insert(...)
 * })
 *
 *
 */
export async function transaction<T>(
  run: (tx: PostgresJsDatabase) => Promise<T>
) {
  const client = postgres(process.env.DATABASE_URL!, { max: 1 });
  const tx = drizzle(client, { logger: new QueryLogger() });

  try {
    await tx.execute(sql`BEGIN`);
    const result = await run(tx);
    await tx.execute(sql`COMMIT`);

    return result;
  } catch (err) {
    await tx.execute(sql`ROLLBACK`);
    throw err;
  } finally {
    client.end();
  }
}

jsonAgg_jsonAggBuildObject.ts

function jsonAggBuildObject<T extends Record<string, AnyColumn>>(shape: T) {
  const chunks: SQL[] = [];
  const filters: SQL[] = [];

  Object.entries(shape).forEach(([key, value]) => {
    if (chunks.length > 0) {
      chunks.push(sql.raw(`,`));
    }
    chunks.push(sql.raw(`'${key}',`));
    chunks.push(sql`${value}`);

    if (filters.length > 0) {
      filters.push(sql.raw(` or `));
    }

    filters.push(sql`${value} is not null`);
  });

  return sql<
    InferColumnsDataTypes<T>[]
  >`coalesce(json_agg(json_build_object(${sql.fromList(
    chunks
  )})) filter (where ${sql.fromList(filters)}), '[]')`;
}

export function jsonAgg<Table extends AnyTable<TableConfig>>(table: Table) {
  return sql<
    InferModel<Table>[]
  >`coalesce(json_agg(${table}) filter (where ${table} is not null), '[]')`;
}

// you use it like that:
const result = await db
      .select({
        post,
        // keep only what you need from table theme
        themes: jsonAggBuildObject({
          id: theme.id,
          label: theme.label,
        }),
      })
      .leftJoin(
        postsThemes,
        eq(postsThemes.theme_id, post.theme_id)
      )
      .leftJoin(theme, eq(theme.id, postsThemes.theme_id))
      .groupBy(post.id);

utils.ts

import { sql, type AnyColumn } from "drizzle-orm";

export function takeFirst<T>(items: T[]) {
  return items.at(0);
}

export function takeFirstOrThrow<T>(items: T[]) {
  const first = takeFirst(items);

  if (!first) {
    throw new Error("First item not found");
  }

  return first;
}

export function distinct<Column extends AnyColumn>(column: Column) {
  return sql<Column["_"]["data"]>`distinct(${column})`;
}