From 4ec5fe07689eacb0456f531d69368b0451ce69a1 Mon Sep 17 00:00:00 2001
From: Doug Martin <doug@dougamartin.com>
Date: Fri, 9 Apr 2021 15:42:53 -0600
Subject: [PATCH] feat(typegoose): Implement `setRelations` to set many
 references

---
 .../services/typegoose-query.service.spec.ts  | 69 +++++++++++++++++++
 .../src/services/reference-query.service.ts   | 19 +++++
 2 files changed, 88 insertions(+)

diff --git a/packages/query-typegoose/__tests__/services/typegoose-query.service.spec.ts b/packages/query-typegoose/__tests__/services/typegoose-query.service.spec.ts
index a6d2cfbf1..c71159b58 100644
--- a/packages/query-typegoose/__tests__/services/typegoose-query.service.spec.ts
+++ b/packages/query-typegoose/__tests__/services/typegoose-query.service.spec.ts
@@ -1269,6 +1269,75 @@ describe('TypegooseQueryService', () => {
     });
   });
 
+  describe('#setRelations', () => {
+    it('set all relations on the entity', async () => {
+      const entity = TEST_ENTITIES[0];
+      const queryService = moduleRef.get(TestEntityService);
+      const relationIds = TEST_REFERENCES.slice(3, 6).map((r) => r._id);
+      const queryResult = await queryService.setRelations(
+        'testReferences',
+        entity._id.toHexString(),
+        relationIds.map((id) => id.toHexString()),
+      );
+      expect(queryResult).toEqual(
+        expect.objectContaining({
+          _id: entity._id,
+          testReferences: expect.arrayContaining(relationIds),
+        }),
+      );
+
+      const relations = await queryService.queryRelations(TestReference, 'testReferences', entity, {});
+      expect(relations.map((r) => r._id)).toEqual(relationIds);
+    });
+
+    it('should remove all relations if the relationIds is empty', async () => {
+      const entity = TEST_ENTITIES[0];
+      const queryService = moduleRef.get(TestEntityService);
+      const queryResult = await queryService.setRelations('testReferences', entity._id.toHexString(), []);
+      expect(queryResult).toEqual(
+        expect.objectContaining({
+          _id: entity._id,
+          testReferences: expect.arrayContaining([]),
+        }),
+      );
+
+      const relations = await queryService.queryRelations(TestReference, 'testReferences', entity, {});
+      expect(relations.map((r) => r._id)).toEqual([]);
+    });
+
+    describe('with modify options', () => {
+      it('should throw an error if the entity is not found with the id and provided filter', async () => {
+        const entity = TEST_ENTITIES[0];
+        const queryService = moduleRef.get(TestEntityService);
+        return expect(
+          queryService.setRelations(
+            'testReferences',
+            entity._id.toHexString(),
+            TEST_REFERENCES.slice(3, 6).map((r) => r._id.toHexString()),
+            {
+              filter: { stringType: { eq: TEST_ENTITIES[1].stringType } },
+            },
+          ),
+        ).rejects.toThrow(`Unable to find TestEntity with id: ${String(entity._id)}`);
+      });
+
+      it('should throw an error if the relations are not found with the relationIds and provided filter', async () => {
+        const entity = TEST_ENTITIES[0];
+        const queryService = moduleRef.get(TestEntityService);
+        return expect(
+          queryService.setRelations<TestReference>(
+            'testReferences',
+            entity._id.toHexString(),
+            TEST_REFERENCES.slice(3, 6).map((r) => r._id.toHexString()),
+            {
+              relationFilter: { referenceName: { like: '%-one' } },
+            },
+          ),
+        ).rejects.toThrow('Unable to find all testReferences to set on TestEntity');
+      });
+    });
+  });
+
   describe('#setRelation', () => {
     it('call select and return the result', async () => {
       const entity = TEST_REFERENCES[0];
diff --git a/packages/query-typegoose/src/services/reference-query.service.ts b/packages/query-typegoose/src/services/reference-query.service.ts
index ff15b79f6..7a9734693 100644
--- a/packages/query-typegoose/src/services/reference-query.service.ts
+++ b/packages/query-typegoose/src/services/reference-query.service.ts
@@ -221,6 +221,25 @@ export abstract class ReferenceQueryService<Entity extends Base> {
     return entity;
   }
 
+  async setRelations<Relation>(
+    relationName: string,
+    id: string,
+    relationIds: (string | number)[],
+    opts?: ModifyRelationOptions<Entity, Relation>,
+  ): Promise<DocumentType<Entity>> {
+    this.checkForReference('AddRelations', relationName, false);
+    const refCount = await this.getRefCount(relationName, relationIds, opts?.relationFilter);
+    if (relationIds.length !== refCount) {
+      throw new Error(`Unable to find all ${relationName} to set on ${this.Model.modelName}`);
+    }
+    const entity = await this.findAndUpdate(
+      id,
+      opts?.filter as Filter<Entity>,
+      { [relationName]: relationIds } as UpdateQuery<DocumentType<Entity>>,
+    );
+    return entity;
+  }
+
   async setRelation<Relation>(
     relationName: string,
     id: string | number,