[VisBuilder] Add Capability to generate dynamic vega

src/plugins/advanced_settings/public/management_app/lib/get_category_name.ts

@@ -48,6 +48,9 @@ const names: Record<string, string> = {
   visualizations: i18n.translate('advancedSettings.categoryNames.visualizationsLabel', {
     defaultMessage: 'Visualizations',
   }),
+  visbuilder: i18n.translate('advancedSettings.categoryNames.visbuilderLabel', {
+    defaultMessage: 'VisBuilder',
+  }),
   discover: i18n.translate('advancedSettings.categoryNames.discoverLabel', {
     defaultMessage: 'Discover',
   }),

src/plugins/expressions/public/index.ts

@@ -134,3 +134,4 @@ export {
   UnmappedTypeStrings,
   ExpressionValueRender as Render,
 } from '../common';
+export { getExpressionsService } from './services';

src/plugins/vis_builder/common/constants.ts

@@ -0,0 +1,6 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+export const VISBUILDER_ENABLE_VEGA_SETTING = 'visbuilder:enableVega';

src/plugins/vis_builder/public/application/components/workspace.tsx

@@ -13,6 +13,7 @@ import { validateSchemaState, validateAggregations } from '../utils/validations'
 import { useTypedDispatch, useTypedSelector, setUIStateState } from '../utils/state_management';
 import { useAggs, useVisualizationType } from '../utils/use';
 import { PersistedState } from '../../../../visualizations/public';
+import { VISBUILDER_ENABLE_VEGA_SETTING } from '../../../common/constants';
 
 import hand_field from '../../assets/hand_field.svg';
 import fields_bg from '../../assets/fields_bg.svg';
@@ -27,6 +28,7 @@ export const WorkspaceUI = () => {
       notifications: { toasts },
       data,
       uiActions,
+      uiSettings,
     },
   } = useOpenSearchDashboards<VisBuilderServices>();
   const { toExpression, ui } = useVisualizationType();
@@ -37,6 +39,7 @@ export const WorkspaceUI = () => {
     filters: data.query.filterManager.getFilters(),
     timeRange: data.query.timefilter.timefilter.getTime(),
   });
+  const useVega = uiSettings.get(VISBUILDER_ENABLE_VEGA_SETTING);
   const rootState = useTypedSelector((state) => state);
   const dispatch = useTypedDispatch();
   // Visualizations require the uiState object to persist even when the expression changes
@@ -81,12 +84,20 @@ export const WorkspaceUI = () => {
         return;
       }
 
-      const exp = await toExpression(rootState, searchContext);
+      const exp = await toExpression(rootState, searchContext, useVega);
       setExpression(exp);
     }
 
     loadExpression();
-  }, [rootState, toExpression, toasts, ui.containerConfig.data.schemas, searchContext, aggConfigs]);
+  }, [
+    rootState,
+    toExpression,
+    toasts,
+    ui.containerConfig.data.schemas,
+    searchContext,
+    aggConfigs,
+    useVega,
+  ]);
 
   useLayoutEffect(() => {
     const subscription = data.query.state$.subscribe(({ state }) => {

src/plugins/vis_builder/public/embeddable/vis_builder_embeddable.tsx

@@ -34,11 +34,13 @@ import {
   getIndexPatterns,
   getTypeService,
   getUIActions,
+  getUISettings,
 } from '../plugin_services';
 import { PersistedState, prepareJson } from '../../../visualizations/public';
 import { VisBuilderSavedVis } from '../saved_visualizations/transforms';
 import { handleVisEvent } from '../application/utils/handle_vis_event';
 import { VisBuilderEmbeddableFactoryDeps } from './vis_builder_embeddable_factory';
+import { VISBUILDER_ENABLE_VEGA_SETTING } from '../../common/constants';
 
 // Apparently this needs to match the saved object type for the clone and replace panel actions to work
 export const VISBUILDER_EMBEDDABLE = VISBUILDER_SAVED_OBJECT;
@@ -150,11 +152,16 @@ export class VisBuilderEmbeddable extends Embeddable<VisBuilderInput, VisBuilder
 
       if (!valid && errorMsg) throw new Error(errorMsg);
 
-      const exp = await toExpression(renderState, {
-        filters: this.filters,
-        query: this.query,
-        timeRange: this.timeRange,
-      });
+      const useVega = getUISettings().get(VISBUILDER_ENABLE_VEGA_SETTING);
+      const exp = await toExpression(
+        renderState,
+        {
+          filters: this.filters,
+          query: this.query,
+          timeRange: this.timeRange,
+        },
+        useVega
+      );
       return exp;
     } catch (error) {
       this.onContainerError(error as Error);

src/plugins/vis_builder/public/plugin.test.ts

@@ -9,6 +9,7 @@ import { dataPluginMock } from '../../data/public/mocks';
 import { embeddablePluginMock } from '../../embeddable/public/mocks';
 import { navigationPluginMock } from '../../navigation/public/mocks';
 import { visualizationsPluginMock } from '../../visualizations/public/mocks';
+import { expressionsPluginMock } from '../../expressions/public/mocks';
 import { PLUGIN_ID, PLUGIN_NAME } from '../common';
 import { VisBuilderPlugin } from './plugin';
 
@@ -29,6 +30,7 @@ describe('VisBuilderPlugin', () => {
         visualizations: visualizationsPluginMock.createSetupContract(),
         embeddable: embeddablePluginMock.createSetupContract(),
         data: dataPluginMock.createSetupContract(),
+        expressions: expressionsPluginMock.createSetupContract(), // Add this line
       };
 
       const setup = plugin.setup(coreSetup, setupDeps);
@@ -41,6 +43,7 @@ describe('VisBuilderPlugin', () => {
           aliasApp: PLUGIN_ID,
         })
       );
+      expect(setupDeps.expressions.registerFunction).toHaveBeenCalled(); // Add this expectation
     });
   });
 });

