From 24df909f5c1a8954ac426ff5c936d3242522af93 Mon Sep 17 00:00:00 2001
From: Sergey Petushkov <petushkov.sergey@gmail.com>
Date: Fri, 15 Dec 2023 18:04:09 +0100
Subject: [PATCH] fix(aggregations): pass object to the get schema; make sure
 get schema handles recursive deps

---
 .../src/components/stage-wizard/index.tsx     |  4 ++-
 .../src/utils/get-schema.spec.ts              | 18 +++++++++++++
 .../src/utils/get-schema.ts                   | 26 ++++++++++++++-----
 3 files changed, 41 insertions(+), 7 deletions(-)

diff --git a/packages/compass-aggregations/src/components/stage-wizard/index.tsx b/packages/compass-aggregations/src/components/stage-wizard/index.tsx
index e7dbc034e81..8e8061e776a 100644
--- a/packages/compass-aggregations/src/components/stage-wizard/index.tsx
+++ b/packages/compass-aggregations/src/components/stage-wizard/index.tsx
@@ -289,7 +289,9 @@ export default connect(
       return { name, type };
     });
     const previousStageFieldsWithSchema = getSchema(
-      previousStage?.previewDocs ?? []
+      previousStage?.previewDocs?.map((doc) => {
+        return doc.generateObject();
+      }) ?? []
     );
 
     const fields =
diff --git a/packages/compass-aggregations/src/utils/get-schema.spec.ts b/packages/compass-aggregations/src/utils/get-schema.spec.ts
index 294cbb3847e..bbf59b443c5 100644
--- a/packages/compass-aggregations/src/utils/get-schema.spec.ts
+++ b/packages/compass-aggregations/src/utils/get-schema.spec.ts
@@ -189,4 +189,22 @@ describe('get schema', function () {
       expect(getSchema(input)).to.deep.equal(output);
     });
   });
+
+  it("doesn't break when getting schema from recursive object", function () {
+    const a = {
+      b: {
+        c: {
+          get a() {
+            return a;
+          },
+        },
+      },
+    };
+
+    expect(getSchema([a])).to.deep.eq([
+      { name: 'b', type: 'Object' },
+      { name: 'b.c', type: 'Object' },
+      { name: 'b.c.a', type: 'Object' },
+    ]);
+  });
 });
diff --git a/packages/compass-aggregations/src/utils/get-schema.ts b/packages/compass-aggregations/src/utils/get-schema.ts
index 69deeed1b1d..e17df42954a 100644
--- a/packages/compass-aggregations/src/utils/get-schema.ts
+++ b/packages/compass-aggregations/src/utils/get-schema.ts
@@ -22,17 +22,28 @@ const toFieldSchemaWithPrefix = (prefix: string) => {
   });
 };
 
-const getSchemaForObject = (document: Document): DocumentSchema => {
+const getSchemaForObject = (
+  document: Document,
+  seen = new WeakSet()
+): DocumentSchema => {
   const schema: DocumentSchema = [];
+
+  if (seen.has(document)) {
+    return schema;
+  }
+
+  seen.add(document);
+
   for (const key in document) {
     const value = document[key];
+
     schema.push({
       name: key,
       type: TypeChecker.type(value),
     });
 
     if (Array.isArray(value)) {
-      const valueSchema = getSchemaForArray(value).map(
+      const valueSchema = getSchemaForArray(value, seen).map(
         toFieldSchemaWithPrefix(key)
       );
       schema.push(...valueSchema);
@@ -41,7 +52,7 @@ const getSchemaForObject = (document: Document): DocumentSchema => {
       value !== null &&
       !value._bsontype
     ) {
-      const valueSchema = getSchemaForObject(value).map(
+      const valueSchema = getSchemaForObject(value, seen).map(
         toFieldSchemaWithPrefix(key)
       );
       schema.push(...valueSchema);
@@ -50,18 +61,21 @@ const getSchemaForObject = (document: Document): DocumentSchema => {
   return schema;
 };
 
-const getSchemaForArray = (records: Document[]): DocumentSchema => {
+const getSchemaForArray = (
+  records: Document[],
+  seen = new WeakSet()
+): DocumentSchema => {
   const schema: DocumentSchema = [];
 
   for (const record of records) {
     if (Array.isArray(record)) {
-      schema.push(...getSchemaForArray(record));
+      schema.push(...getSchemaForArray(record, seen));
     } else if (
       typeof record === 'object' &&
       record !== null &&
       !record._bsontype
     ) {
-      schema.push(...getSchemaForObject(record));
+      schema.push(...getSchemaForObject(record, seen));
     }
   }