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

// demo https://drizzle.run/se2noay5mhdu24va3xhv0lqo

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

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

    chunks.push(sql.raw(`'${key}',`));

    // json_build_object formats to ISO 8601 ...
    if (is(value, PgTimestampString)) {
      chunks.push(sql`timezone('UTC', ${value})`);
    } else {
      chunks.push(sql`${value}`);
    }
  });

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

export function jsonAggBuildObject<
  T extends SelectedFields<any, any>,
  Column extends AnyColumn,
>(
  shape: T,
  options?: { orderBy?: { colName: Column; direction: "ASC" | "DESC" } },
) {
  return sql<SelectResultFields<T>[]>`coalesce(
    json_agg(${jsonBuildObject(shape)}
    ${
      options?.orderBy
        ? sql`ORDER BY ${options.orderBy.colName} ${sql.raw(
            options.orderBy.direction,
          )}`
        : undefined
    })
    FILTER (WHERE ${and(
      sql.join(
        Object.values(shape).map((value) => sql`${sql`${value}`} IS NOT NULL`),
        sql` AND `,
      ),
    )})
    ,'${sql`[]`}')`;
}

// with filter non-null + distinct
export function jsonAgg<Column extends AnyColumn>(column: Column) {
  return coalesce<GetColumnData<Column, "raw">[]>(
    sql`json_agg(distinct ${sql`${column}`}) filter (where ${column} is not null)`,
    sql`'[]'`
  );
}

// generalist
export function jsonAgg<Column extends AnyColumn>(column: Column) {
  return coalesce<GetColumnData<Column, "raw">[]>(
    sql`json_agg(${sql`${column}`})`,
    sql`'[]'`
  );
}

// Sometimes you want an array and not a json
export function arrayAgg<Column extends AnyColumn>(column: Column) {
  return sql<
    GetColumnData<Column, "raw">[]
  >`array_agg(distinct ${sql`${column}`}) filter (where ${column} is not null)`;
}

// To be completed
type PGCastTypes = "uuid" | "uuid[]" | "text" | "text[]";

type PGArrayCastTypes = {
  [P in PGCastTypes]: P extends `${infer _T}[]` ? P : never;
}[PGCastTypes];

// Transform an array of values (from a function params) into a postgres array
export function toArray<Values>(values: Values[], cast: PGArrayCastTypes) {
  const chunks: SQL[] = [];

  values.forEach((column) => {
    if (chunks.length > 0) {
      chunks.push(sql.raw(`,`));
    }

    chunks.push(sql`${column}`);
  });

  return sql`array[${sql.join(chunks)}]::${sql.raw(cast)}`;
}
// exemple:
await db
  .select()
  .from(existingDiscussionQuery)
  .where(
    arrayContained(
      // from Drizzle
      existingDiscussionQuery.participants,
      toArray(
        [
          "c3b1399f-2c6b-40d7-9d37-cfaf9a7c6164",
          "77c75084-7123-481b-a326-49c9ebceb431",
        ],
        "uuid[]"
      )
    )
  );

// 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);

sqlite.ts

import { SQL, sql } from "drizzle-orm";
import { SelectResultFields } from "drizzle-orm/query-builders/select.types";
import { SelectedFields } from "drizzle-orm/sqlite-core";

export function jsonObject<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_object(${sql.join(
    chunks,
  )}),  ${sql`json_object()`})`;
}

export function jsonAggObject<T extends SelectedFields>(shape: T) {
  return sql<SelectResultFields<T>[]>`coalesce(json_group_array(${jsonObject(
    shape,
  )}), ${sql`json_array()`})`.mapWith(
    (v) => JSON.parse(v) as SelectResultFields<T>[],
  );
}

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
 * Ex default array: themes: coalesce(pubThemeListQuery.themes, sql`'[]'`)
 * Ex default number: votesCount: coalesce(PubPollAnswersQuery.count, sql`0`)
 */
export function coalesce<T>(value: SQL.Aliased<T> | SQL<T>, defaultValue: SQL) {
  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)}
	  )`;
}

// Like getTableColumns but for sub queries https://orm.drizzle.team/docs/goodies#get-typed-table-columns
/**
* @deprecated - use  subQuery._.selectedFields
*/
export function getSubQueryColumns<
  S extends ColumnsSelection,
  A extends string,
>(subQuery: SubqueryWithSelection<S, A> | WithSubqueryWithSelection<S, A>) {
  const { selection } = subQuery as unknown as {
    selection: (typeof subQuery)["_"]["selectedFields"];
  };

  return selection;
}
// can now used like that
subQuery._.selectedFields

export type InferSubQueryModel<T extends SQL.Aliased> = T["_"]["type"];

type PgColumnJson<T> = PgColumn<ColumnBaseConfig<'json', string> & { data: T }>;
export function isNotNullJsonb<T> (column: PgColumnJson<T>, key: keyof T): SQL | undefined {
  return and(isNotNull(column), isNotNull(sql`${column}->>${key}`));
}
// usage: https://drizzle.run/lm4zy2ohxoc1sxzizzgf1kt0 Author: nowifi4u (discord)


// usage: https://drizzle.run/lkd38uqtk5broj117asmxkah
function mergeJson<
  Colum extends PgColumn,
  CustomType extends Colum["_"]["data"] = Colum["_"]["data"],
>(column: Colum, data: Partial<CustomType> | null): SQL;
function mergeJson<
  Colum extends PgColumn,
  CustomType extends Colum["_"]["data"] = Colum["_"]["data"],
  Key extends keyof CustomType = keyof CustomType,
>(column: Colum, field: Key, data: Partial<CustomType[Key]> | null): SQL;
function mergeJson<
  Colum extends PgColumn,
  CustomType extends Colum["_"]["data"] = Colum["_"]["data"],
  Key extends keyof CustomType = keyof CustomType,
>(
  column: Colum,
  fieldOrData: Key | Partial<CustomType>,
  data?: Partial<CustomType[Key]> | null,
) {
  if (typeof fieldOrData === "string") {
    return sql`jsonb_set(${column}, '{${sql.raw(String(fieldOrData))}}', ${data ? sql`${column} -> '${sql.raw(String(fieldOrData))}' || ${JSON.stringify(data)}` : "null"})`;
  }
  return sql`coalesce(${column}, '{}') || ${fieldOrData ? JSON.stringify(fieldOrData) : null}`;
}

// example
await db.update(config).set({
  roles: mergeJson(config.roles, "president", {
    label: "new-P",
  }),
});

await db.update(config).set({
  roles: mergeJson(config.roles, {
    member: { color: "new-blue", label: "mbr" },
  }),
});

Are those snippets still relevant? or most have been integrated to Drizzle already?