fix: add relationship fields items and anyOf support COMPASS-9649 (#7161)

Author: paula-stacho

packages/compass-data-modeling/src/store/diagram.spec.ts

@@ -7,6 +7,7 @@ import {
   openDiagram,
   redoEdit,
   undoEdit,
+  selectFieldsForCurrentModel,
 } from './diagram';
 import type {
   Edit,
@@ -295,4 +296,150 @@ describe('Data Modeling store', function () {
     expect(diagramAfterRedo.edits[0]).to.deep.include(loadedDiagram.edits[0]);
     expect(diagramAfterRedo.edits[1]).to.deep.include(edit);
   });
+
+  describe('selectFieldsForCurrentModel', function () {
+    it('should select fields from a flat schema', function () {
+      const edits: MongoDBDataModelDescription['edits'] = [
+        {
+          id: 'first-edit',
+          timestamp: new Date().toISOString(),
+          type: 'SetModel',
+          model: {
+            collections: [
+              {
+                ns: 'collection1',
+                indexes: [],
+                displayPosition: [0, 0],
+                shardKey: {},
+                jsonSchema: {
+                  bsonType: 'object',
+                  properties: {
+                    field1: { bsonType: 'string' },
+                    field2: { bsonType: 'int' },
+                    field3: { bsonType: 'int' },
+                  },
+                },
+              },
+            ],
+            relationships: [],
+          },
+        },
+      ];
+      const selectedFields = selectFieldsForCurrentModel(edits);
+
+      expect(selectedFields).to.deep.equal({
+        collection1: [['field1'], ['field2'], ['field3']],
+      });
+    });
+
+    it('should select fields from a nested schema', function () {
+      const edits: MongoDBDataModelDescription['edits'] = [
+        {
+          id: 'first-edit',
+          timestamp: new Date().toISOString(),
+          type: 'SetModel',
+          model: {
+            collections: [
+              {
+                ns: 'collection1',
+                indexes: [],
+                displayPosition: [0, 0],
+                shardKey: {},
+                jsonSchema: {
+                  bsonType: 'object',
+                  properties: {
+                    prop1: { bsonType: 'string' },
+                    // Deeply nested properties
+                    prop2: {
+                      bsonType: 'object',
+                      properties: {
+                        prop2A: { bsonType: 'string' },
+                        prop2B: {
+                          bsonType: 'object',
+                          properties: {
+                            prop2B1: { bsonType: 'string' },
+                            prop2B2: { bsonType: 'int' },
+                          },
+                        },
+                      },
+                    },
+                    // Array of objects
+                    prop3: {
+                      bsonType: 'array',
+                      items: {
+                        bsonType: 'object',
+                        properties: {
+                          prop3A: { bsonType: 'string' },
+                        },
+                      },
+                    },
+                    // Mixed type with objects
+                    prop4: {
+                      anyOf: [
+                        {
+                          bsonType: 'object',
+                          properties: {
+                            prop4A: { bsonType: 'string' },
+                          },
+                        },
+                        {
+                          bsonType: 'object',
+                          properties: {
+                            prop4B: { bsonType: 'string' },
+                          },
+                        },
+                      ],
+                    },
+                    // Mixed array with objects
+                    prop5: {
+                      bsonType: 'array',
+                      items: [
+                        {
+                          bsonType: 'object',
+                          properties: {
+                            prop5A: { bsonType: 'string' },
+                          },
+                        },
+                        {
+                          bsonType: 'object',
+                          properties: {
+                            prop5B: { bsonType: 'number' },
+                          },
+                        },
+                      ],
+                    },
+                  },
+                },
+              },
+            ],
+            relationships: [],
+          },
+        },
+      ];
+      const selectedFields = selectFieldsForCurrentModel(edits);
+
+      expect(selectedFields).to.have.property('collection1');
+      expect(selectedFields.collection1).to.deep.include(['prop1']);
+      expect(selectedFields.collection1).to.deep.include(['prop2']);
+      expect(selectedFields.collection1).to.deep.include(['prop2', 'prop2A']);
+      expect(selectedFields.collection1).to.deep.include([
+        'prop2',
+        'prop2B',
+        'prop2B1',
+      ]);
+      expect(selectedFields.collection1).to.deep.include([
+        'prop2',
+        'prop2B',
+        'prop2B2',
+      ]);
+      expect(selectedFields.collection1).to.deep.include(['prop3']);
+      expect(selectedFields.collection1).to.deep.include(['prop3', 'prop3A']);
+      expect(selectedFields.collection1).to.deep.include(['prop4']);
+      expect(selectedFields.collection1).to.deep.include(['prop4', 'prop4A']);
+      expect(selectedFields.collection1).to.deep.include(['prop4', 'prop4B']);
+      expect(selectedFields.collection1).to.deep.include(['prop5']);
+      expect(selectedFields.collection1).to.deep.include(['prop5', 'prop5A']);
+      expect(selectedFields.collection1).to.deep.include(['prop5', 'prop5B']);
+    });
+  });
 });

packages/compass-data-modeling/src/store/diagram.ts

@@ -610,6 +610,19 @@ function extractFields(
   parentKey?: string[],
   fields: string[][] = []
 ) {
+  if ('anyOf' in parentSchema && parentSchema.anyOf) {
+    for (const schema of parentSchema.anyOf) {
+      extractFields(schema, parentKey, fields);
+    }
+  }
+  if ('items' in parentSchema && parentSchema.items) {
+    const items = Array.isArray(parentSchema.items)
+      ? parentSchema.items
+      : [parentSchema.items];
+    for (const schema of items) {
+      extractFields(schema, parentKey, fields);
+    }
+  }
   if ('properties' in parentSchema && parentSchema.properties) {
     for (const [key, value] of Object.entries(parentSchema.properties)) {
       const fullKey = parentKey ? [...parentKey, key] : [key];