Drizzle snippets

clear-db.ts

// Credits to Louistiti from Drizzle Discord: https://discord.com/channels/1043890932593987624/1130802621750448160/1143083373535973406

import { sql } from 'drizzle-orm';

const clearDb = async (): Promise<void> => {
    const query = sql<string>`SELECT table_name
      FROM information_schema.tables
      WHERE table_schema = 'public'
        AND table_type = 'BASE TABLE';
    `;

    const tables = await db.execute(query);
    
    // @LauraKirby
    for (let table of tables.rows) {
        const query = sql.raw(`TRUNCATE TABLE ${table.table_name} CASCADE;`);
        await db.execute(query);
    }

//     previous version
//     for (let table of tables) {
//       const query = sql.raw(`TRUNCATE TABLE ${table.table_name} CASCADE;`);
//       await db.execute(query);
//     }
  }

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() });

jsonAgg_jsonAggBuildObject.ts

import {
	sql,
	type AnyColumn,
	type SQL,
	type InferSelectModel,
} from "drizzle-orm";
import {
	type SelectedFields,
	type PgTable,
	type TableConfig,
} from "drizzle-orm/pg-core";
import { type SelectResultFields } from "node_modules/drizzle-orm/query-builders/select.types";

export function jsonBuildObject<T extends SelectedFields>(shape: T) {
    const chunks: SQL[] = [];

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

    return sql<SelectResultFields<T>>`coalesce(json_build_object(${sql.join(
        chunks,
    )}), '{}')`;
}

export function jsonAggBuildObject<
    T extends SelectedFields,
    Column extends AnyColumn,
>(
    shape: T,
    options?: { orderBy?: { colName: Column; direction: "ASC" | "DESC" } },
) {
    return sql<SelectResultFields<T>[]>`coalesce(jsonb_agg(${jsonBuildObject(
        shape,
    )}${
        options?.orderBy
            ? sql`order by ${options.orderBy.colName} ${sql.raw(
                  options.orderBy.direction,
              )}`
            : undefined
    }), '[]')`;
}

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

// generalist
export function jsonAgg<Table extends AnyTable<TableConfig>>(table: Table) {
    return sql<InferModel<Table>[]>`coalesce(json_agg(${table}), '[]')`;
}

// 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,
	type SQL,
	type InferSelectModel,
} from "drizzle-orm";
import {
	type SelectedFields,
	type PgTable,
	type TableConfig,
} from "drizzle-orm/pg-core";
import { type SelectResultFields } from "node_modules/drizzle-orm/query-builders/select.types";

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})`;
}

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

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

export function count<Column extends AnyColumn>(column: Column) {
    return sql<number>`cast(count(${column}) as integer)`;
}

/**
 * Coalesce a value to a default value if the value is null
 *
 */
export function coalesce<T>(value: SQL.Aliased<T>, defaultValue: T) {
	return sql<T>`coalesce(${value}, ${defaultValue})`;
}

type Unit = "minutes" | "minute";
type Operator = "+" | "-";

export function now(interval?: `${Operator} interval ${number} ${Unit}`) {
	return sql<string>`now() ${interval || ""}`;
}


// 
// example where table.data type is { id: string; type: 'document' | 'image' }
// eq(eqJsonb(table.data, {
//                              id: 'some value',
//                              type: "document",
//                          }))
export function eqJsonb<T extends PgColumn>(
    column: T,
    value: Partial<GetColumnData<T, "raw">>,
) {
    return sql`${column} @> ${value}`;
}

// Select a JSONB field
// example:
// const results = await db
//            .select({
//                myField: pickJsonbField<
//                    MyDataType, // the one you use for jsonb("data").$type<MyDataType>().notNull(),
//                    "fieldKey" // one of MyDataType
//                >(table.data, "fieldKey"),
//            })
export function pickJsonbField<
    U,
    K extends keyof U,
    T extends PgColumn = PgColumn,
>(column: T, field: K, cast?: "uuid") {
    return sql<U[K]>`((${column}->${field})${
        cast ? sql.raw(`::${cast}`) : undefined
    })`;
}

// .where(inJsonArray(subQueryWithJsonAggregate.anArray, "keyName", [valueFromParams]))
export function inJsonArray<T extends SQL.Aliased<unknown[]>>(
    jsonArray: T,
    key: keyof T["_"]["type"][number],
    values: string[],
) {
    const element = sql.raw(`${String(key)}_array_element`);

    return sql`EXISTS (
		SELECT 1
		FROM jsonb_array_elements(${jsonArray}) AS ${element}
		WHERE ${inArray(sql`${element}->>${key}`, values)}
	  )`;
}