How to return the entity with its relations after saving it?

I am building a storage application, with GraphQL as the backend, using Typegraphql and TypeORM.

The categories need to be added separately and then when adding a product, you choose from a dropdown one of the available categories. This in turn passes the categoryId to the product in a one-to-many/many-to-one relationship.

Here is my Category entity:

import {
  Entity,
  PrimaryColumn,
  Column,
  BaseEntity,
  Generated,
  OneToMany
} from 'typeorm';
import Product from './Product';

@ObjectType()
@Entity('categories')
export default class Category extends BaseEntity {
  @Field()
  @PrimaryColumn()
  @Generated('uuid')
  categoryId: string;

  @Field()
  @Column()
  categoryName: string;

  @OneToMany(() => Product, (product: Product) => product.category)
  products: Product[];
}

and here is my Product entity

import {
  Entity,
  PrimaryColumn,
  Column,
  BaseEntity,
  Generated,
  ManyToOne,
  JoinColumn
} from 'typeorm';
import Category from './Category';

@ObjectType()
@Entity('products')
export default class Product extends BaseEntity {
  @Field()
  @PrimaryColumn()
  @Generated('uuid')
  productID: string;

  @Field()
  @Column()
  productName: string;

  @Field(() => Category)
  @ManyToOne(() => Category, (category: Category) => category.products, {
    cascade: true,
    lazy: true
  })
  @JoinColumn()
  category: Category;

  @Field()
  @Column()
  productQuantity: number;

  @Field()
  @Column({ type: 'decimal', precision: 2 })
  productPrice: number;

  @Field()
  @Column({ type: 'decimal', precision: 2 })
  productPriceRA: number;

  @Field()
  @Column({ type: 'decimal', precision: 2 })
  productPriceKK: number;

  @Field()
  @Column('varchar', { length: 255 })
  productSupplier: string;

  @Field()
  @Column('varchar', { length: 255 })
  productOrderLink: string;

  @Field()
  @Column('longtext')
  productImage: string;
}

For the save mutation, I've created an Input type as well:

export default class ProductInput implements Partial<Product> {
    @Field()
    productName: string;

    @Field(() => String)
    category: Category;

    @Field()
    productQuantity: number;

    @Field()
    productPrice: number;

    @Field()
    productPriceRA: number;

    @Field()
    productPriceKK: number;

    @Field()
    productSupplier: string;

    @Field()
    productOrderLink: string;

    @Field()
    productImage: string;
}

The relations work, as I am able to query the products, along with their category data with the following query:

{
  getProducts {
    productID
    productName
    category {
       categoryId
       categoryName
    }
  }
}

However, when saving a product it always returns "message": "Cannot return null for non-nullable field Category.categoryName."

This is the Mutation's code in the Resolver:

@Mutation(() => Product, { description: 'Add new product' })
  async addProduct(
    @Arg('product') productInput: ProductInput
  ): Promise<Product | any> {
    try {
      const product = await Product.create(productInput).save();

      console.log('product: ', product);
      return product;
    } catch (error) {
      return error;
    }
  }

I've been trying different things, however nothing seems to work and I am wondering if it's even possible to directly return the entity with its relations. If it's not, the other option I can think of is to return true/false based on the result and re-query all of the data. But this seems very inefficient and I am actively trying to avoid going this route.

Any help will be much appreciated.

After some more research and I decided to go with the following approach:

try {
      const { productID } = await Product.create(productInput).save();

      return await Product.findOne(productID);
    } catch (error) {
      return error;
    }

This allows me to directly return the product, based on the productID after it's saved in the database and properly returns the object with it's relationship.