Merge remote-tracking branch 'origin/main' into beta-releases

mongodb-js · Jun 5, 2024 · bccc14c · bccc14c
2 parents ee9f64a + cf56288
commit bccc14c
Show file tree

Hide file tree

Showing 15 changed files with 140 additions and 698 deletions.
diff --git a/THIRD-PARTY-NOTICES.md b/THIRD-PARTY-NOTICES.md
diff --git a/package-lock.json b/package-lock.json
diff --git a/packages/compass-generative-ai/scripts/ai-accuracy-tests/ai-accuracy-tests.ts b/packages/compass-generative-ai/scripts/ai-accuracy-tests/ai-accuracy-tests.ts
@@ -29,6 +29,10 @@ import util from 'util';
 import { execFile as callbackExecFile } from 'child_process';
 import decomment from 'decomment';
 
+import {
+  validateAIQueryResponse,
+  validateAIAggregationResponse,
+} from '../../src/atlas-ai-service';
 import { loadFixturesToDB } from './fixtures';
 import type { Fixtures } from './fixtures';
 import { AtlasAPI } from './ai-backend';
@@ -229,6 +233,10 @@ const runOnce = async (
     if (assertResult) {
       let cursor;
 
+      type === 'query'
+        ? validateAIQueryResponse(response)
+        : validateAIAggregationResponse(response);
+
       if (
         type === 'aggregation' ||
         (type === 'query' &&

diff --git a/packages/compass-generative-ai/src/atlas-ai-service.ts b/packages/compass-generative-ai/src/atlas-ai-service.ts
@@ -94,6 +94,107 @@ function buildQueryOrAggregationMessageBody(
   return msgBody;
 }
 
+function hasExtraneousKeys(obj: any, expectedKeys: string[]) {
+  return Object.keys(obj).some((key) => !expectedKeys.includes(key));
+}
+
+export function validateAIQueryResponse(
+  response: any
+): asserts response is AIQuery {
+  const { content } = response ?? {};
+
+  if (typeof content !== 'object' || content === null) {
+    throw new Error('Unexpected response: expected content to be an object');
+  }
+
+  if (hasExtraneousKeys(content, ['query', 'aggregation'])) {
+    throw new Error(
+      'Unexpected keys in response: expected query and aggregation'
+    );
+  }
+
+  const { query, aggregation } = content;
+
+  if (!query && !aggregation) {
+    throw new Error(
+      'Unexpected response: expected query or aggregation, got none'
+    );
+  }
+
+  if (query && typeof query !== 'object') {
+    throw new Error('Unexpected response: expected query to be an object');
+  }
+
+  if (
+    hasExtraneousKeys(query, [
+      'filter',
+      'project',
+      'collation',
+      'sort',
+      'skip',
+      'limit',
+    ])
+  ) {
+    throw new Error(
+      'Unexpected keys in response: expected filter, project, collation, sort, skip, limit, aggregation'
+    );
+  }
+
+  for (const field of [
+    'filter',
+    'project',
+    'collation',
+    'sort',
+    'skip',
+    'limit',
+  ]) {
+    if (query[field] && typeof query[field] !== 'string') {
+      throw new Error(
+        `Unexpected response: expected field ${field} to be a string, got ${JSON.stringify(
+          query[field],
+          null,
+          2
+        )}`
+      );
+    }
+  }
+
+  if (aggregation && typeof aggregation.pipeline !== 'string') {
+    throw new Error(
+      `Unexpected response: expected aggregation pipeline to be a string, got ${JSON.stringify(
+        aggregation,
+        null,
+        2
+      )}`
+    );
+  }
+}
+
+export function validateAIAggregationResponse(
+  response: any
+): asserts response is AIAggregation {
+  const { content } = response;
+
+  if (typeof content !== 'object' || content === null) {
+    throw new Error('Unexpected response: expected content to be an object');
+  }
+
+  if (hasExtraneousKeys(content, ['aggregation'])) {
+    throw new Error('Unexpected keys in response: expected aggregation');
+  }
+
+  if (content.aggregation && typeof content.aggregation.pipeline !== 'string') {
+    // Compared to queries where we will always get the `query` field, for
+    // aggregations backend deletes the whole `aggregation` key if pipeline is
+    // empty, so we only validate `pipeline` key if `aggregation` key is present
+    throw new Error(
+      `Unexpected response: expected aggregation to be a string, got ${String(
+        content.aggregation.pipeline
+      )}`
+    );
+  }
+}
+
 export class AtlasAiService {
   private initPromise: Promise<void> | null = null;
 
@@ -240,110 +341,18 @@ export class AtlasAiService {
     return this.getQueryOrAggregationFromUserInput(
       AGGREGATION_URI,
       input,
-      this.validateAIAggregationResponse.bind(this)
+      validateAIAggregationResponse
     );
   }
 
   async getQueryFromUserInput(input: GenerativeAiInput) {
     return this.getQueryOrAggregationFromUserInput(
       QUERY_URI,
       input,
-      this.validateAIQueryResponse.bind(this)
+      validateAIQueryResponse
     );
   }
 
-  private validateAIQueryResponse(response: any): asserts response is AIQuery {
-    const { content } = response ?? {};
-
-    if (typeof content !== 'object' || content === null) {
-      throw new Error('Unexpected response: expected content to be an object');
-    }
-
-    if (this.hasExtraneousKeys(content, ['query', 'aggregation'])) {
-      throw new Error(
-        'Unexpected keys in response: expected query and aggregation'
-      );
-    }
-
-    const { query, aggregation } = content;
-
-    if (typeof query !== 'object' || query === null) {
-      throw new Error('Unexpected response: expected query to be an object');
-    }
-
-    if (
-      this.hasExtraneousKeys(query, [
-        'filter',
-        'project',
-        'collation',
-        'sort',
-        'skip',
-        'limit',
-      ])
-    ) {
-      throw new Error(
-        'Unexpected keys in response: expected filter, project, collation, sort, skip, limit, aggregation'
-      );
-    }
-
-    for (const field of [
-      'filter',
-      'project',
-      'collation',
-      'sort',
-      'skip',
-      'limit',
-    ]) {
-      if (query[field] && typeof query[field] !== 'string') {
-        throw new Error(
-          `Unexpected response: expected field ${field} to be a string, got ${JSON.stringify(
-            query[field],
-            null,
-            2
-          )}`
-        );
-      }
-    }
-
-    if (aggregation && typeof aggregation.pipeline !== 'string') {
-      throw new Error(
-        `Unexpected response: expected aggregation pipeline to be a string, got ${JSON.stringify(
-          aggregation,
-          null,
-          2
-        )}`
-      );
-    }
-  }
-
-  private validateAIAggregationResponse(
-    response: any
-  ): asserts response is AIAggregation {
-    const { content } = response;
-
-    if (typeof content !== 'object' || content === null) {
-      throw new Error('Unexpected response: expected content to be an object');
-    }
-
-    if (this.hasExtraneousKeys(content, ['aggregation'])) {
-      throw new Error('Unexpected keys in response: expected aggregation');
-    }
-
-    if (
-      content.aggregation &&
-      typeof content.aggregation.pipeline !== 'string'
-    ) {
-      // Compared to queries where we will always get the `query` field, for
-      // aggregations backend deletes the whole `aggregation` key if pipeline is
-      // empty, so we only validate `pipeline` key if `aggregation` key is present
-      throw new Error(
-        `Unexpected response: expected aggregation to be a string, got ${String(
-          content.aggregation.pipeline
-        )}`
-      );
-    }
-  }
-
   private validateAIFeatureEnablementResponse(
     response: any
   ): asserts response is AIFeatureEnablement {
@@ -352,8 +361,4 @@ export class AtlasAiService {
       throw new Error('Unexpected response: expected features to be an object');
     }
   }
-
-  private hasExtraneousKeys(obj: any, expectedKeys: string[]) {
-    return Object.keys(obj).some((key) => !expectedKeys.includes(key));
-  }
 }
diff --git a/packages/compass-import-export/package.json b/packages/compass-import-export/package.json
@@ -74,8 +74,7 @@
     "redux": "^4.2.1",
     "redux-thunk": "^2.4.2",
     "stream-json": "^1.7.5",
-    "strip-bom-stream": "^4.0.0",
-    "temp": "^0.9.4"
+    "strip-bom-stream": "^4.0.0"
   },
   "devDependencies": {
     "@mongodb-js/compass-test-server": "^0.1.16",
@@ -106,6 +105,7 @@
     "react-dom": "^17.0.2",
     "sinon": "^9.2.3",
     "sinon-chai": "^3.7.0",
+    "temp": "^0.9.4",
     "typescript": "^5.0.4",
     "xvfb-maybe": "^0.2.1"
   },

diff --git a/packages/compass-import-export/src/export/export-csv.ts b/packages/compass-import-export/src/export/export-csv.ts
@@ -2,7 +2,6 @@ import fs from 'fs';
 import { EJSON } from 'bson';
 import type { Document } from 'bson';
 import { pipeline } from 'stream/promises';
-import temp from 'temp';
 import { Transform } from 'stream';
 import type { Readable, Writable } from 'stream';
 import toNS from 'mongodb-ns';
@@ -11,6 +10,8 @@ import type { PreferencesAccess } from 'compass-preferences-model/provider';
 import { capMaxTimeMSAtPreferenceLimit } from 'compass-preferences-model/provider';
 import Parser from 'stream-json/Parser';
 import StreamValues from 'stream-json/streamers/StreamValues';
+import path from 'path';
+import os from 'os';
 
 import { lookupValueForPath, ColumnRecorder } from './export-utils';
 import {
@@ -31,6 +32,12 @@ import type { AggregationCursor, FindCursor } from 'mongodb';
 
 const debug = createDebug('export-csv');
 
+const generateTempFilename = (suffix: string) => {
+  const randomString = Math.random().toString(36).substring(2, 15);
+  const filename = `temp-${randomString}${suffix}`;
+  return path.join(os.tmpdir(), filename);
+};
+
 // First we download all the docs for the query/aggregation to a temporary file
 // while determining the unique set of columns we'll need and their order
 // (DOWNLOAD), then we write the header row, then process that temp file in
@@ -223,7 +230,7 @@ async function loadEJSONFileAndColumns({
   // while simultaneously determining the unique set of columns in the order
   // we'll have to write to the file.
   const inputStream = cursor.stream();
-  const filename = temp.path({ suffix: '.jsonl' });
+  const filename = generateTempFilename('.jsonl');
   const output = fs.createWriteStream(filename);
 
   const columnStream = new ColumnStream(progressCallback);

diff --git a/packages/compass-import-export/src/modules/export.spec.ts b/packages/compass-import-export/src/modules/export.spec.ts
@@ -1,15 +1,12 @@
 import os from 'os';
 import { expect } from 'chai';
-import temp from 'temp';
 import fs from 'fs';
 import path from 'path';
 import Sinon from 'sinon';
 import type { DataService } from 'mongodb-data-service';
 import { connect } from 'mongodb-data-service';
 import AppRegistry from 'hadron-app-registry';
 
-temp.track();
-
 import {
   openExport,
   addFieldToExport,