Skip to content

Commit bccc14c

Browse files
Merge remote-tracking branch 'origin/main' into beta-releases
2 parents ee9f64a + cf56288 commit bccc14c

File tree

15 files changed

+140
-698
lines changed

15 files changed

+140
-698
lines changed

THIRD-PARTY-NOTICES.md

Lines changed: 1 addition & 364 deletions
Large diffs are not rendered by default.

package-lock.json

Lines changed: 8 additions & 2 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

packages/compass-generative-ai/scripts/ai-accuracy-tests/ai-accuracy-tests.ts

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,10 @@ import util from 'util';
2929
import { execFile as callbackExecFile } from 'child_process';
3030
import decomment from 'decomment';
3131

32+
import {
33+
validateAIQueryResponse,
34+
validateAIAggregationResponse,
35+
} from '../../src/atlas-ai-service';
3236
import { loadFixturesToDB } from './fixtures';
3337
import type { Fixtures } from './fixtures';
3438
import { AtlasAPI } from './ai-backend';
@@ -229,6 +233,10 @@ const runOnce = async (
229233
if (assertResult) {
230234
let cursor;
231235

236+
type === 'query'
237+
? validateAIQueryResponse(response)
238+
: validateAIAggregationResponse(response);
239+
232240
if (
233241
type === 'aggregation' ||
234242
(type === 'query' &&

packages/compass-generative-ai/src/atlas-ai-service.ts

Lines changed: 103 additions & 98 deletions
Original file line numberDiff line numberDiff line change
@@ -94,6 +94,107 @@ function buildQueryOrAggregationMessageBody(
9494
return msgBody;
9595
}
9696

97+
function hasExtraneousKeys(obj: any, expectedKeys: string[]) {
98+
return Object.keys(obj).some((key) => !expectedKeys.includes(key));
99+
}
100+
101+
export function validateAIQueryResponse(
102+
response: any
103+
): asserts response is AIQuery {
104+
const { content } = response ?? {};
105+
106+
if (typeof content !== 'object' || content === null) {
107+
throw new Error('Unexpected response: expected content to be an object');
108+
}
109+
110+
if (hasExtraneousKeys(content, ['query', 'aggregation'])) {
111+
throw new Error(
112+
'Unexpected keys in response: expected query and aggregation'
113+
);
114+
}
115+
116+
const { query, aggregation } = content;
117+
118+
if (!query && !aggregation) {
119+
throw new Error(
120+
'Unexpected response: expected query or aggregation, got none'
121+
);
122+
}
123+
124+
if (query && typeof query !== 'object') {
125+
throw new Error('Unexpected response: expected query to be an object');
126+
}
127+
128+
if (
129+
hasExtraneousKeys(query, [
130+
'filter',
131+
'project',
132+
'collation',
133+
'sort',
134+
'skip',
135+
'limit',
136+
])
137+
) {
138+
throw new Error(
139+
'Unexpected keys in response: expected filter, project, collation, sort, skip, limit, aggregation'
140+
);
141+
}
142+
143+
for (const field of [
144+
'filter',
145+
'project',
146+
'collation',
147+
'sort',
148+
'skip',
149+
'limit',
150+
]) {
151+
if (query[field] && typeof query[field] !== 'string') {
152+
throw new Error(
153+
`Unexpected response: expected field ${field} to be a string, got ${JSON.stringify(
154+
query[field],
155+
null,
156+
2
157+
)}`
158+
);
159+
}
160+
}
161+
162+
if (aggregation && typeof aggregation.pipeline !== 'string') {
163+
throw new Error(
164+
`Unexpected response: expected aggregation pipeline to be a string, got ${JSON.stringify(
165+
aggregation,
166+
null,
167+
2
168+
)}`
169+
);
170+
}
171+
}
172+
173+
export function validateAIAggregationResponse(
174+
response: any
175+
): asserts response is AIAggregation {
176+
const { content } = response;
177+
178+
if (typeof content !== 'object' || content === null) {
179+
throw new Error('Unexpected response: expected content to be an object');
180+
}
181+
182+
if (hasExtraneousKeys(content, ['aggregation'])) {
183+
throw new Error('Unexpected keys in response: expected aggregation');
184+
}
185+
186+
if (content.aggregation && typeof content.aggregation.pipeline !== 'string') {
187+
// Compared to queries where we will always get the `query` field, for
188+
// aggregations backend deletes the whole `aggregation` key if pipeline is
189+
// empty, so we only validate `pipeline` key if `aggregation` key is present
190+
throw new Error(
191+
`Unexpected response: expected aggregation to be a string, got ${String(
192+
content.aggregation.pipeline
193+
)}`
194+
);
195+
}
196+
}
197+
97198
export class AtlasAiService {
98199
private initPromise: Promise<void> | null = null;
99200

@@ -240,110 +341,18 @@ export class AtlasAiService {
240341
return this.getQueryOrAggregationFromUserInput(
241342
AGGREGATION_URI,
242343
input,
243-
this.validateAIAggregationResponse.bind(this)
344+
validateAIAggregationResponse
244345
);
245346
}
246347

247348
async getQueryFromUserInput(input: GenerativeAiInput) {
248349
return this.getQueryOrAggregationFromUserInput(
249350
QUERY_URI,
250351
input,
251-
this.validateAIQueryResponse.bind(this)
352+
validateAIQueryResponse
252353
);
253354
}
254355

255-
private validateAIQueryResponse(response: any): asserts response is AIQuery {
256-
const { content } = response ?? {};
257-
258-
if (typeof content !== 'object' || content === null) {
259-
throw new Error('Unexpected response: expected content to be an object');
260-
}
261-
262-
if (this.hasExtraneousKeys(content, ['query', 'aggregation'])) {
263-
throw new Error(
264-
'Unexpected keys in response: expected query and aggregation'
265-
);
266-
}
267-
268-
const { query, aggregation } = content;
269-
270-
if (typeof query !== 'object' || query === null) {
271-
throw new Error('Unexpected response: expected query to be an object');
272-
}
273-
274-
if (
275-
this.hasExtraneousKeys(query, [
276-
'filter',
277-
'project',
278-
'collation',
279-
'sort',
280-
'skip',
281-
'limit',
282-
])
283-
) {
284-
throw new Error(
285-
'Unexpected keys in response: expected filter, project, collation, sort, skip, limit, aggregation'
286-
);
287-
}
288-
289-
for (const field of [
290-
'filter',
291-
'project',
292-
'collation',
293-
'sort',
294-
'skip',
295-
'limit',
296-
]) {
297-
if (query[field] && typeof query[field] !== 'string') {
298-
throw new Error(
299-
`Unexpected response: expected field ${field} to be a string, got ${JSON.stringify(
300-
query[field],
301-
null,
302-
2
303-
)}`
304-
);
305-
}
306-
}
307-
308-
if (aggregation && typeof aggregation.pipeline !== 'string') {
309-
throw new Error(
310-
`Unexpected response: expected aggregation pipeline to be a string, got ${JSON.stringify(
311-
aggregation,
312-
null,
313-
2
314-
)}`
315-
);
316-
}
317-
}
318-
319-
private validateAIAggregationResponse(
320-
response: any
321-
): asserts response is AIAggregation {
322-
const { content } = response;
323-
324-
if (typeof content !== 'object' || content === null) {
325-
throw new Error('Unexpected response: expected content to be an object');
326-
}
327-
328-
if (this.hasExtraneousKeys(content, ['aggregation'])) {
329-
throw new Error('Unexpected keys in response: expected aggregation');
330-
}
331-
332-
if (
333-
content.aggregation &&
334-
typeof content.aggregation.pipeline !== 'string'
335-
) {
336-
// Compared to queries where we will always get the `query` field, for
337-
// aggregations backend deletes the whole `aggregation` key if pipeline is
338-
// empty, so we only validate `pipeline` key if `aggregation` key is present
339-
throw new Error(
340-
`Unexpected response: expected aggregation to be a string, got ${String(
341-
content.aggregation.pipeline
342-
)}`
343-
);
344-
}
345-
}
346-
347356
private validateAIFeatureEnablementResponse(
348357
response: any
349358
): asserts response is AIFeatureEnablement {
@@ -352,8 +361,4 @@ export class AtlasAiService {
352361
throw new Error('Unexpected response: expected features to be an object');
353362
}
354363
}
355-
356-
private hasExtraneousKeys(obj: any, expectedKeys: string[]) {
357-
return Object.keys(obj).some((key) => !expectedKeys.includes(key));
358-
}
359364
}

packages/compass-import-export/package.json

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -74,8 +74,7 @@
7474
"redux": "^4.2.1",
7575
"redux-thunk": "^2.4.2",
7676
"stream-json": "^1.7.5",
77-
"strip-bom-stream": "^4.0.0",
78-
"temp": "^0.9.4"
77+
"strip-bom-stream": "^4.0.0"
7978
},
8079
"devDependencies": {
8180
"@mongodb-js/compass-test-server": "^0.1.16",
@@ -106,6 +105,7 @@
106105
"react-dom": "^17.0.2",
107106
"sinon": "^9.2.3",
108107
"sinon-chai": "^3.7.0",
108+
"temp": "^0.9.4",
109109
"typescript": "^5.0.4",
110110
"xvfb-maybe": "^0.2.1"
111111
},

packages/compass-import-export/src/export/export-csv.ts

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,6 @@ import fs from 'fs';
22
import { EJSON } from 'bson';
33
import type { Document } from 'bson';
44
import { pipeline } from 'stream/promises';
5-
import temp from 'temp';
65
import { Transform } from 'stream';
76
import type { Readable, Writable } from 'stream';
87
import toNS from 'mongodb-ns';
@@ -11,6 +10,8 @@ import type { PreferencesAccess } from 'compass-preferences-model/provider';
1110
import { capMaxTimeMSAtPreferenceLimit } from 'compass-preferences-model/provider';
1211
import Parser from 'stream-json/Parser';
1312
import StreamValues from 'stream-json/streamers/StreamValues';
13+
import path from 'path';
14+
import os from 'os';
1415

1516
import { lookupValueForPath, ColumnRecorder } from './export-utils';
1617
import {
@@ -31,6 +32,12 @@ import type { AggregationCursor, FindCursor } from 'mongodb';
3132

3233
const debug = createDebug('export-csv');
3334

35+
const generateTempFilename = (suffix: string) => {
36+
const randomString = Math.random().toString(36).substring(2, 15);
37+
const filename = `temp-${randomString}${suffix}`;
38+
return path.join(os.tmpdir(), filename);
39+
};
40+
3441
// First we download all the docs for the query/aggregation to a temporary file
3542
// while determining the unique set of columns we'll need and their order
3643
// (DOWNLOAD), then we write the header row, then process that temp file in
@@ -223,7 +230,7 @@ async function loadEJSONFileAndColumns({
223230
// while simultaneously determining the unique set of columns in the order
224231
// we'll have to write to the file.
225232
const inputStream = cursor.stream();
226-
const filename = temp.path({ suffix: '.jsonl' });
233+
const filename = generateTempFilename('.jsonl');
227234
const output = fs.createWriteStream(filename);
228235

229236
const columnStream = new ColumnStream(progressCallback);

packages/compass-import-export/src/modules/export.spec.ts

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,15 +1,12 @@
11
import os from 'os';
22
import { expect } from 'chai';
3-
import temp from 'temp';
43
import fs from 'fs';
54
import path from 'path';
65
import Sinon from 'sinon';
76
import type { DataService } from 'mongodb-data-service';
87
import { connect } from 'mongodb-data-service';
98
import AppRegistry from 'hadron-app-registry';
109

11-
temp.track();
12-
1310
import {
1411
openExport,
1512
addFieldToExport,

0 commit comments

Comments
 (0)