Polymorphic Associations in NestJS with PostgreSQL & Drizzle ORM

Introduction

In a database, it’s not uncommon for a single entity to relate to several similar types of entities. For example, a user’s comment could be associated with a photo, an article, or an audio file. The most intuitive solution for this scenario might be to create a specific comments table for each content type, such as PhotoComment and ArticleComment. However, this approach leads to code duplication, as the structure of a comment on a photo is similar to that of a comment on an article.

To simplify this design and avoid duplication, we can use polymorphic associations, a design pattern that allows a single table to associate with multiple different tables, allowing one entity, such as a comment, to link to various types of entities.

Implementing Polymorphic Associations

Incorrect Approach

A simple solution to implement a polymorphic association would be to use a single column in the comments table, such as targetId, which points to the ID of either a photo or an article. However, this approach is unreliable, as PostgreSQL treats targetId as an ordinary number and cannot guarantee that it points to a valid record in the photos or articles tables. Therefore, if a photo or an article is deleted, we would need to remember to manually delete all related comments.

import { serial, text, pgTable, integer } from 'drizzle-orm/pg-core';

export const photos = pgTable('photos', {
  id: serial('id').primaryKey(),
  imageUrl: text('photo_url').notNull(),
});

export const articles = pgTable('articles', {
  id: serial('id').primaryKey(),
  title: text('title').notNull(),
  content: text('content').notNull(),
});

export const comments = pgTable('comments', {
  id: serial('id').primaryKey(),
  content: text('content').notNull(),
  targetId: integer('target_id'),
});

export const databaseSchema = {
  articles,
  photos,
  comments,
};

An Improved Solution

To improve data integrity, we can split targetId into two separate columns: photoId and articleId. With Drizzle ORM, these columns can be related to their respective tables and defined as foreign keys. This strategy creates independent constraints for each relationship, and PostgreSQL can ensure that only one of these columns has a non-null value in each row.

import { serial, text, pgTable, integer } from 'drizzle-orm/pg-core';
import { relations } from 'drizzle-orm';

export const photos = pgTable('photos', {
  id: serial('id').primaryKey(),
  imageUrl: text('photo_url').notNull(),
});

export const articles = pgTable('articles', {
  id: serial('id').primaryKey(),
  title: text('title').notNull(),
  content: text('content').notNull(),
});

export const comments = pgTable('comments', {
  id: serial('id').primaryKey(),
  content: text('content').notNull(),
  photoId: integer('photo_id').references(() => photos.id),
  articleId: integer('article_id').references(() => articles.id),
});

export const commentsRelations = relations(comments, ({ one }) => ({
  photo: one(photos, {
    fields: [comments.photoId],
    references: [photos.id],
  }),
  article: one(articles, {
    fields: [comments.articleId],
    references: [articles.id],
  }),
}));

export const databaseSchema = {
  articles,
  photos,
  comments,
  commentsRelations,
};

In the SQL migration generated by Drizzle ORM, the constraint ensures that photoId and articleId columns allow null values. This enables a comment to be related to either a photo or an article, but never both at the same time. PostgreSQL’s num_nonnulls function guarantees that exactly one of the columns has a non-null value:

ALTER TABLE "comments"
ADD CONSTRAINT check_if_only_one_is_not_null CHECK 
(num_nonnulls("photo_id", "article_id") = 1);

Data Validation

To ensure that the user provides only one of the foreign keys (photoId or articleId), we can create a custom decorator in class-validator. This decorator checks that one of the fields has a valid integer value and that only one of them is present.

import {
  IsNotEmpty,
  IsString,
  registerDecorator,
  ValidationArguments,
} from 'class-validator';

const idKeys: (keyof CreateCommentDto)[] = ['photoId', 'articleId'];

export function ContainsValidForeignKeys() {
  return function (object: object, propertyName: string) {
    registerDecorator({
      name: 'containsValidForeignKeys',
      target: object.constructor,
      propertyName: propertyName,
      options: {
        message: `You need to provide exactly one of the following properties: ${idKeys.join(', ')}`,
      },
      validator: {
        validate(value: unknown, validationArguments: ValidationArguments) {
          const comment = validationArguments.object as CreateCommentDto;

          if (value && !Number.isInteger(value)) {
            return false;
          }

          const includedIdKeys = idKeys.filter((key) => comment[key]);
          return includedIdKeys.length === 1;
        },
      },
    });
  };
}

export class CreateCommentDto {
  @IsString()
  @IsNotEmpty()
  content: string;

  @ContainsValidForeignKeys()
  photoId?: number;

  @ContainsValidForeignKeys()
  articleId?: number;
}

This decorator ensures that the user provides only one photoId or articleId per comment. Thus, if a user tries to create a comment without any valid reference or with both, an appropriate error message is generated.

Error Handling in the Service

To enhance user experience, we implement error handling to capture constraint violations and send appropriate messages:

import { BadRequestException, Injectable } from '@nestjs/common';
import { DrizzleService } from '../database/drizzle.service';
import { databaseSchema } from '../database/database-schema';
import { CreateCommentDto } from './dto/create-comment.dto';
import { isDatabaseError } from '../database/databse-error';
import { PostgresErrorCode } from '../database/postgres-error-code.enum';

@Injectable()
export class CommentsService {
  constructor(private readonly drizzleService: DrizzleService) {}

  async create(comment: CreateCommentDto) {
    try {
      const createdComments = await this.drizzleService.db
        .insert(databaseSchema.comments)
        .values(comment)
        .returning();

      return createdComments.pop();
    } catch (error) {
      if (!isDatabaseError(error)) {
        throw error;
      }

      if (error.code === PostgresErrorCode.ForeignKeyViolation) {
        throw new BadRequestException('Provide a valid foreign key');
      }

      if (error.code === PostgresErrorCode.CheckViolation) {
        throw new BadRequestException('Provide exactly one foreign key');
      }
      throw error;
    }
  }
}

Conclusion

In this article, we implemented a polymorphic association using Drizzle ORM and PostgreSQL to manage comments linked to photos or articles, ensuring database integrity. This design pattern simplifies the schema by reducing the number of tables, offering flexibility in entity relationships. However, it’s essential to use this strategy carefully, as it can make debugging and maintenance more challenging for developers unfamiliar with it.