From d0a165eaad9e7271da6fedb494286091808f6c28 Mon Sep 17 00:00:00 2001
From: Kawika Avilla <kavilla414@gmail.com>
Date: Thu, 5 Sep 2024 19:34:08 +0000
Subject: [PATCH] [enhance] access datasource information with dot notation to
 fix MQL

Deconstructing the dataSource caused an exception when no dataSource was set
with the query.

Accessing by dot notation with null operator fixes this issue.

Adding tests to prevent this happening again

Signed-off-by: Kawika Avilla <kavilla414@gmail.com>
---
 .../server/utils/facet.test.ts                | 119 ++++++++++++++++++
 .../query_enhancements/server/utils/facet.ts  |   4 +-
 2 files changed, 121 insertions(+), 2 deletions(-)
 create mode 100644 src/plugins/query_enhancements/server/utils/facet.test.ts

diff --git a/src/plugins/query_enhancements/server/utils/facet.test.ts b/src/plugins/query_enhancements/server/utils/facet.test.ts
new file mode 100644
index 00000000000..20ae78612c1
--- /dev/null
+++ b/src/plugins/query_enhancements/server/utils/facet.test.ts
@@ -0,0 +1,119 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { Logger } from 'opensearch-dashboards/server';
+import { Facet, FacetProps } from './facet';
+
+describe('Facet', () => {
+  let facet: Facet;
+  let mockClient: jest.Mock;
+  let mockLogger: jest.Mocked<Logger>;
+  let mockContext: any;
+  let mockRequest: any;
+
+  beforeEach(() => {
+    mockClient = jest.fn();
+    mockLogger = ({
+      error: jest.fn(),
+    } as unknown) as jest.Mocked<Logger>;
+
+    const props: FacetProps = {
+      client: { asScoped: jest.fn().mockReturnValue({ callAsCurrentUser: mockClient }) },
+      logger: mockLogger,
+      endpoint: 'test-endpoint',
+    };
+
+    facet = new Facet(props);
+
+    mockContext = {
+      dataSource: {
+        opensearch: {
+          legacy: {
+            getClient: jest.fn().mockReturnValue({ callAPI: mockClient }),
+          },
+        },
+      },
+    };
+
+    mockRequest = {
+      body: {
+        query: {
+          query: 'test query',
+          dataset: {
+            dataSource: {
+              id: 'test-id',
+              meta: {
+                name: 'test-name',
+                sessionId: 'test-session',
+              },
+            },
+          },
+        },
+        format: 'jdbc',
+        lang: 'sql',
+      },
+    };
+  });
+
+  describe('describeQuery', () => {
+    it('should handle request with complete dataset information', async () => {
+      mockClient.mockResolvedValue({ result: 'success' });
+
+      const result = await facet.describeQuery(mockContext, mockRequest);
+
+      expect(result).toEqual({ success: true, data: { result: 'success' } });
+      expect(mockClient).toHaveBeenCalledWith('test-endpoint', {
+        body: {
+          query: 'test query',
+          datasource: 'test-name',
+          sessionId: 'test-session',
+          lang: 'sql',
+        },
+      });
+    });
+
+    it('should handle request with missing dataSource', async () => {
+      mockRequest.body.query.dataset.dataSource = undefined;
+      mockClient.mockResolvedValue({ result: 'success' });
+
+      const result = await facet.describeQuery(mockContext, mockRequest);
+
+      expect(result).toEqual({ success: true, data: { result: 'success' } });
+      expect(mockClient).toHaveBeenCalledWith('test-endpoint', {
+        body: {
+          query: 'test query',
+          lang: 'sql',
+        },
+      });
+    });
+
+    it('should handle request with missing dataset', async () => {
+      mockRequest.body.query.dataset = undefined;
+      mockClient.mockResolvedValue({ result: 'success' });
+
+      const result = await facet.describeQuery(mockContext, mockRequest);
+
+      expect(result).toEqual({ success: true, data: { result: 'success' } });
+      expect(mockClient).toHaveBeenCalledWith('test-endpoint', {
+        body: {
+          query: 'test query',
+          lang: 'sql',
+        },
+      });
+    });
+
+    it('should handle errors', async () => {
+      const error = new Error('Test error');
+      mockClient.mockRejectedValue(error);
+
+      const result = await facet.describeQuery(mockContext, mockRequest);
+
+      expect(result).toEqual({ success: false, data: error });
+      expect(mockLogger.error).toHaveBeenCalledWith(
+        'Facet fetch: test-endpoint: Error: Test error'
+      );
+    });
+  });
+});
diff --git a/src/plugins/query_enhancements/server/utils/facet.ts b/src/plugins/query_enhancements/server/utils/facet.ts
index f86a07b5432..0b6dd52407c 100644
--- a/src/plugins/query_enhancements/server/utils/facet.ts
+++ b/src/plugins/query_enhancements/server/utils/facet.ts
@@ -38,8 +38,8 @@ export class Facet {
   ): Promise<FacetResponse> => {
     try {
       const query: Query = request.body.query;
-      const { dataSource } = query.dataset!;
-      const { meta } = dataSource!;
+      const dataSource = query.dataset?.dataSource;
+      const meta = dataSource?.meta;
       const { format, lang } = request.body;
       const params = {
         body: {