Type safety on JSONB fields from PostgreSQL

[DenisBessa]

Hi,

I really like Drizzle, and I think I managed to solve the problem to get nested values from JSONB fields in a type-safe way.

Here is my code:

import { StringChunk, SQL, sql } from "drizzle-orm";

type NestedKeyOf<ObjectType extends object> = {
  [Key in keyof ObjectType & (string | number)]: ObjectType[Key] extends object ? `${Key}` | `${Key}.${NestedKeyOf<ObjectType[Key]>}` : `${Key}`;
}[keyof ObjectType & (string | number)];

function jsonbField<T extends object, U>(column: AnyColumn, path: NestedKeyOf<T>) {
  const pathParts = path.split(".");
  const emptySQL = new SQL([]);
  const query = pathParts.reduce((acc, part, index) => {
    if (index === 0) {
      return new SQL([new StringChunk(column.name), new StringChunk("::jsonb"), new StringChunk(`->'${part}'`)]);
    }
    return new SQL([acc, new StringChunk(`->'${part}'`)]);
  }, emptySQL);

  return sql<U>`${query}`;
}

and I can use it like this:

const data = await db.select({
      jsonbField: jsonbField<JsonbDataType, string>(table.column, "nested.field.property") // the second parameter is type-safe
    })
(...)

[stereobooster]

I wonder if it possible to take this idea to the next level. Something like this:

schema.ts:

// zod just as an example - it can be any suitable schema
const metaSchema = z.object({
  username: z.string(),
});

export const document = sqliteTable("documents", {
  meta: json("meta", metaSchema),

And then query could look like this:

const data = await db.select({ username: document.meta.username })
// or with special operator to avoid conflicts with other methods/fields
// kind of reminds `json_extract` syntax '$.username'
const data = await db.select({ username: document.meta.$.username })
// still can be used as any other column
const data = await db.select({ meta: document.meta })

PS

• duplicate Made an auxiliary type to get nested values inside json objects on PostgreSQL #785 ?
• related https://orm.drizzle.team/docs/column-types/sqlite#customizing-column-data-type
• related https://orm.drizzle.team/docs/custom-types

[DenisBessa]

I have made a few improvements on the function that generates a type-safe JSONB field on Drizzle. Here they are:

import type { ColumnBaseConfig } from "drizzle-orm";
import { SQL, StringChunk } from "drizzle-orm";
import type { PgColumn } from "drizzle-orm/pg-core";

type NestedKeyOf<ObjectType> = {
  [Key in keyof ObjectType & (string | number)]: ObjectType[Key] extends Array<infer ArrayType>
    ? `${Key}` | `${Key}.${number}` | (ArrayType extends object ? `${Key}.${number}.${NestedKeyOf<ArrayType>}` : never)
    : ObjectType[Key] extends object
    ? `${Key}` | `${Key}.${NestedKeyOf<ObjectType[Key]>}`
    : `${Key}`;
}[keyof ObjectType & (string | number)];

// Helper type to recursively infer the type at a given path.
type AtPath<T, Path extends string> = Path extends `${infer Key}.${infer Rest}`
  ? Key extends keyof T
    ? Rest extends NestedKeyOf<T[Key]>
      ? AtPath<T[Key], Rest>
      : never
    : Key extends `${number}`
    ? T extends Array<infer ArrayType>
      ? ArrayType extends object
        ? AtPath<ArrayType, Rest>
        : never
      : never
    : never
  : Path extends keyof T
  ? T[Path]
  : never;

export function jsonbField<T extends PgColumn<ColumnBaseConfig<"json", "PgJsonb">>, P extends NestedKeyOf<T["_"]["data"]>>(column: T, path: P) {
  const pathParts = path.split(".");
  let concatenatedSql = "";

  pathParts.forEach((part, index) => {
    if (index === pathParts.length - 1) {
      concatenatedSql += ` ->> '${part}'`;
      return;
    }
    concatenatedSql += `-> '${part}'`;
  });

  return new SQL<AtPath<T["_"]["data"], P>>([column, new StringChunk(concatenatedSql)]);
}

Here's how to use it:

  const data = await db
    .select({
      column: jsonbField(table.column, "nested.field"), // this is type safe
    })
    .from(table)

(you need to define the type of the column in the schema definition with $type function).

[CanRau]

The same for MySQL json fields would be fantastic to have 😍