From fc0e7d2b190ce64b8be519b58a161687ff1aab90 Mon Sep 17 00:00:00 2001
From: Maurizio Casimirri <maurizio.cas@gmail.com>
Date: Wed, 20 Oct 2021 08:59:08 +0200
Subject: [PATCH] feat: add db.sql helper for ADL $sql MONGOSH-900 (#1100)

* wip: add db.sql

* fix server versions

* Add doc for db.sql

* Update packages/i18n/src/locales/en_US.ts

Mark SQL as experimental in docs

* add warning and error message

* fixup: skip session test

Co-authored-by: Anna Henningsen <anna@addaleax.net>
---
 packages/i18n/src/locales/en_US.ts      |  5 +++
 packages/shell-api/src/database.spec.ts | 52 ++++++++++++++++++++++++-
 packages/shell-api/src/database.ts      | 34 ++++++++++++++++
 3 files changed, 90 insertions(+), 1 deletion(-)

diff --git a/packages/i18n/src/locales/en_US.ts b/packages/i18n/src/locales/en_US.ts
index 61df1072f..3faaa2a39 100644
--- a/packages/i18n/src/locales/en_US.ts
+++ b/packages/i18n/src/locales/en_US.ts
@@ -946,6 +946,11 @@ const translations: Catalog = {
         help: {
           description: 'Database Class',
           attributes: {
+            sql: {
+              link: 'https://docs.mongodb.com/manual/reference/method/db.sql',
+              description: '(Experimental) Runs a SQL query against Atlas Data Lake. Note: this is an experimental feature that may be subject to change in future releases.',
+              example: 'const cursor = db.sql("SELECT * FROM myCollection", options)'
+            },
             watch: {
               link: 'https://docs.mongodb.com/manual/reference/method/db.watch',
               description: 'Opens a change stream cursor on the database',
diff --git a/packages/shell-api/src/database.spec.ts b/packages/shell-api/src/database.spec.ts
index 15273c132..01529a7b2 100644
--- a/packages/shell-api/src/database.spec.ts
+++ b/packages/shell-api/src/database.spec.ts
@@ -2495,6 +2495,55 @@ describe('Database', () => {
         expect.fail('Failed to throw');
       });
     });
+    describe('sql', () => {
+      it('runs a $sql aggregation', async() => {
+        const serviceProviderCursor = stubInterface<ServiceProviderAggCursor>();
+        serviceProvider.aggregateDb.returns(serviceProviderCursor as any);
+        await database.sql('SELECT * FROM somecollection;', { options: true });
+        expect(serviceProvider.aggregateDb).to.have.been.calledWith(
+          database._name,
+          [{
+            $sql: {
+              dialect: 'mongosql',
+              format: 'jdbc',
+              formatVersion: 1,
+              statement: 'SELECT * FROM somecollection;'
+            }
+          }],
+          { options: true }
+        );
+      });
+
+      it('throws if aggregateDb fails', async() => {
+        serviceProvider.aggregateDb.throws(new Error('err'));
+        const error: any = await database.sql('SELECT * FROM somecollection;').catch(err => err);
+        expect(error.message).to.be.equal('err');
+      });
+
+      it('throws if connecting to an unsupported server', async() => {
+        const serviceProviderCursor = stubInterface<ServiceProviderAggCursor>();
+        serviceProvider.aggregateDb.returns(serviceProviderCursor as any);
+        serviceProviderCursor.hasNext.throws(Object.assign(new Error(), { code: 40324 }));
+        const error: any = await database.sql('SELECT * FROM somecollection;').catch(err => err);
+        expect(error.message).to.match(/db\.sql currently only works when connected to a Data Lake/);
+      });
+
+      it('forwards other driver errors', async() => {
+        const serviceProviderCursor = stubInterface<ServiceProviderAggCursor>();
+        serviceProvider.aggregateDb.returns(serviceProviderCursor as any);
+        serviceProviderCursor.hasNext.throws(Object.assign(new Error('any error'), { code: 12345 }));
+        const error: any = await database.sql('SELECT * FROM somecollection;').catch(err => err);
+        expect(error.message).to.be.equal('any error');
+      });
+
+      it('forwards generic cursor errors', async() => {
+        const serviceProviderCursor = stubInterface<ServiceProviderAggCursor>();
+        serviceProvider.aggregateDb.returns(serviceProviderCursor as any);
+        serviceProviderCursor.hasNext.throws(Object.assign(new Error('any error')));
+        const error: any = await database.sql('SELECT * FROM somecollection;').catch(err => err);
+        expect(error.message).to.be.equal('any error');
+      });
+    });
   });
   describe('with session', () => {
     let serviceProvider: StubbedInstance<ServiceProvider>;
@@ -2534,7 +2583,8 @@ describe('Database', () => {
       'cloneCollection',
       'copyDatabase',
       'getReplicationInfo',
-      'setSecondaryOk'
+      'setSecondaryOk',
+      'sql'
     ];
     const args = [ {}, {}, {} ];
     beforeEach(() => {
diff --git a/packages/shell-api/src/database.ts b/packages/shell-api/src/database.ts
index eb1bd2bb9..9011cadd4 100644
--- a/packages/shell-api/src/database.ts
+++ b/packages/shell-api/src/database.ts
@@ -1452,4 +1452,38 @@ export default class Database extends ShellApiWithMongoClass {
     this._mongo._instanceState.currentCursor = cursor;
     return cursor;
   }
+
+  @serverVersions(['4.4.0', ServerVersions.latest])
+  @returnsPromise
+  @returnType('AggregationCursor')
+  async sql(sqlString: string, options?: Document): Promise<AggregationCursor> {
+    this._emitDatabaseApiCall('sql', { sqlString: sqlString, options });
+    await this._instanceState.shellApi.print(
+      'Note: this is an experimental feature that may be subject to change in future releases.'
+    );
+
+    const cursor = await this.aggregate([{
+      $sql: {
+        statement: sqlString,
+        format: 'jdbc',
+        dialect: 'mongosql',
+        formatVersion: 1
+      }
+    }], options);
+
+    try {
+      await cursor.hasNext();
+    } catch (err: any) {
+      if (err.code?.valueOf() === 40324) { // unrecognized stage error
+        throw new MongoshRuntimeError(
+          'db.sql currently only works when connected to a Data Lake',
+          CommonErrors.CommandFailed
+        );
+      }
+
+      throw err;
+    }
+
+    return cursor;
+  }
 }