fix(compass-import-export): handle export flake COMPASS-8839 (#6622)

* early return when signal is aborted * test for delay abort * close the stream
mongodb-js · Jan 16, 2025 · bd38dd8 · bd38dd8
1 parent ba5c36f
commit bd38dd8
Show file tree

Hide file tree

Showing 2 changed files with 51 additions and 14 deletions.
diff --git a/packages/compass-import-export/src/export/export-json.spec.ts b/packages/compass-import-export/src/export/export-json.spec.ts
@@ -274,16 +274,52 @@ describe('exportJSON', function () {
     expect(result.docsWritten).to.equal(0);
     expect(result.aborted).to.be.true;
 
-    let resultText;
     try {
-      resultText = await fs.promises.readFile(resultPath, 'utf8');
+      await fs.promises.readFile(resultPath, 'utf8');
+      expect.fail('Expected file to not exist');
     } catch (err) {
-      console.log(resultPath);
-      throw err;
+      // noop
     }
+    // close the stream so that afterEach hook can clear the tmpdir
+    // otherwise it will throw an error (for windows)
+    output.close();
+  });
 
-    const expectedText = '';
-    expect(resultText).to.deep.equal(expectedText);
+  it('responds to abortSignal.aborted - with delayed abort', async function () {
+    const abortController = new AbortController();
+
+    const resultPath = path.join(tmpdir, 'test-abort.exported.ejson');
+
+    const importedText = await fs.promises.readFile(
+      fixtures.json.complex,
+      'utf8'
+    );
+    const ejsonToInsert = EJSON.parse(importedText);
+    await dataService.insertMany(testNS, ejsonToInsert);
+
+    const output = fs.createWriteStream(resultPath);
+    const promise = exportJSONFromQuery({
+      dataService,
+      ns: testNS,
+      output,
+      variant: 'default',
+      abortSignal: abortController.signal,
+    });
+    abortController.abort();
+
+    const result = await promise;
+    const data = await fs.promises.readFile(resultPath, 'utf8');
+
+    try {
+      JSON.parse(data);
+      expect.fail('Expected file to not be valid JSON');
+    } catch (err) {
+      // With signal part of streams pipeline the file is created and if
+      // the signal is aborted the stream is destroyed and file is not
+      // writable anymore and as a result its not able to write trailing ] to the file.
+    }
+    expect(result.aborted).to.be.true;
+    expect(result.docsWritten).to.equal(0);
   });
 
   it('exports aggregations', async function () {

diff --git a/packages/compass-import-export/src/export/export-json.ts b/packages/compass-import-export/src/export/export-json.ts
@@ -52,10 +52,15 @@ export async function exportJSON({
 
   const ejsonOptions = getEJSONOptionsForVariant(variant);
 
-  if (!abortSignal?.aborted) {
-    output.write('[');
+  if (abortSignal?.aborted) {
+    return {
+      docsWritten,
+      aborted: true,
+    };
   }
 
+  output.write('[');
+
   const docStream = new Transform({
     objectMode: true,
     transform: function (chunk: any, encoding, callback) {
@@ -88,21 +93,17 @@ export async function exportJSON({
       [inputStream, docStream, output],
       ...(abortSignal ? [{ signal: abortSignal }] : [])
     );
-    output.write(']\n', 'utf8');
   } catch (err: any) {
     if (err.code === 'ABORT_ERR') {
-      // Finish the JSON file as the `final` of the docs stream didn't run.
-      output.write(']\n', 'utf8');
-
       return {
         docsWritten,
         aborted: true,
       };
     }
-    output.write(']\n', 'utf8');
-
     throw err;
   } finally {
+    // Finish the array output
+    output.write(']\n');
     void input.close();
   }