feat(data-modeling): automatically infer relationships from indexes and object ids

Author: gribnoysup

packages/compass-data-modeling/src/store/analysis-process.ts

@@ -3,11 +3,11 @@ import { isAction } from './util';
 import type { DataModelingThunkAction } from './reducer';
 import { analyzeDocuments, type MongoDBJSONSchema } from 'mongodb-schema';
 import { getCurrentDiagramFromState } from './diagram';
-import type { Document } from 'bson';
-import type { AggregationCursor } from 'mongodb';
+import { UUID } from 'bson';
 import type { Relationship } from '../services/data-model-storage';
 import { applyLayout } from '@mongodb-js/diagramming';
 import { collectionToDiagramNode } from '../utils/nodes-and-edges';
+import { inferLocalToForeignRelationshipsForCollection } from './relationships';
 
 export type AnalysisProcessState = {
   currentAnalysisOptions:
@@ -161,18 +161,18 @@ export function startAnalysis(
       options,
     });
     try {
+      let relations: Relationship[] = [];
       const dataService =
         services.connections.getDataServiceForConnection(connectionId);
+
       const collections = await Promise.all(
         namespaces.map(async (ns) => {
-          const sample: AggregationCursor<Document> = dataService.sampleCursor(
+          const sample = await dataService.sample(
             ns,
             { size: 100 },
+            { promoteValues: false },
             {
-              signal: cancelController.signal,
-              promoteValues: false,
-            },
-            {
+              abortSignal: cancelController.signal,
               fallbackReadPreference: 'secondaryPreferred',
             }
           );
@@ -194,12 +194,38 @@ export function startAnalysis(
             type: AnalysisProcessActionTypes.NAMESPACE_SCHEMA_ANALYZED,
             namespace: ns,
           });
-          return { ns, schema };
+          return { ns, schema, sample };
         })
       );
 
       if (options.automaticallyInferRelations) {
-        // TODO
+        relations = (
+          await Promise.all(
+            collections.map(
+              ({
+                ns,
+                schema,
+                sample,
+              }): Promise<Relationship['relationship'][]> => {
+                return inferLocalToForeignRelationshipsForCollection(
+                  ns,
+                  schema,
+                  sample,
+                  collections,
+                  dataService
+                );
+              }
+            )
+          )
+        ).flatMap((relationships) => {
+          return relationships.map((relationship) => {
+            return {
+              id: new UUID().toHexString(),
+              relationship,
+              isInferred: true,
+            };
+          });
+        });
       }
 
       if (cancelController.signal.aborted) {
@@ -229,7 +255,7 @@ export function startAnalysis(
           const position = node ? node.position : { x: 0, y: 0 };
           return { ...coll, position };
         }),
-        relations: [],
+        relations,
       });
 
       services.track('Data Modeling Diagram Created', {

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

@@ -314,7 +314,9 @@ describe('Data Modeling store', function () {
                 jsonSchema: {
                   bsonType: 'object',
                   properties: {
-                    field1: { bsonType: 'string' },
+                    field1: {
+                      anyOf: [{ bsonType: 'string' }, { bsonType: 'int' }],
+                    },
                     field2: { bsonType: 'int' },
                     field3: { bsonType: 'int' },
                   },

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

@@ -714,32 +714,52 @@ export const selectCurrentModelFromState = (state: DataModelingState) => {
   return selectCurrentModel(selectCurrentDiagramFromState(state).edits);
 };
 
-function extractFields(
-  parentSchema: MongoDBJSONSchema,
-  parentKey?: string[],
-  fields: string[][] = []
+/**
+ * A very simplistic depth-first traversing function that only handles a subset
+ * of real JSON schema keywords that is applicable to our MongoDB JSON schema
+ * format
+ */
+export function traverseMongoDBJSONSchema(
+  schema: MongoDBJSONSchema,
+  visitor: (schema: MongoDBJSONSchema, path: string[]) => void,
+  path: string[] = []
 ) {
-  if ('anyOf' in parentSchema && parentSchema.anyOf) {
-    for (const schema of parentSchema.anyOf) {
-      extractFields(schema, parentKey, fields);
+  if (schema.anyOf) {
+    for (const s of schema.anyOf) {
+      traverseMongoDBJSONSchema(s, visitor, path);
     }
+    return;
   }
-  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);
+
+  visitor(schema, path);
+
+  if (schema.items) {
+    for (const s of Array.isArray(schema.items)
+      ? schema.items
+      : [schema.items]) {
+      traverseMongoDBJSONSchema(s, visitor, path);
     }
-  }
-  if ('properties' in parentSchema && parentSchema.properties) {
-    for (const [key, value] of Object.entries(parentSchema.properties)) {
-      const fullKey = parentKey ? [...parentKey, key] : [key];
-      fields.push(fullKey);
-      extractFields(value, fullKey, fields);
+  } else if (schema.properties) {
+    for (const [key, s] of Object.entries(schema.properties)) {
+      traverseMongoDBJSONSchema(s, visitor, [...path, key]);
     }
   }
-  return fields;
+}
+
+function extractFields(schema: MongoDBJSONSchema) {
+  // Keeping fields stringified in a set to make sure that we don't include
+  // duplicates
+  const fields = new Set<string>();
+  traverseMongoDBJSONSchema(schema, (_schema, path) => {
+    // Skip the root path
+    if (path.length === 0) {
+      return;
+    }
+    fields.add(path.join('__FIELD_PATH_SEPARATOR__'));
+  });
+  return Array.from(fields.values(), (pathStr) => {
+    return pathStr.split('__FIELD_PATH_SEPARATOR__');
+  });
 }
 
 function getFieldsForCurrentModel(