src/plugins/vis_builder/public/plugin.ts

@@ -56,6 +56,7 @@ import {
   withNotifyOnErrors,
 } from '../../opensearch_dashboards_utils/public';
 import { opensearchFilters } from '../../data/public';
+import { createRawDataVisFn } from './visualizations/vega/utils/expression_helper';
 
 export class VisBuilderPlugin
   implements
@@ -74,7 +75,7 @@ export class VisBuilderPlugin
 
   public setup(
     core: CoreSetup<VisBuilderPluginStartDependencies, VisBuilderStart>,
-    { embeddable, visualizations, data }: VisBuilderPluginSetupDependencies
+    { embeddable, visualizations, data, expressions: exp }: VisBuilderPluginSetupDependencies
   ) {
     const { appMounted, appUnMounted, stop: stopUrlTracker } = createOsdUrlTracker({
       baseUrl: core.http.basePath.prepend(`/app/${PLUGIN_ID}`),
@@ -107,6 +108,7 @@ export class VisBuilderPlugin
     // Register Default Visualizations
     const typeService = this.typeService;
     registerDefaultTypes(typeService.setup());
+    exp.registerFunction(createRawDataVisFn());
 
     // Register the plugin to core
     core.application.register({

src/plugins/vis_builder/public/services/type_service/types.ts

@@ -31,6 +31,7 @@ export interface VisualizationTypeOptions<T = any> {
   };
   readonly toExpression: (
     state: RenderState,
-    searchContext: IExpressionLoaderParams['searchContext']
+    searchContext: IExpressionLoaderParams['searchContext'],
+    useVega: boolean
   ) => Promise<string | undefined>;
 }

src/plugins/vis_builder/public/services/type_service/visualization_type.tsx

@@ -18,7 +18,8 @@ export class VisualizationType implements IVisualizationType {
   public readonly ui: IVisualizationType['ui'];
   public readonly toExpression: (
     state: RenderState,
-    searchContext: IExpressionLoaderParams['searchContext']
+    searchContext: IExpressionLoaderParams['searchContext'],
+    useVega: boolean
   ) => Promise<string | undefined>;
 
   constructor(options: VisualizationTypeOptions) {

src/plugins/vis_builder/public/types.ts

@@ -8,7 +8,7 @@ import { SavedObject, SavedObjectsStart } from '../../saved_objects/public';
 import { EmbeddableSetup, EmbeddableStart } from '../../embeddable/public';
 import { DashboardStart } from '../../dashboard/public';
 import { VisualizationsSetup } from '../../visualizations/public';
-import { ExpressionsStart } from '../../expressions/public';
+import { ExpressionsStart, ExpressionsPublicPlugin } from '../../expressions/public';
 import { NavigationPublicPluginStart } from '../../navigation/public';
 import { DataPublicPluginStart } from '../../data/public';
 import { TypeServiceSetup, TypeServiceStart } from './services/type_service';
@@ -18,6 +18,7 @@ import { IOsdUrlStateStorage } from '../../opensearch_dashboards_utils/public';
 import { DataPublicPluginSetup } from '../../data/public';
 import { UiActionsStart } from '../../ui_actions/public';
 import { Capabilities } from '../../../core/public';
+import { IUiSettingsClient } from '../../../core/public';
 
 export type VisBuilderSetup = TypeServiceSetup;
 export interface VisBuilderStart extends TypeServiceStart {
@@ -28,6 +29,7 @@ export interface VisBuilderPluginSetupDependencies {
   embeddable: EmbeddableSetup;
   visualizations: VisualizationsSetup;
   data: DataPublicPluginSetup;
+  expressions: ReturnType<ExpressionsPublicPlugin['setup']>;
 }
 export interface VisBuilderPluginStartDependencies {
   embeddable: EmbeddableStart;
@@ -37,6 +39,7 @@ export interface VisBuilderPluginStartDependencies {
   dashboard: DashboardStart;
   expressions: ExpressionsStart;
   uiActions: UiActionsStart;
+  uiSettings: IUiSettingsClient;
 }
 
 export interface VisBuilderServices extends CoreStart {

src/plugins/vis_builder/public/visualizations/common/expression_helpers.ts

@@ -7,20 +7,21 @@ import { cloneDeep } from 'lodash';
 import { OpenSearchaggsExpressionFunctionDefinition } from '../../../../data/public';
 import { ExpressionFunctionOpenSearchDashboards } from '../../../../expressions';
 import { buildExpressionFunction } from '../../../../expressions/public';
-import { VisualizationState } from '../../application/utils/state_management';
+import { VisualizationState, StyleState } from '../../application/utils/state_management';
 import { getSearchService, getIndexPatterns } from '../../plugin_services';
-import { StyleState } from '../../application/utils/state_management';
+import { IExpressionLoaderParams } from '../../../../expressions/public';
 
 export const getAggExpressionFunctions = async (
   visualization: VisualizationState,
-  style?: StyleState
+  style?: StyleState,
+  useVega: boolean = false,
+  searchContext?: IExpressionLoaderParams['searchContext']
 ) => {
   const { activeVisualization, indexPattern: indexId = '' } = visualization;
   const { aggConfigParams } = activeVisualization || {};
 
   const indexPatternsService = getIndexPatterns();
   const indexPattern = await indexPatternsService.get(indexId);
-  // aggConfigParams is the serealizeable aggConfigs that need to be reconstructed here using the agg servce
   const aggConfigs = getSearchService().aggs.createAggConfigs(
     indexPattern,
     cloneDeep(aggConfigParams)
@@ -31,7 +32,6 @@ export const getAggExpressionFunctions = async (
     {}
   );
 
-  // soon this becomes: const opensearchaggs = vis.data.aggs!.toExpressionAst();
   const opensearchaggs = buildExpressionFunction<OpenSearchaggsExpressionFunctionDefinition>(
     'opensearchaggs',
     {
@@ -43,9 +43,20 @@ export const getAggExpressionFunctions = async (
     }
   );
 
+  let expressionFns = [opensearchDashboards, opensearchaggs];
+
+  if (useVega === true && searchContext) {
+    const opensearchDashboardsContext = buildExpressionFunction('opensearch_dashboards_context', {
+      timeRange: JSON.stringify(searchContext.timeRange || {}),
+      filters: JSON.stringify(searchContext.filters || []),
+      query: JSON.stringify(searchContext.query || []),
+    });
+    expressionFns = [opensearchDashboards, opensearchDashboardsContext, opensearchaggs];
+  }
+
   return {
     aggConfigs,
     indexPattern,
-    expressionFns: [opensearchDashboards, opensearchaggs],
+    expressionFns,
   };
 };

src/plugins/vis_builder/public/visualizations/vega/components/encoding.test.ts

@@ -0,0 +1,31 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { buildVegaLiteEncoding } from './encoding';
+
+describe('encoding.ts', () => {
+  describe('buildVegaLiteEncoding', () => {
+    it('should build correct encoding for x and y axes', () => {
+      const dimensions = { x: [{}], y: [{}] };
+      const formats = { xAxisLabel: 'X Label', yAxisLabel: 'Y Label' };
+      const result = buildVegaLiteEncoding(dimensions, formats);
+
+      expect(result.x).toBeDefined();
+      expect(result.y).toBeDefined();
+      expect(result.x!.axis!.title).toBe('X Label');
+      expect(result.y!.axis!.title).toBe('Y Label');
+    });
+
+    it('should include color encoding when y dimension is present', () => {
+      const dimensions = { x: [{}], y: [{}] };
+      const formats = {};
+      const result = buildVegaLiteEncoding(dimensions, formats);
+
+      expect(result.color).toBeDefined();
+      expect(result.color!.field).toBe('series');
+      expect(result.color!.type).toBe('nominal');
+    });
+  });
+});