diff --git a/src/plugins/data/public/query/query_service.ts b/src/plugins/data/public/query/query_service.ts
index 80f6bd542e5d..0f902aa4bcde 100644
--- a/src/plugins/data/public/query/query_service.ts
+++ b/src/plugins/data/public/query/query_service.ts
@@ -36,7 +36,6 @@ import { TimefilterService, TimefilterSetup } from './timefilter';
 import { createSavedQueryService } from './saved_query/saved_query_service';
 import { createQueryStateObservable } from './state_sync/create_global_query_observable';
 import { QueryStringManager, QueryStringContract } from './query_string';
-import { DataSetContract, DataSetManager } from './dataset_manager';
 import {
   buildOpenSearchQuery,
   DataStorage,
diff --git a/src/plugins/data/public/query/query_string/query_string_manager.ts b/src/plugins/data/public/query/query_string/query_string_manager.ts
index fd0fe063b361..bde708cbf698 100644
--- a/src/plugins/data/public/query/query_string/query_string_manager.ts
+++ b/src/plugins/data/public/query/query_string/query_string_manager.ts
@@ -31,7 +31,7 @@
 import { BehaviorSubject } from 'rxjs';
 import { skip } from 'rxjs/operators';
 import { CoreStart } from 'opensearch-dashboards/public';
-import { DataStorage, Query, SimpleDataSet, TimeRange, UI_SETTINGS } from '../../../common';
+import { Dataset, DataStorage, Query, TimeRange, UI_SETTINGS } from '../../../common';
 import { createHistory, QueryHistory } from './query_history';
 import { DatasetContract, DatasetManager } from './dataset_manager';
 
@@ -107,7 +107,7 @@ export class QueryStringManager {
   };
 
   // Todo: update this function to use the Query object when it is udpated, Query object should include time range and dataset
-  public addToQueryHistory(dataSet: SimpleDataSet, query: Query, timeRange?: TimeRange) {
+  public addToQueryHistory(dataSet: Dataset, query: Query, timeRange?: TimeRange) {
     if (query.query) {
       this.queryHistory.addQueryToHistory(dataSet, query, timeRange);
     }
diff --git a/src/plugins/query_enhancements/public/search/sql_search_interceptor.ts b/src/plugins/query_enhancements/public/search/sql_search_interceptor.ts
index 9376d4f3bd7c..a73a1cee8206 100644
--- a/src/plugins/query_enhancements/public/search/sql_search_interceptor.ts
+++ b/src/plugins/query_enhancements/public/search/sql_search_interceptor.ts
@@ -7,12 +7,7 @@ import { trimEnd } from 'lodash';
 import { Observable, throwError } from 'rxjs';
 import { i18n } from '@osd/i18n';
 import { concatMap, map } from 'rxjs/operators';
-import {
-  DATA_FRAME_TYPES,
-  getRawDataFrame,
-  getRawQueryString,
-  SIMPLE_DATA_SET_TYPES,
-} from '../../../data/common';
+import { DATA_FRAME_TYPES, getRawDataFrame, getRawQueryString } from '../../../data/common';
 import {
   DataPublicPluginStart,
   IOpenSearchDashboardsSearchRequest,
@@ -66,10 +61,10 @@ export class SQLSearchInterceptor extends SearchInterceptor {
       ...dataFrame.meta,
       queryConfig: {
         ...dataFrame.meta.queryConfig,
-        ...(this.queryService.dataSetManager.getDataSet() && {
-          dataSourceId: this.queryService.dataSetManager.getDataSet()?.dataSourceRef?.id,
-          dataSourceName: this.queryService.dataSetManager.getDataSet()?.dataSourceRef?.name,
-          timeFieldName: this.queryService.dataSetManager.getDataSet()?.timeFieldName,
+        ...(this.queryService.dataSetManager.getDataset() && {
+          dataSourceId: this.queryService.dataSetManager.getDataset()?.dataSource?.id,
+          dataSourceName: this.queryService.dataSetManager.getDataset()?.dataSource?.title,
+          timeFieldName: this.queryService.dataSetManager.getDataset()?.timeFieldName,
         }),
       },
     };
@@ -109,10 +104,10 @@ export class SQLSearchInterceptor extends SearchInterceptor {
     }
 
     const queryString = getRawQueryString(searchRequest) ?? '';
-    const dataSourceRef = this.queryService.dataSetManager.getDataSet()
+    const dataSourceRef = this.queryService.dataSetManager.getDataset()
       ? {
-          dataSourceId: this.queryService.dataSetManager.getDataSet()?.dataSourceRef?.id,
-          dataSourceName: this.queryService.dataSetManager.getDataSet()?.dataSourceRef?.name,
+          dataSourceId: this.queryService.dataSetManager.getDataset()?.dataSource?.id,
+          dataSourceName: this.queryService.dataSetManager.getDataset()?.dataSource?.title,
         }
       : {};
 
@@ -187,8 +182,8 @@ export class SQLSearchInterceptor extends SearchInterceptor {
   }
 
   public search(request: IOpenSearchDashboardsSearchRequest, options: ISearchOptions) {
-    const dataSet = this.queryService.dataSetManager.getDataSet();
-    if (dataSet?.type === SIMPLE_DATA_SET_TYPES.TEMPORARY_ASYNC) {
+    const dataSet = this.queryService.dataSetManager.getDataset();
+    if (dataSet?.type === 'TEMPORARY_ASYNC') {
       return this.runSearchAsync(request, options.abortSignal, SEARCH_STRATEGY.SQL_ASYNC);
     }
     return this.runSearch(request, options.abortSignal, SEARCH_STRATEGY.SQL);