feat: analyze a whole database and find its relationships [WIP] [Skunkworks]

package.json

@@ -56,6 +56,7 @@
     "@types/mocha": "^10.0.1",
     "@types/node": "^18.11.18",
     "@types/reservoir": "^0.1.0",
+    "@types/yargs": "^17.0.29",
     "@typescript-eslint/eslint-plugin": "^5.47.1",
     "@typescript-eslint/parser": "^5.47.1",
     "bson": "^5.0.1",

scripts/analyze-database.ts

@@ -0,0 +1,78 @@
+#!/usr/bin/env npx ts-node
+
+import { MongoClient, Document } from 'mongodb';
+import yargs from 'yargs';
+import { hideBin } from 'yargs/helpers';
+
+import type { Schema, Relationship } from '../src';
+import { SchemaAnalyzer, findRelationshipsForSchema } from '../src';
+
+async function analyzeCollection(documents: AsyncIterable<Document>) {
+  const analyzer = new SchemaAnalyzer({
+    storeValues: true
+  });
+  for await (const doc of documents) {
+    analyzer.analyzeDoc(doc);
+  }
+  return analyzer;
+}
+
+let client: MongoClient;
+async function run() {
+  const argv = await yargs(hideBin(process.argv))
+    .option('sampleSize', { type: 'number', default: 1000 })
+    .argv;
+
+  const [uri, databaseName] = argv._ as [string, string];
+  if (!(uri && databaseName)) {
+    throw new Error('USAGE: analyze-database.ts connectionURI databaseName');
+  }
+
+  client = new MongoClient(uri);
+  await client.connect();
+
+  const db = client.db(databaseName);
+
+  const collectionInfos = await db.listCollections().toArray();
+  console.dir(collectionInfos);
+
+  const collections: Record<string, Schema> = {};
+
+  const relationships: Relationship[] = [];
+
+  const collectionNames = collectionInfos.map((c) => c.name);
+
+  for (const coll of collectionInfos) {
+    console.log(coll.name);
+    const collection = db.collection(coll.name);
+    const cursor = collection.aggregate([{
+      $sample: {
+        size: argv.sampleSize
+      }
+    }], {
+      allowDiskUse: true
+    });
+
+    const analyzer = await analyzeCollection(cursor);
+
+    const schema = analyzer.getResult();
+    collections[coll.name] = schema;
+
+    relationships.push(...await findRelationshipsForSchema(db, coll.name, collectionNames, schema));
+
+    console.log(); // newline
+  }
+
+  console.dir(relationships, { depth: null });
+}
+
+if (require.main === module) {
+  run()
+    .finally(() => {
+      client?.close();
+    })
+    .catch((err) => {
+      console.error(err.stack);
+      process.exit(1);
+    });
+}

src/database-analyzer.ts

@@ -0,0 +1,95 @@
+
+import type { Schema } from './schema-analyzer';
+import type { Db } from 'mongodb';
+
+type CollectionFieldReference = {
+  collection: string;
+  fieldPath: string[];
+}
+
+type FieldReferenceWithValues = CollectionFieldReference & {
+  values: any[]
+}
+
+export type Relationship = {
+  from: CollectionFieldReference;
+  to: CollectionFieldReference;
+}
+
+function shuffleArray(array: any[]) {
+  for (let i = array.length - 1; i > 0; i--) {
+    const j = Math.floor(Math.random() * (i + 1));
+    [array[i], array[j]] = [array[j], array[i]];
+  }
+}
+
+function findCandidateReferencesForSchema(collectionName: string, schema: Schema) {
+  const candidatePaths: FieldReferenceWithValues[] = [];
+
+  for (const field of schema.fields) {
+    if (field.name === '_id') {
+      continue;
+    }
+
+    // TODO: also consider anything matching a known naming convention like /_id$/
+    // TODO: we might also want to consider any large integers if there are lots of different values?
+
+    const values: any[] = [];
+    for (const typeInfo of field.types) {
+      if (['ObjectId', 'UUID'].includes(typeInfo.bsonType)) {
+        values.push(...(typeInfo as { values: any[]}).values ?? []);
+      }
+    }
+    if (values.length) {
+      // in case the sample came from limit()* and wasn't already sorted randomly
+      shuffleArray(values);
+
+      candidatePaths.push({
+        collection: collectionName,
+        fieldPath: field.path,
+        values
+      });
+      console.log(field.path);
+    }
+  }
+
+  return candidatePaths;
+}
+
+async function findRelationshipsCandidate(db: Db, collectionNames: string[], candidatePaths: FieldReferenceWithValues[]) {
+  const relationships: Relationship[] = [];
+
+  // not the most efficient..
+  for (const { collection, fieldPath, values } of candidatePaths) {
+    for (const target of collectionNames) {
+      const ids = values.slice(0, 10);
+      const result = (await db.collection(target).aggregate([
+        { $match: { _id: { $in: ids } } },
+        { $count: 'matches' }
+      ]).toArray());
+
+      if (result.length) {
+        console.log(collection, fieldPath, result);
+        relationships.push({
+          from: {
+            collection,
+            fieldPath
+          },
+          to: {
+            collection: target,
+            fieldPath: ['_id']
+          }
+        });
+        // no point checking the collections - we assume this is a many to one
+        break;
+      }
+    }
+  }
+
+  return relationships;
+}
+
+export async function findRelationshipsForSchema(db: Db, collectionName: string, collectionNames: string[], schema: Schema) {
+  const candidatePaths = findCandidateReferencesForSchema(collectionName, schema);
+  return await findRelationshipsCandidate(db, collectionNames, candidatePaths);
+}

src/index.ts

@@ -23,6 +23,12 @@ import type {
   SimplifiedSchema
 } from './schema-analyzer';
 import * as schemaStats from './stats';
+import type {
+  Relationship
+} from './database-analyzer';
+import {
+  findRelationshipsForSchema
+} from './database-analyzer';
 
 type MongoDBCursor = AggregationCursor | FindCursor;
 
@@ -109,7 +115,8 @@ export type {
   SimplifiedSchemaDocumentType,
   SimplifiedSchemaType,
   SimplifiedSchemaField,
-  SimplifiedSchema
+  SimplifiedSchema,
+  Relationship
 };
 
 export {
@@ -119,5 +126,6 @@ export {
   getSchemaPaths,
   getSimplifiedSchema,
   SchemaAnalyzer,
-  schemaStats
+  schemaStats,
+  findRelationshipsForSchema
 };