feat(compass-aggregations): refactor for get-schema to also add type information - COMPASS-6793 (#4355)

Author: himanshusinghs

packages/compass-aggregations/src/components/aggregation-side-panel/stage-wizard-use-cases/index.ts

@@ -6,12 +6,9 @@ import BasicGroupUseCase from './group/basic-group';
 import GroupWithStatistics from './group/group-with-statistics';
 import MatchUseCase from './match/match';
 import GroupWithSubset from './group/group-with-subset';
-import type { TypeCastTypes } from 'hadron-type-checker';
+import type { FieldSchema } from '../../../utils/get-schema';
 
-export type StageWizardFields = {
-  name: string;
-  type: TypeCastTypes;
-}[];
+export type StageWizardFields = FieldSchema[];
 
 export type WizardComponentProps = {
   fields: StageWizardFields;

packages/compass-aggregations/src/components/stage-wizard/index.tsx

@@ -27,6 +27,8 @@ import type {
 import { getSchema } from '../../utils/get-schema';
 import { getStageHelpLink } from '../../utils/stage';
 import type { SortableProps } from '../pipeline-builder-workspace/pipeline-builder-ui-workspace/sortable-list';
+import type { FieldSchema } from '../../utils/get-schema';
+import type { TypeCastTypes } from 'hadron-type-checker';
 
 const containerStyles = css({
   display: 'flex',
@@ -178,21 +180,30 @@ export default connect(
       .reverse()
       .find((x): x is StoreStage => x.type === 'stage' && !x.disabled);
 
-    const mappedInitialFields = initialFields.map(
-      (x: { name: string; description: string }) => ({
-        name: x.name,
-        // parsed schema has the bson type Object replaced with
-        // Document to avoid collision with JS Objects but that
-        // shouldn't be a problem for us because we use these
-        // as string values alongside well defined casters.
-        type: x.description === 'Document' ? 'Object' : x.description,
-      })
+    const mappedInitialFields = (
+      initialFields as {
+        name: string;
+        description?: string;
+      }[]
+    ).map<FieldSchema>(({ name, description }) => {
+      // parsed schema has the bson type Object replaced with
+      // Document to avoid collision with JS Objects but that
+      // shouldn't be a problem for us because we use these
+      // as string values alongside well defined casters.
+      const type =
+        description === 'Document'
+          ? 'Object'
+          : ((description ?? 'String') as TypeCastTypes);
+
+      return { name, type };
+    });
+    const previousStageFieldsWithSchema = getSchema(
+      previousStage?.previewDocs ?? []
     );
-    const previousStageFields = getSchema(previousStage?.previewDocs ?? []);
 
     const fields =
-      previousStageFields.length > 0 && autoPreview
-        ? previousStageFields.map((name) => ({ name, type: 'String' }))
+      previousStageFieldsWithSchema.length > 0 && autoPreview
+        ? previousStageFieldsWithSchema
         : mappedInitialFields;
 
     return {

packages/compass-aggregations/src/modules/collections-fields.ts

@@ -184,7 +184,7 @@ export const fetchCollectionFields = (
         collection,
         data: {
           ...collectionInfo,
-          fields: getSchema(documents),
+          fields: getSchema(documents).map(({ name }) => name),
           isLoading: false,
         },
       });

packages/compass-aggregations/src/utils/get-schema.spec.ts

@@ -3,6 +3,11 @@ import bson from 'bson';
 import { getSchema } from './get-schema';
 
 const DATA = [
+  {
+    useCase: 'does nothing when input is empty',
+    input: [],
+    output: [],
+  },
   {
     useCase: '_id is always the first one',
     input: [
@@ -11,7 +16,16 @@ const DATA = [
         _id: 123456,
       },
     ],
-    output: ['_id', 'data'],
+    output: [
+      {
+        name: '_id',
+        type: 'Int32',
+      },
+      {
+        name: 'data',
+        type: 'String',
+      },
+    ],
   },
   {
     useCase: 'simple json object',
@@ -21,7 +35,50 @@ const DATA = [
         data: 'hello',
       },
     ],
-    output: ['data', 'name'],
+    output: [
+      {
+        name: 'data',
+        type: 'String',
+      },
+      {
+        name: 'name',
+        type: 'String',
+      },
+    ],
+  },
+  {
+    useCase: 'simple json object with falsy values',
+    input: [
+      {
+        name: '',
+        downloads: 0,
+        popular: false,
+        phoneNumbers: [],
+        addresses: null,
+      },
+    ],
+    output: [
+      {
+        name: 'addresses',
+        type: 'Null',
+      },
+      {
+        name: 'downloads',
+        type: 'Int32',
+      },
+      {
+        name: 'name',
+        type: 'String',
+      },
+      {
+        name: 'phoneNumbers',
+        type: 'Array',
+      },
+      {
+        name: 'popular',
+        type: 'Boolean',
+      },
+    ],
   },
   {
     useCase: 'nested json object',
@@ -39,13 +96,13 @@ const DATA = [
       },
     ],
     output: [
-      'address',
-      'address.city',
-      'address.street',
-      'address.street.name',
-      'address.street.number',
-      'data',
-      'name',
+      { name: 'address', type: 'Object' },
+      { name: 'address.city', type: 'String' },
+      { name: 'address.street', type: 'Object' },
+      { name: 'address.street.name', type: 'String' },
+      { name: 'address.street.number', type: 'Int32' },
+      { name: 'data', type: 'String' },
+      { name: 'name', type: 'String' },
     ],
   },
   {
@@ -77,29 +134,35 @@ const DATA = [
       },
     ],
     output: [
-      'data',
-      'name',
-      'streets',
-      'streets._id',
-      'streets.city',
-      'streets.name',
-      'streets.number',
-      'streets.zip',
+      { name: 'data', type: 'String' },
+      { name: 'name', type: 'String' },
+      { name: 'streets', type: 'Array' },
+      { name: 'streets._id', type: 'Int32' },
+      { name: 'streets.city', type: 'String' },
+      { name: 'streets.name', type: 'String' },
+      { name: 'streets.number', type: 'Int32' },
+      { name: 'streets.zip', type: 'Int32' },
     ],
   },
   {
     useCase: 'handles bson values',
     input: [
       {
         _id: new bson.ObjectId(),
-        data: new bson.Int32(123),
+        data: new bson.Double(123),
         address: {
           street: 'Alt-Moabit',
           number: new bson.Int32(18),
         },
       },
     ],
-    output: ['_id', 'address', 'address.number', 'address.street', 'data'],
+    output: [
+      { name: '_id', type: 'ObjectId' },
+      { name: 'address', type: 'Object' },
+      { name: 'address.number', type: 'Int32' },
+      { name: 'address.street', type: 'String' },
+      { name: 'data', type: 'Double' },
+    ],
   },
   {
     useCase: 'nested array with scaler values',
@@ -112,7 +175,11 @@ const DATA = [
         },
       },
     ],
-    output: ['meta', 'meta.common', 'meta.common.artists'],
+    output: [
+      { name: 'meta', type: 'Object' },
+      { name: 'meta.common', type: 'Object' },
+      { name: 'meta.common.artists', type: 'Array' },
+    ],
   },
 ];
 

packages/compass-aggregations/src/utils/get-schema.ts

@@ -1,41 +1,75 @@
+import TypeChecker from 'hadron-type-checker';
+import { sortedUniqBy, sortBy } from 'lodash';
+
+import type { TypeCastTypes } from 'hadron-type-checker';
 import type { Document } from 'mongodb';
 
-const getArrayKeys = (records: Document[]) => {
-  return records
-    .filter((x) => typeof x === 'object')
-    .map((item) => getObjectKeys(item))
-    .flat()
-    .filter((x, i, a) => a.indexOf(x) === i)
-    .sort();
+export type FieldSchema = {
+  name: string;
+  type: TypeCastTypes;
+};
+
+export type DocumentSchema = FieldSchema[];
+
+/**
+ * Mapper function that maps a FieldSchema with the name prefixed with the
+ * provided prefix, separated by a dot
+ * */
+const toFieldSchemaWithPrefix = (prefix: string) => {
+  return (fieldSchema: FieldSchema): FieldSchema => ({
+    name: `${prefix}.${fieldSchema.name}`,
+    type: fieldSchema.type,
+  });
 };
 
-const getObjectKeys = (record: Document) => {
-  const keys: string[] = [];
+const getSchemaForObject = (document: Document): DocumentSchema => {
+  const schema: DocumentSchema = [];
+  for (const key in document) {
+    const value = document[key];
+    schema.push({
+      name: key,
+      type: TypeChecker.type(value),
+    });
 
-  if (!record) {
-    return keys;
+    if (Array.isArray(value)) {
+      const valueSchema = getSchemaForArray(value).map(
+        toFieldSchemaWithPrefix(key)
+      );
+      schema.push(...valueSchema);
+    } else if (
+      typeof value === 'object' &&
+      value !== null &&
+      !value._bsontype
+    ) {
+      const valueSchema = getSchemaForObject(value).map(
+        toFieldSchemaWithPrefix(key)
+      );
+      schema.push(...valueSchema);
+    }
   }
+  return schema;
+};
+
+const getSchemaForArray = (records: Document[]): DocumentSchema => {
+  const schema: DocumentSchema = [];
 
-  for (const key in record) {
-    keys.push(key);
-    const value = record[key];
-
-    if (value && typeof value === 'object') {
-      const isBson = value._bsontype;
-      if (!isBson) {
-        const nestedKeys = Array.isArray(value)
-          ? getArrayKeys(value)
-          : getObjectKeys(value);
-        nestedKeys.forEach((nestedKey) => {
-          keys.push(`${key}.${nestedKey}`);
-        });
-      }
+  for (const record of records) {
+    if (Array.isArray(record)) {
+      schema.push(...getSchemaForArray(record));
+    } else if (
+      typeof record === 'object' &&
+      record !== null &&
+      !record._bsontype
+    ) {
+      schema.push(...getSchemaForObject(record));
     }
   }
 
-  return keys;
+  return schema;
 };
 
-export const getSchema = (data: Document[]) => {
-  return getArrayKeys(data);
+export const getSchema = (records: Document[]): DocumentSchema => {
+  const schema = getSchemaForArray(records);
+  const sortedSchema = sortBy(schema, 'name');
+  return sortedUniqBy(sortedSchema, 'name');
 };