Drizzle ORM, deep sub queries

helpers.ts

/* -------------------------------------------------------------------------- */
/*                                 More here;                                 */
/* -------------------------------------------------------------------------- */
//gist.github.com/rphlmr/0d1722a794ed5a16da0fdf6652902b15

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

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

    // 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>>`coalesce(json_build_object(${sql.join(
    chunks
  )}), '{}')`;
}

export function jsonAggBuildObject<T extends SelectedFields>(shape: T) {
  return sql<SelectResultFields<T>[]>`coalesce(jsonb_agg(${jsonBuildObject(
    shape
  )}), '${sql`[]`}')`;
}

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

query.ts

const commentQuery = db
		.select({
			id: distinctOn(comment.id).mapWith(String).as("comment_id"),
			fileId: sql`${comment.fileId}`.mapWith(String).as("file_id"),
			text: comment.text,
			commenter: {
				id: sql`${user.id}`.mapWith(String).as("commenter_id"),
				name: user.name,
			},
		})
		.from(comment)
		.innerJoin(user, eq(comment.commenterId, user.id))
		.orderBy(comment.id)
		.as("comment_query");

	const commentsQuery = db
		.select({
			fileId: commentQuery.fileId,
			comments: jsonAggBuildObject({
				id: commentQuery.id,
				text: commentQuery.text,
				commenter: jsonBuildObject({
					id: commentQuery.commenter.id,
					name: commentQuery.commenter.name,
				}),
			}).as("comments"),
		})
		.from(commentQuery)
		.groupBy(commentQuery.fileId)
		.as("comments_query");

	const tagQuery = db
		.select({
			id: distinctOn(tagsPivot.id).mapWith(String).as("tag_link_id"),
			tagId: sql`${tagsPivot.tagId}`.mapWith(String).as("tag_id"),
			fileId: sql`${tagsPivot.fileId}`
				.mapWith(String)
				.as("tagged_file_id"),
			name: tag.name,
		})
		.from(tagsPivot)
		.innerJoin(tag, eq(tag.id, tagsPivot.tagId))
		.orderBy(tagsPivot.id)
		.as("tag_query");

	const tagsQuery = db
		.select({
			fileId: tagQuery.fileId,
			tags: jsonAggBuildObject({
				id: tagQuery.tagId,
				name: tagQuery.name,
			}).as("tags"),
		})
		.from(tagQuery)
		.groupBy(tagQuery.fileId)
		.as("tags_query");

	const result = await db
		.select({
			...getTableColumns(file),
			comments: commentsQuery.comments,
			tags: tagsQuery.tags,
		})
		.from(file)

		.leftJoin(commentsQuery, eq(commentsQuery.fileId, file.id))
		.leftJoin(tagsQuery, eq(tagsQuery.fileId, file.id));

	const filterByTagId = await db
		.select({
			...getTableColumns(file),
			comments: commentsQuery.comments,
			tags: tagsQuery.tags,
		})
		.from(file)
	        // this line is how to filter
		.where(inJsonArray(tagsQuery.tags, "id", [tagIdFilter]))
		.leftJoin(commentsQuery, eq(commentsQuery.fileId, file.id))
		.leftJoin(tagsQuery, eq(tagsQuery.fileId, file.id));

	console.log(JSON.stringify(result, null, 2));

result.json

[
  {
    "id": "84d88f3c-ff46-4b52-90b8-3d8838fa0c75",
    "name": "Jedi code",
    "comments": [
      {
        "id": "b7190643-4b2d-4412-a438-3e440f4f5322",
        "text": "Learn that my apprentice",
        "commenter": {
          "id": "01d1fd7a-10a8-4760-900d-72b201d51b45",
          "name": "Obi-Wan"
        }
      }
    ],
    "tags": [
      {
        "id": "f19f5f1c-5efb-4e50-ba53-882274529659",
        "name": "padawan"
      }
    ]
  },
  {
    "id": "592998ec-f0d3-43ff-8ed3-ab10ee8522de",
    "name": "Sith code",
    "comments": [
      {
        "id": "46d1b4c3-89dd-49f1-9159-53eb32372b79",
        "text": "Agree I am",
        "commenter": {
          "id": "6c9f3c61-9f18-44ed-90f7-90b97a1776fa",
          "name": "Yoda"
        }
      },
      {
        "id": "e3d3d03c-4ae3-4e38-b866-224e29c35fbb",
        "text": "We should hide that from padawan",
        "commenter": {
          "id": "01d1fd7a-10a8-4760-900d-72b201d51b45",
          "name": "Obi-Wan"
        }
      }
    ],
    "tags": [
      {
        "id": "0917cf03-a56b-4056-bacd-b6ac84e3adf4",
        "name": "knight"
      },
      {
        "id": "330b7583-ce67-4ce0-aacc-8e34030f75f0",
        "name": "master"
      }
    ]
  }
]

schema.ts

export const user = pgTable("user", {
	id: uuid("id").primaryKey().defaultRandom(),
	name: text("name").notNull(),
});

export const file = pgTable("file", {
	id: uuid("id").primaryKey().defaultRandom(),
	name: text("name").notNull(),
});

export const tag = pgTable("tag", {
	id: uuid("id").primaryKey().defaultRandom(),
	name: text("name").notNull(),
});

export const tagsPivot = pgTable(
	"tags_pivot",
	{
		id: uuid("id").notNull().defaultRandom(),
		fileId: uuid("file_id")
			.references(() => file.id)
			.notNull(),
		tagId: uuid("tag_id")
			.references(() => tag.id)
			.notNull(),
	},
	(t) => ({
		cpk: primaryKey({ columns: [t.fileId, t.tagId] }),
	}),
);

export const comment = pgTable("comment", {
	id: uuid("id").primaryKey().defaultRandom(),
	commenterId: uuid("commenter_id")
		.references(() => user.id)
		.notNull(),
	fileId: uuid("file_id")
		.references(() => file.id)
		.notNull(),
	text: text("text").notNull(),
});

inJsonArray now use built-in drizzle helpers + takes an array to make it easier to use.

@rphlmr First, thanks for this. These helpers are awesome.
I came across a use case that I was struggling with, I found the solution while trying to formulate a question on how to do the types for it. Including it here for others and/or for comment if you have suggestions on how to improve it.

Use case: Show [] when the result of a join would result in null. Ex: A file does not have any tags assigned to it.

export function coalesceToArray<T extends SQL.Aliased>(aliased: T) {
  return sql<T['_']['type']>`coalesce(${aliased}, '[]')`;
}

You may know a way to do the types more cleanly but this appears to work.

The use case in context:

        // ...

	const result = await db
		.select({
			...getTableColumns(file),
			comments: commentsQuery.comments,
                         // Here we would have an [] instead of null when the leftJoin() does not have a match
			tags: coalesceToArray(tagsQuery.tags),
		})
		.from(file)

		.leftJoin(commentsQuery, eq(commentsQuery.fileId, file.id))
		.leftJoin(tagsQuery, eq(tagsQuery.fileId, file.id));

        // ...

EDIT: Remove a typo in the return generic

Hey 👋 @austincm you are right, thanks.
You may love this gist : https://gist.github.com/rphlmr/0d1722a794ed5a16da0fdf6652902b15#file-utils-ts-L48

Hey 👋 @austincm you are right, thanks. You may love this gist : https://gist.github.com/rphlmr/0d1722a794ed5a16da0fdf6652902b15#file-utils-ts-L48

Thanks! Supper useful stuff. I fixed a typo in my solution but I think I like the one you have documented better overall.

Hmm interesting, is there any way to have jsonAggBuildObject support having all the columns from a subquery spread into it?

    const { embedding, ...mediaColumns } = getTableColumns(media);

    const mediaQuery = db
      .select(mediaColumns)
      .from(attachments)
      .innerJoin(media, eq(media.id, attachments.mediaId))
      .where(eq(attachments.entityId, submissions.id))
      .as("media");

    // Doesn't work
    const query = db
      .select({
        media: jsonAggBuildObject({
          ...mediaQuery,
        }),
      })
      .from(submissions)

      .limit(3);

Hmm interesting, is there any way to have jsonAggBuildObject support having all the columns from a subquery spread into it?

    const { embedding, ...mediaColumns } = getTableColumns(media);

    const mediaQuery = db
      .select(mediaColumns)
      .from(attachments)
      .innerJoin(media, eq(media.id, attachments.mediaId))
      .where(eq(attachments.entityId, submissions.id))
      .as("media");

    // Doesn't work
    const query = db
      .select({
        media: jsonAggBuildObject({
          ...mediaQuery,
        }),
      })
      .from(submissions)

      .limit(3);

Good question. The hardest part would be to find the good Typescript definition to handle both ways + be able to extract the subquery result shape.

@lithdew

function getSubQueryColumns<S extends ColumnsSelection, A extends string>(
    subQuery: SubqueryWithSelection<S, A> | WithSubqueryWithSelection<S, A>,
) {
    const { selection } = (subQuery as any)[SubqueryConfig] as {
        selection: (typeof subQuery)["_"]["selectedFields"];
    };

    return selection;
}

    const { embedding, ...mediaColumns } = getTableColumns(media);

    const mediaQuery = db
      .select(mediaColumns)
      .from(attachments)
      .innerJoin(media, eq(media.id, attachments.mediaId))
      .where(eq(attachments.entityId, submissions.id))
      .as("media");

    // Doesn't work
    const query = db
      .select({
        media: jsonAggBuildObject({
          ...getSubQueryColumns(mediaQuery),
        }),
      })
      .from(submissions)

      .limit(3);

Maybe 👀