feat: allow text and value query variables

fr-ser · Mar 8, 2024 · 3e83eea · 3e83eea
1 parent 074ebbf
commit 3e83eea
Show file tree

Hide file tree

Showing 2 changed files with 63 additions and 5 deletions.
diff --git a/src/DataSource.test.ts b/src/DataSource.test.ts
@@ -31,7 +31,7 @@ describe('DataSource', () => {
   });
 
   describe('query variables', () => {
-    it('returns the result of a query', async () => {
+    it('returns the result of a query for a single column', async () => {
       const ds = new DataSource({} as any);
       const mockResponse = {
         data: [
@@ -47,6 +47,51 @@ describe('DataSource', () => {
       expect(result).toStrictEqual([{ text: 1 }, { text: 2 }]);
     });
 
+    it('returns the result of a query for a text and values', async () => {
+      const ds = new DataSource({} as any);
+      const mockResponse = {
+        data: [
+          new MutableDataFrame({
+            fields: [
+              { name: '__value', type: FieldType.number, values: [1, 2] },
+              { name: '__text', type: FieldType.number, values: ['a', 'b'] },
+            ],
+          }),
+        ],
+      };
+      ds.query = jest.fn(() => ({ toPromise: async () => mockResponse })) as any;
+
+      const result = await ds.metricFindQuery("SELECT 'my-query'", { variable: { datasource: 'sqlite' } });
+
+      expect(result).toStrictEqual([
+        { value: 1, text: 'a' },
+        { value: 2, text: 'b' },
+      ]);
+    });
+
+    it('throws for 2 columns if __text or __value is missing', async () => {
+      const ds = new DataSource({} as any);
+      const mockResponse = {
+        data: [
+          new MutableDataFrame({
+            fields: [
+              { name: 'value', type: FieldType.number, values: [1, 2] },
+              { name: 'text', type: FieldType.number, values: ['a', 'b'] },
+            ],
+          }),
+        ],
+      };
+      ds.query = jest.fn(() => ({ toPromise: async () => mockResponse })) as any;
+
+      try {
+        await ds.metricFindQuery("SELECT 'my-query'", { variable: { datasource: 'sqlite' } });
+        fail('did not receive an error');
+      } catch (error) {
+        const errorMessage = (error as Error).toString();
+        expect(errorMessage).toContain('No columns named "__text" and "__value" were found');
+      }
+    });
+
     it('throws for multiple columns', async () => {
       const ds = new DataSource({} as any);
       const mockResponse = {
@@ -55,6 +100,7 @@ describe('DataSource', () => {
             fields: [
               { name: 'value', type: FieldType.number, values: [1, 2] },
               { name: 'name', type: FieldType.number, values: ['a', 'b'] },
+              { name: 'label', type: FieldType.number, values: ['c', 'd'] },
             ],
           }),
         ],
@@ -66,7 +112,7 @@ describe('DataSource', () => {
         fail('did not receive an error');
       } catch (error) {
         const errorMessage = (error as Error).toString();
-        expect(errorMessage).toContain('Received more than one (2) fields');
+        expect(errorMessage).toContain('Received more than two (3) fields');
       }
     });
 

diff --git a/src/DataSource.ts b/src/DataSource.ts
@@ -39,11 +39,23 @@ export class DataSource extends DataSourceWithBackend<SQLiteQuery, MyDataSourceO
     }
 
     const data = response.data[0] as DataFrame;
-    if (data.fields.length !== 1) {
+    if (data.fields.length === 1) {
+      return data.fields[0].values.toArray().map((text) => ({ text }));
+    } else if (data.fields.length === 2) {
+      const textIndex = data.fields.findIndex((x) => x.name === '__text');
+      const valueIndex = data.fields.findIndex((x) => x.name === '__value');
+      if (textIndex === -1 || valueIndex === -1) {
+        throw new Error(
+          `No columns named "__text" and "__value" were found. Columns: ${data.fields.map((x) => x.name).join(',')}`
+        );
+      }
+
+      const valueArray = data.fields[valueIndex].values.toArray();
+      return data.fields[textIndex].values.toArray().map((text, index) => ({ text, value: valueArray[index] }));
+    } else {
       throw new Error(
-        `Received more than one (${data.fields.length}) fields: ${data.fields.map((x) => x.name).join(',')}`
+        `Received more than two (${data.fields.length}) fields: ${data.fields.map((x) => x.name).join(',')}`
       );
     }
-    return data.fields[0].values.toArray().map((text) => ({ text }));
   }
 }