feat(asset-calibration): workspace grouping

ni · Oct 4, 2024 · 0663c99 · 0663c99
1 parent 7acdca2
commit 0663c99
Show file tree

Hide file tree

Showing 5 changed files with 169 additions and 13 deletions.
diff --git a/src/datasources/asset-calibration/AssetCalibrationDataSource.test.ts b/src/datasources/asset-calibration/AssetCalibrationDataSource.test.ts
@@ -120,6 +120,28 @@ const monthLocationGroupCalibrationForecastResponseMock: CalibrationForecastResp
   }
 }
 
+const modelWorkspaceGroupCalibrationForecastResponseMock: CalibrationForecastResponse =
+{
+  calibrationForecast: {
+    columns: [
+      { name: "", values: [1], columnDescriptors: [{ value: "Model1", type: ColumnDescriptorType.StringValue }, { value: "Workspace1", type: ColumnDescriptorType.WorkspaceId }] },
+      { name: "", values: [2], columnDescriptors: [{ value: "Model2", type: ColumnDescriptorType.StringValue }, { value: "Workspace1", type: ColumnDescriptorType.WorkspaceId }] }
+    ]
+  }
+}
+
+const monthWorkspaceGroupCalibrationForecastResponseMock: CalibrationForecastResponse =
+{
+  calibrationForecast: {
+    columns: [
+      { name: "", values: ["2022-01-01T00:00:00.0000000Z", "2022-02-01T00:00:00.0000000Z", "2022-03-01T00:00:00.0000000Z"], columnDescriptors: [{ value: "Month", type: ColumnDescriptorType.Time }] },
+      { name: "", values: [1, 2, 3], columnDescriptors: [{ value: "Workspace1", type: ColumnDescriptorType.WorkspaceId }] },
+      { name: "", values: [2, 4, 1], columnDescriptors: [{ value: "Workspace2", type: ColumnDescriptorType.WorkspaceId }] },
+      { name: "", values: [4, 3, 1], columnDescriptors: [{ value: "Workspace3", type: ColumnDescriptorType.WorkspaceId }] }
+    ]
+  }
+}
+
 const monthBasedCalibrationForecastQueryMock: AssetCalibrationQuery = {
   refId: '',
   groupBy: [AssetCalibrationTimeBasedGroupByType.Month],
@@ -155,6 +177,16 @@ const monthLocationBasedCalibrationForecastQueryMock: AssetCalibrationQuery = {
   groupBy: [AssetCalibrationTimeBasedGroupByType.Month, AssetCalibrationPropertyGroupByType.Location],
 }
 
+const modelWorkspaceBasedCalibrationForecastQueryMock: AssetCalibrationQuery = {
+  refId: '',
+  groupBy: [AssetCalibrationPropertyGroupByType.Model, AssetCalibrationPropertyGroupByType.Workspace],
+}
+
+const monthWorkspaceBasedCalibrationForecastQueryMock: AssetCalibrationQuery = {
+  refId: '',
+  groupBy: [AssetCalibrationTimeBasedGroupByType.Month, AssetCalibrationPropertyGroupByType.Workspace],
+}
+
 const buildCalibrationForecastQuery = getQueryBuilder<AssetCalibrationQuery>()({
   refId: '',
   groupBy: []
@@ -253,7 +285,7 @@ describe('queries', () => {
 
     const result = await datastore.query(buildCalibrationForecastQuery(modelBasedCalibrationForecastQueryMock))
 
-    expect(result.data).toMatchSnapshot()
+    expect(result.data).toMatchSnapshot();
   })
 
   test('calibration forecast with model and location groupBy', async () => {
@@ -263,7 +295,7 @@ describe('queries', () => {
 
     const result = await datastore.query(buildCalibrationForecastQuery(modelLocationBasedCalibrationForecastQueryMock))
 
-    expect(result.data).toMatchSnapshot()
+    expect(result.data).toMatchSnapshot();
   })
 
   test('calibration forecast with month and location groupBy', async () => {
@@ -273,6 +305,26 @@ describe('queries', () => {
 
     const result = await datastore.query(buildCalibrationForecastQuery(monthLocationBasedCalibrationForecastQueryMock))
 
+    expect(result.data).toMatchSnapshot();
+  })
+
+  test('calibration forecast with model and workspace groupBy', async () => {
+    backendServer.fetch
+      .calledWith(requestMatching({ url: '/niapm/v1/assets/calibration-forecast' }))
+      .mockReturnValue(createFetchResponse(modelWorkspaceGroupCalibrationForecastResponseMock as CalibrationForecastResponse))
+
+    const result = await datastore.query(buildCalibrationForecastQuery(modelWorkspaceBasedCalibrationForecastQueryMock))
+
+    expect(result.data).toMatchSnapshot();
+  })
+
+  test('calibration forecast with month and workspace groupBy', async () => {
+    backendServer.fetch
+      .calledWith(requestMatching({ url: '/niapm/v1/assets/calibration-forecast' }))
+      .mockReturnValue(createFetchResponse(monthWorkspaceGroupCalibrationForecastResponseMock as CalibrationForecastResponse))
+
+    const result = await datastore.query(buildCalibrationForecastQuery(monthWorkspaceBasedCalibrationForecastQueryMock))
+
     expect(result.data).toMatchSnapshot()
   })
 

diff --git a/src/datasources/asset-calibration/AssetCalibrationDataSource.ts b/src/datasources/asset-calibration/AssetCalibrationDataSource.ts
@@ -22,6 +22,7 @@ import TTLCache from '@isaacs/ttlcache';
 import { metadataCacheTTL } from 'datasources/data-frame/constants';
 import { SystemMetadata } from 'datasources/system/types';
 import { defaultOrderBy, defaultProjection } from 'datasources/system/constants';
+import { Workspace } from 'core/types';
 
 export class AssetCalibrationDataSource extends DataSourceBase<AssetCalibrationQuery> {
   public defaultQuery = {
@@ -40,9 +41,11 @@ export class AssetCalibrationDataSource extends DataSourceBase<AssetCalibrationQ
   baseUrl = this.instanceSettings.url + '/niapm/v1';
 
   systemAliasCache: TTLCache<string, SystemMetadata> = new TTLCache<string, SystemMetadata>({ ttl: metadataCacheTTL });
+  workspacesCache: TTLCache<string, Workspace> = new TTLCache<string, Workspace>({ ttl: metadataCacheTTL });
 
   async runQuery(query: AssetCalibrationQuery, options: DataQueryRequest): Promise<DataFrameDTO> {
     await this.loadSystems();
+    await this.loadWorkspaces();
 
     if (query.filter) {
       query.filter = this.templateSrv.replace(transformComputedFieldsQuery(query.filter, AssetComputedDataFields), options.scopedVars);
@@ -113,7 +116,13 @@ export class AssetCalibrationDataSource extends DataSourceBase<AssetCalibrationQ
       if (descriptor.type === ColumnDescriptorType.MinionId && this.systemAliasCache) {
         const system = this.systemAliasCache.get(descriptor.value);
 
-        return system?.alias || descriptor.value
+        return system?.alias || descriptor.value;
+      }
+
+      if (descriptor.type === ColumnDescriptorType.WorkspaceId && this.workspacesCache) {
+        const workspace = this.workspacesCache.get(descriptor.value);
+
+        return workspace?.name || descriptor.value
       }
       return descriptor.value
     }).join(' - ');
@@ -181,6 +190,15 @@ export class AssetCalibrationDataSource extends DataSourceBase<AssetCalibrationQ
     systems.forEach(system => this.systemAliasCache.set(system.id, system));
   }
 
+  private async loadWorkspaces(): Promise<void> {
+    if (this.workspacesCache.size > 0) {
+      return;
+    }
+
+    const workspaces = await this.getWorkspaces();
+    workspaces.forEach(workspace => this.workspacesCache.set(workspace.id, workspace));
+  }
+
   async testDatasource(): Promise<TestDataSourceResponse> {
     await this.get(this.baseUrl + '/assets?take=1');
     return { status: 'success', message: 'Data source connected and authentication successful!' };

diff --git a/src/datasources/asset-calibration/__snapshots__/AssetCalibrationDataSource.test.ts.snap b/src/datasources/asset-calibration/__snapshots__/AssetCalibrationDataSource.test.ts.snap
@@ -190,6 +190,30 @@ exports[`queries calibration forecast with model and location groupBy 1`] = `
 ]
 `;
 
+exports[`queries calibration forecast with model and workspace groupBy 1`] = `
+[
+  {
+    "fields": [
+      {
+        "name": "Group",
+        "values": [
+          "Model1 - Workspace1",
+          "Model2 - Workspace1",
+        ],
+      },
+      {
+        "name": "Assets",
+        "values": [
+          1,
+          2,
+        ],
+      },
+    ],
+    "refId": "A",
+  },
+]
+`;
+
 exports[`queries calibration forecast with model groupBy 1`] = `
 [
   {
@@ -280,6 +304,72 @@ exports[`queries calibration forecast with month and location groupBy 1`] = `
 ]
 `;
 
+exports[`queries calibration forecast with month and workspace groupBy 1`] = `
+[
+  {
+    "fields": [
+      {
+        "columnDescriptors": [
+          {
+            "type": "TIME",
+            "value": "Month",
+          },
+        ],
+        "name": "Month",
+        "values": [
+          "January 2022",
+          "February 2022",
+          "March 2022",
+        ],
+      },
+      {
+        "columnDescriptors": [
+          {
+            "type": "WORKSPACE_ID",
+            "value": "Workspace1",
+          },
+        ],
+        "name": "Workspace1",
+        "values": [
+          1,
+          2,
+          3,
+        ],
+      },
+      {
+        "columnDescriptors": [
+          {
+            "type": "WORKSPACE_ID",
+            "value": "Workspace2",
+          },
+        ],
+        "name": "Workspace2",
+        "values": [
+          2,
+          4,
+          1,
+        ],
+      },
+      {
+        "columnDescriptors": [
+          {
+            "type": "WORKSPACE_ID",
+            "value": "Workspace3",
+          },
+        ],
+        "name": "Workspace3",
+        "values": [
+          4,
+          3,
+          1,
+        ],
+      },
+    ],
+    "refId": "A",
+  },
+]
+`;
+
 exports[`queries calibration forecast with month groupBy returns empty results 1`] = `
 [
   {

diff --git a/src/datasources/asset-calibration/components/AssetCalibrationQueryEditor.tsx b/src/datasources/asset-calibration/components/AssetCalibrationQueryEditor.tsx
@@ -7,24 +7,20 @@ import { enumToOptions } from '../../../core/utils';
 import _ from 'lodash';
 import { AssetCalibrationQueryBuilder } from './AssetCalibrationQueryBuilder';
 import { Workspace } from 'core/types';
-import { FloatingError, parseErrorMessage } from 'core/errors';
 
 type Props = QueryEditorProps<AssetCalibrationDataSource, AssetCalibrationQuery>;
 
 export function AssetCalibrationQueryEditor({ query, onChange, onRunQuery, datasource }: Props) {
   query = datasource.prepareQuery(query) as AssetCalibrationQuery;
 
+  const areWorkspacesLoaded = datasource.workspacesCache.size > 0;
   const [workspaces, setWorkspaces] = useState<Workspace[]>([]);
-  const [error, setError] = useState<string>('');
 
   useEffect(() => {
-    const getWorkspaces = async () => {
-      const workspaces = await datasource.getWorkspaces();
-      setWorkspaces(workspaces);
+    if (areWorkspacesLoaded) {
+      setWorkspaces(Array.from(datasource.workspacesCache.values()));
     }
-
-    getWorkspaces().catch(error => setError(parseErrorMessage(error) || 'Failed to fetch workspaces'));
-  }, [datasource]);
+  }, [datasource, areWorkspacesLoaded]);
 
   const handleQueryChange = (value: AssetCalibrationQuery, runQuery: boolean): void => {
     onChange(value);
@@ -84,8 +80,6 @@ export function AssetCalibrationQueryEditor({ query, onChange, onRunQuery, datas
         workspaces={workspaces}
         onChange={(event: any) => onParameterChange(event)}>
       </AssetCalibrationQueryBuilder>
-
-      <FloatingError message={error} />
     </div>
   );
 }

diff --git a/src/datasources/asset-calibration/types.ts b/src/datasources/asset-calibration/types.ts
@@ -15,6 +15,7 @@ export enum AssetQueryLabel {
 export enum AssetCalibrationPropertyGroupByType {
   Location = "LOCATION",
   Model = "MODEL",
+  Workspace = "WORKSPACE",
 }
 
 export enum AssetCalibrationTimeBasedGroupByType {
@@ -71,6 +72,7 @@ export enum ColumnDescriptorType {
   Count = "COUNT",
   StringValue = "STRING_VALUE",
   MinionId = "MINION_ID",
+  WorkspaceId = "WORKSPACE_ID",
 }
 
 export interface QBField extends QueryBuilderField {