@finos/legend-application-query 13.8.0 → 13.8.2

package/src/components/data-space/LegendQueryDataSpaceQueryBuilder.tsx

@@ -0,0 +1,377 @@
+/**
+ * Copyright (c) 2020-present, Goldman Sachs
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import { useApplicationStore } from '@finos/legend-application';
+import {
+  AnchorLinkIcon,
+  CheckIcon,
+  compareLabelFn,
+  ControlledDropdownMenu,
+  createFilter,
+  CustomSelectorInput,
+  MenuContent,
+  MenuContentItem,
+  MenuContentItemIcon,
+  MenuContentItemLabel,
+  MoreVerticalIcon,
+  PanelHeader,
+  PanelHeaderActionItem,
+  PanelHeaderActions,
+  SearchIcon,
+} from '@finos/legend-art';
+import {
+  ResolvedDataSpaceEntityWithOrigin,
+  resolveUsableDataSpaceClasses,
+} from '@finos/legend-extension-dsl-data-space/application';
+import {
+  buildExecutionContextOption,
+  DataSpaceAdvancedSearchModal,
+  type ExecutionContextOption,
+} from '@finos/legend-extension-dsl-data-space/application-query';
+import {
+  getMappingCompatibleRuntimes,
+  PackageableElementExplicitReference,
+  RuntimePointer,
+  type PackageableRuntime,
+  type Runtime,
+} from '@finos/legend-graph';
+import {
+  buildRuntimeValueOption,
+  getRuntimeOptionFormatter,
+  QueryBuilderClassSelector,
+} from '@finos/legend-query-builder';
+import { guaranteeNonNullable, guaranteeType } from '@finos/legend-shared';
+import { DepotEntityWithOrigin } from '@finos/legend-storage';
+import { flowResult } from 'mobx';
+import { observer } from 'mobx-react-lite';
+import { useEffect } from 'react';
+import type { DataProductWithLegacyOption } from '../../stores/data-space/DataProductSelectorState.js';
+import { formatDataProductOrSpaceOptionLabel } from '../shared/LegendQueryDataProductOptionLabel.js';
+import type { LegendQueryDataSpaceQueryBuilderState } from '../../stores/data-space/query-builder/LegendQueryDataSpaceQueryBuilderState.js';
+
+const resolveExecutionContextRuntimes = (
+  queryBuilderState: LegendQueryDataSpaceQueryBuilderState,
+): PackageableRuntime[] => {
+  if (queryBuilderState.dataSpaceAnalysisResult) {
+    const executionContext = Array.from(
+      queryBuilderState.dataSpaceAnalysisResult.executionContextsIndex.values(),
+    ).find(
+      (e) =>
+        e.mapping.path ===
+        queryBuilderState.executionContext.mapping.value.path,
+    );
+    return guaranteeNonNullable(executionContext).compatibleRuntimes;
+  }
+  return getMappingCompatibleRuntimes(
+    queryBuilderState.executionContext.mapping.value,
+    queryBuilderState.graphManagerState.usableRuntimes,
+  );
+};
+
+/**
+ * LegendQuery-specific DataSpace query builder setup panel. Unlike the base
+ * DataSpaceQueryBuilderSetupPanelContent, this component explicitly knows about
+ * both legacy data spaces and data products as combined options in the first
+ * dropdown, mirroring the same awareness in LegendQueryDataProductQueryBuilderState.
+ */
+const LegendQueryDataSpaceQueryBuilderSetupPanelContent = observer(
+  (props: { queryBuilderState: LegendQueryDataSpaceQueryBuilderState }) => {
+    const { queryBuilderState } = props;
+    const applicationStore = useApplicationStore();
+
+    // combined first dropdown: legacy data spaces + data products from depot
+    const legacyOptions = queryBuilderState.dataSpaceOptions; // applies prioritizeEntityFunc
+    const productOptions: DataProductWithLegacyOption[] = (
+      queryBuilderState.productSelectorState.dataProducts ?? []
+    ).map((e) => ({ label: e.name, value: e }));
+    const allOptions: DataProductWithLegacyOption[] = [
+      ...legacyOptions,
+      ...productOptions,
+    ];
+
+    const selectedOption = queryBuilderState.selectedDataSpaceOption;
+
+    const onOptionChange = (
+      option: DataProductWithLegacyOption | null,
+    ): void => {
+      if (!option) {
+        return;
+      }
+      const value = option.value;
+      if (value instanceof ResolvedDataSpaceEntityWithOrigin) {
+        queryBuilderState.queryChatState?.abort();
+        queryBuilderState
+          .onDataSpaceChange(value)
+          .catch(queryBuilderState.applicationStore.alertUnhandledError);
+      } else if (value instanceof DepotEntityWithOrigin) {
+        queryBuilderState.queryChatState?.abort();
+        queryBuilderState.onDataProductChange(value);
+      }
+    };
+
+    const openDataSpaceAdvancedSearch = (): void => {
+      if (queryBuilderState.isAdvancedDataSpaceSearchEnabled) {
+        queryBuilderState.showAdvancedSearchPanel(queryBuilderState.dataSpace);
+      }
+    };
+
+    // execution context
+    const executionContextOptions =
+      queryBuilderState.dataSpace.executionContexts
+        .map(buildExecutionContextOption)
+        .sort(compareLabelFn);
+    const showExecutionContextOptions = executionContextOptions.length > 1;
+    const selectedExecutionContextOption = buildExecutionContextOption(
+      queryBuilderState.executionContext,
+    );
+
+    const onExecutionContextOptionChange = async (
+      option: ExecutionContextOption,
+    ): Promise<void> => {
+      if (option.value === queryBuilderState.executionContext) {
+        return;
+      }
+      const currentMapping =
+        queryBuilderState.executionContext.mapping.value.path;
+      queryBuilderState.setExecutionContext(option.value);
+      await queryBuilderState.propagateExecutionContextChange(
+        currentMapping === option.value.mapping.value.path,
+      );
+      queryBuilderState.onExecutionContextChange?.(option.value);
+    };
+
+    const handleExecutionContextOptionChange = (
+      option: ExecutionContextOption,
+    ): void => {
+      flowResult(onExecutionContextOptionChange(option));
+    };
+
+    // runtime
+    const runtimeOptions = resolveExecutionContextRuntimes(queryBuilderState)
+      .map(
+        (rt) =>
+          new RuntimePointer(PackageableElementExplicitReference.create(rt)),
+      )
+      .map(buildRuntimeValueOption)
+      .sort(compareLabelFn);
+    const selectedRuntimeOption = queryBuilderState.executionContextState
+      .runtimeValue
+      ? buildRuntimeValueOption(
+          queryBuilderState.executionContextState.runtimeValue,
+        )
+      : null;
+    const changeRuntime = (option: { value: Runtime }): void => {
+      if (
+        option.value === queryBuilderState.executionContextState.runtimeValue
+      ) {
+        return;
+      }
+      queryBuilderState.changeRuntime(option.value);
+      queryBuilderState.onRuntimeChange?.(option.value);
+    };
+    const runtimeFilterOption = createFilter({
+      ignoreCase: true,
+      ignoreAccents: false,
+      stringify: (option: { data: { value: Runtime } }): string =>
+        guaranteeType(option.data.value, RuntimePointer).packageableRuntime
+          .value.path,
+    });
+
+    // class
+    const classes = resolveUsableDataSpaceClasses(
+      queryBuilderState.dataSpace,
+      queryBuilderState.executionContext.mapping.value,
+      queryBuilderState.graphManagerState,
+      queryBuilderState,
+    );
+
+    useEffect(() => {
+      flowResult(queryBuilderState.loadEntities()).catch(
+        applicationStore.alertUnhandledError,
+      );
+    }, [queryBuilderState, applicationStore]);
+
+    const copyDataSpaceLinkToClipboard = (): void => {
+      if (queryBuilderState.isDataSpaceLinkable) {
+        queryBuilderState.copyDataSpaceLinkToClipboard();
+      }
+    };
+
+    return (
+      <div className="query-builder__setup__config-group">
+        <PanelHeader title="properties">
+          <PanelHeaderActions>
+            <PanelHeaderActionItem
+              title="copy data product query set up link to clipboard"
+              onClick={copyDataSpaceLinkToClipboard}
+              disabled={!queryBuilderState.isDataSpaceLinkable}
+            >
+              <AnchorLinkIcon />
+            </PanelHeaderActionItem>
+            <ControlledDropdownMenu
+              className="panel__header__action query-builder__setup__config-group__header__dropdown-trigger"
+              title="Show Settings..."
+              content={
+                <MenuContent>
+                  <MenuContentItem
+                    onClick={(): void =>
+                      queryBuilderState.setShowRuntimeSelector(
+                        !queryBuilderState.showRuntimeSelector,
+                      )
+                    }
+                  >
+                    <MenuContentItemIcon>
+                      {queryBuilderState.showRuntimeSelector ? (
+                        <CheckIcon />
+                      ) : null}
+                    </MenuContentItemIcon>
+                    <MenuContentItemLabel>
+                      Show Runtime Selector
+                    </MenuContentItemLabel>
+                  </MenuContentItem>
+                </MenuContent>
+              }
+              menuProps={{
+                anchorOrigin: { vertical: 'bottom', horizontal: 'right' },
+                transformOrigin: { vertical: 'top', horizontal: 'right' },
+              }}
+            >
+              <MoreVerticalIcon className="query-builder__icon__more-options" />
+            </ControlledDropdownMenu>
+          </PanelHeaderActions>
+        </PanelHeader>
+        <div className="query-builder__setup__config-group__content">
+          <div className="query-builder__setup__config-group__item">
+            <label
+              className="btn--sm query-builder__setup__config-group__data-product"
+              title="data product"
+              htmlFor="query-builder__setup__data-space-selector"
+            >
+              Data Product
+            </label>
+            <CustomSelectorInput
+              inputId="query-builder__setup__data-space-selector"
+              className="panel__content__form__section__dropdown query-builder__setup__config-group__item__selector"
+              options={allOptions}
+              isLoading={queryBuilderState.loadEntitiesState.isInProgress}
+              onChange={onOptionChange}
+              value={selectedOption}
+              placeholder="Search for data product..."
+              escapeClearsValue={true}
+              darkMode={
+                !applicationStore.layoutService
+                  .TEMPORARY__isLightColorThemeEnabled
+              }
+              formatOptionLabel={formatDataProductOrSpaceOptionLabel}
+            />
+            {queryBuilderState.isAdvancedDataSpaceSearchEnabled && (
+              <>
+                <button
+                  tabIndex={-1}
+                  className="query-builder__setup__data-space-searcher__btn btn--dark"
+                  onClick={openDataSpaceAdvancedSearch}
+                  title="Open advanced search for data product..."
+                >
+                  <SearchIcon />
+                </button>
+                {queryBuilderState.advancedSearchState && (
+                  <DataSpaceAdvancedSearchModal
+                    searchState={queryBuilderState.advancedSearchState}
+                    onClose={() => queryBuilderState.hideAdvancedSearchPanel()}
+                  />
+                )}
+              </>
+            )}
+          </div>
+          {Boolean(showExecutionContextOptions) && (
+            <div className="query-builder__setup__config-group__item">
+              <label
+                className="btn--sm query-builder__setup__config-group__item__label"
+                title="execution context"
+                htmlFor="query-builder__setup__context-selector"
+              >
+                Context
+              </label>
+              <CustomSelectorInput
+                inputId="query-builder__setup__context-selector"
+                className="panel__content__form__section__dropdown query-builder__setup__config-group__item__selector"
+                placeholder="Choose an execution context..."
+                options={executionContextOptions}
+                disabled={
+                  executionContextOptions.length < 1 ||
+                  (executionContextOptions.length === 1 &&
+                    Boolean(selectedExecutionContextOption))
+                }
+                onChange={handleExecutionContextOptionChange}
+                value={selectedExecutionContextOption}
+                darkMode={
+                  !applicationStore.layoutService
+                    .TEMPORARY__isLightColorThemeEnabled
+                }
+              />
+            </div>
+          )}
+          {queryBuilderState.showRuntimeSelector && (
+            <div className="query-builder__setup__config-group__item">
+              <label
+                className="btn--sm query-builder__setup__config-group__item__label"
+                title="runtime"
+                htmlFor="query-builder__setup__runtime-selector"
+              >
+                Runtime
+              </label>
+              <CustomSelectorInput
+                inputId="query-builder__setup__runtime-selector"
+                className="panel__content__form__section__dropdown query-builder__setup__config-group__item__selector"
+                placeholder="Choose a runtime..."
+                noMatchMessage="No compatible runtime found for specified execution context"
+                options={runtimeOptions}
+                onChange={changeRuntime}
+                value={selectedRuntimeOption}
+                darkMode={
+                  !applicationStore.layoutService
+                    .TEMPORARY__isLightColorThemeEnabled
+                }
+                filterOption={runtimeFilterOption}
+                formatOptionLabel={getRuntimeOptionFormatter({
+                  darkMode:
+                    !applicationStore.layoutService
+                      .TEMPORARY__isLightColorThemeEnabled,
+                })}
+              />
+            </div>
+          )}
+          <div className="query-builder__setup__config-group__item">
+            <QueryBuilderClassSelector
+              queryBuilderState={queryBuilderState}
+              classes={classes}
+              onClassChange={queryBuilderState.onClassChange}
+              noMatchMessage="No compatible entity found for specified execution context"
+            />
+          </div>
+        </div>
+      </div>
+    );
+  },
+);
+
+export const renderLegendQueryDataSpaceQueryBuilderSetupPanelContent = (
+  queryBuilderState: LegendQueryDataSpaceQueryBuilderState,
+): React.ReactNode => (
+  <LegendQueryDataSpaceQueryBuilderSetupPanelContent
+    queryBuilderState={queryBuilderState}
+  />
+);

package/src/components/shared/LegendQueryDataProductOptionLabel.tsx

@@ -0,0 +1,80 @@
+/**
+ * Copyright (c) 2020-present, Goldman Sachs
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import { ResolvedDataSpaceEntityWithOrigin } from '@finos/legend-extension-dsl-data-space/application';
+import {
+  type DepotEntityWithOrigin,
+  generateGAVCoordinates,
+} from '@finos/legend-storage';
+import type { DataProductWithLegacyOption } from '../../stores/data-space/DataProductSelectorState.js';
+
+/**
+ * Inline badge rendered next to every DataProduct (Lakehouse) entry
+ * in a combined data-space / data-product selector dropdown.
+ */
+export const DataProductLakehouseBadge = (): React.ReactNode => (
+  <span className="query-builder__setup__data-product__option__tag query-builder__setup__data-product__option__tag--lakehouse">
+    Lakehouse
+  </span>
+);
+
+/**
+ * Option label for a combined dropdown that shows both legacy DataSpaces and
+ * depot DataProducts. DataProduct entries get a "Lakehouse" badge.
+ */
+export const formatDataProductOrSpaceOptionLabel = (
+  option: DataProductWithLegacyOption,
+): React.ReactNode => {
+  const { label, value } = option;
+  const isDataProduct = !(value instanceof ResolvedDataSpaceEntityWithOrigin);
+  const title = isDataProduct
+    ? `${label} - ${value.path} - ${generateGAVCoordinates(value.origin?.groupId ?? '', value.origin?.artifactId ?? '', value.origin?.versionId ?? '')}`
+    : `${label} - ${value.path}${
+        value.origin
+          ? ` - ${generateGAVCoordinates(value.origin.groupId, value.origin.artifactId, value.origin.versionId)}`
+          : ''
+      }`;
+  return (
+    <div className="query-builder__setup__data-product__option" title={title}>
+      <div className="query-builder__setup__data-product__option__label">
+        {label}
+      </div>
+      {isDataProduct && <DataProductLakehouseBadge />}
+    </div>
+  );
+};
+
+/**
+ * Option label for a DataProduct-only dropdown. Entries that are legacy
+ * DataSpaces (ResolvedDataSpaceEntityWithOrigin) do NOT get the badge; only
+ * actual Lakehouse DataProduct entries do.
+ */
+export const formatDataProductOptionLabel = (option: {
+  label: string;
+  value: DepotEntityWithOrigin;
+}): React.ReactNode => {
+  const { label, value } = option;
+  const isDataProduct = !(value instanceof ResolvedDataSpaceEntityWithOrigin);
+  const title = `${label} - ${value.path} - ${generateGAVCoordinates(value.origin?.groupId ?? '', value.origin?.artifactId ?? '', value.origin?.versionId ?? '')}`;
+  return (
+    <div className="query-builder__setup__data-product__option" title={title}>
+      <div className="query-builder__setup__data-product__option__label">
+        {label}
+      </div>
+      {isDataProduct && <DataProductLakehouseBadge />}
+    </div>
+  );
+};

package/src/stores/QueryEditorStore.ts

@@ -215,7 +215,7 @@ export class QueryCreatorState {
         'Query builder state required to build query to edit',
       );
       this.createQueryState.inProgress();
-      const rawLambda = queryBuilderState.buildQuery();
+      const rawLambda = queryBuilderState.buildQueryForPersistence();
       const config = this.editorStore.getPersistConfiguration(rawLambda, {
         update: true,
       });
@@ -1047,27 +1047,16 @@ export abstract class QueryEditorStore {
     const modelGroups = dataProduct.accessPointGroups.filter(
       filterByType(ModelAccessPointGroup),
     );
+    if (!modelGroups.length) {
+      throw new UnsupportedOperationError(
+        `Only Data Product with Model Access Points are currently supported in Query. ${dataProduct.path} does not have any Model Access Point Groups.`,
+      );
+    }
     if (accessId) {
       const matchingGroup = modelGroups.find((g) => g.id === accessId);
       if (matchingGroup) {
         return matchingGroup;
       }
-      // Search native execution contexts
-      const matchingNative =
-        dataProduct.nativeModelAccess?.nativeModelExecutionContexts.find(
-          (ctx) => ctx.key === accessId,
-        );
-      if (matchingNative) {
-        return matchingNative;
-      }
-    } else {
-      // No accessId: fall back to defaults
-      if (dataProduct.nativeModelAccess) {
-        return dataProduct.nativeModelAccess.defaultExecutionContext;
-      }
-      if (modelGroups.length > 0) {
-        return guaranteeNonNullable(modelGroups[0]);
-      }
     }
     throw new UnsupportedOperationError(
       `Can't resolve execution state for data product '${dataProduct.path}'${accessId ? ` with access ID '${accessId}'` : ''}`,
@@ -1172,6 +1161,9 @@ export abstract class QueryEditorStore {
     accessId: string,
     dataProductAccessType: DataProductAccessType,
     onDataProductChange: (val: DepotEntityWithOrigin) => Promise<void>,
+    onLegacyDataSpaceChange?:
+      | ((val: ResolvedDataSpaceEntityWithOrigin) => void)
+      | undefined,
     productSelectorState?: DataProductSelectorState | undefined,
   ): Promise<LegendQueryDataProductQueryBuilderState> {
     // 3. Build minimal graph and get analysis result
@@ -1229,6 +1221,7 @@ export abstract class QueryEditorStore {
           this.depotServerClient,
           this.applicationStore,
         ),
+      onLegacyDataSpaceChange,
       undefined,
       undefined,
       this.applicationStore.config.options.queryBuilderConfig,
@@ -1628,7 +1621,7 @@ export class ExistingQueryUpdateState {
         this.editorStore.queryBuilderState,
         'Query builder state required to build query to edit',
       );
-      const rawLambda = queryBuilderState.buildQuery();
+      const rawLambda = queryBuilderState.buildQueryForPersistence();
       const config = this.editorStore.getPersistConfiguration(rawLambda, {
         update: true,
       });

package/src/stores/data-product/query-builder/DataProductArtifactHelper.ts

@@ -34,14 +34,14 @@ export const resolveDefaultDataProductAccessType = (
     };
   }
 
-  const native =
-    dataProductArtifact.nativeModelAccess?.nativeModelExecutionContexts[0];
-  if (native) {
-    return {
-      type: DataProductAccessType.NATIVE,
-      id: native.key,
-    };
-  }
+  // const native =
+  //   dataProductArtifact.nativeModelAccess?.nativeModelExecutionContexts[0];
+  // if (native) {
+  //   return {
+  //     type: DataProductAccessType.NATIVE,
+  //     id: native.key,
+  //   };
+  // }
   throw new Error(
     `Data Product not supported for querying on legend query ${dataProductArtifact.dataProduct.path}. Must contain a model access point or native model access.`,
   );

package/src/stores/data-product/query-builder/LegendQueryDataProductQueryBuilderState.ts

@@ -22,7 +22,6 @@ import {
   type QueryBuilderActionConfig,
   type QueryBuilderConfig,
   type QueryBuilderWorkflowState,
-  type ExtraOptionsConfig,
 } from '@finos/legend-query-builder';
 import { renderLegendDataProductQueryBuilderSetupPanelContent } from '../../../components/data-product/LegendQueryDataProductQueryBuilder.js';
 import type { LegendQueryApplicationStore } from '../../LegendQueryBaseStore.js';
@@ -51,6 +50,7 @@ import {
 } from '../../../__lib__/LegendQueryNavigation.js';
 import { compareLabelFn } from '@finos/legend-art';
 import type { DataProductSelectorState } from '../../data-space/DataProductSelectorState.js';
+import { ResolvedDataSpaceEntityWithOrigin } from '@finos/legend-extension-dsl-data-space/application';
 import type { GeneratorFn } from '@finos/legend-shared';
 import { flowResult } from 'mobx';
 
@@ -58,7 +58,9 @@ export class LegendQueryDataProductQueryBuilderState extends DataProductQueryBui
   declare applicationStore: LegendQueryApplicationStore;
   depotServerClient: DepotServerClient;
   project: ProjectGAVCoordinates;
-  declare extraOptionsConfig: ExtraOptionsConfig<DepotEntityWithOrigin>;
+  readonly onLegacyDataSpaceChange:
+    | ((val: ResolvedDataSpaceEntityWithOrigin) => void)
+    | undefined;
   productSelectorState: DataProductSelectorState;
 
   constructor(
@@ -73,6 +75,9 @@ export class LegendQueryDataProductQueryBuilderState extends DataProductQueryBui
     project: ProjectGAVCoordinates,
     onDataProductChange: (val: DepotEntityWithOrigin) => Promise<void>,
     productSelectorState: DataProductSelectorState,
+    onLegacyDataSpaceChange?:
+      | ((val: ResolvedDataSpaceEntityWithOrigin) => void)
+      | undefined,
     onExecutionContextChange?:
       | ((val: NativeModelExecutionContext) => void)
       | undefined,
@@ -98,6 +103,18 @@ export class LegendQueryDataProductQueryBuilderState extends DataProductQueryBui
     this.project = project;
     this.depotServerClient = depotServerClient;
     this.productSelectorState = productSelectorState;
+    this.onLegacyDataSpaceChange = onLegacyDataSpaceChange;
+  }
+
+  override handleDataProductChange(val: DepotEntityWithOrigin): void {
+    if (
+      val instanceof ResolvedDataSpaceEntityWithOrigin &&
+      this.onLegacyDataSpaceChange
+    ) {
+      this.onLegacyDataSpaceChange(val);
+    } else {
+      super.handleDataProductChange(val);
+    }
   }
 
   override get dataProductOptions(): DataProductOption[] {
@@ -115,8 +132,6 @@ export class LegendQueryDataProductQueryBuilderState extends DataProductQueryBui
   }
 
   override *loadEntities(): GeneratorFn<void> {
-    yield flowResult(super.loadEntities());
-    // also ensure the shared selector has loaded depot-level entities
     if (!this.productSelectorState.isCompletelyLoaded) {
       yield flowResult(this.productSelectorState.loadProducts());
     }

package/src/stores/data-space/DataProductQueryCreatorStore.ts

@@ -60,10 +60,11 @@ import {
 import type { LegendQueryApplicationStore } from '../LegendQueryBaseStore.js';
 import {
   type DataSpaceQueryBuilderState,
-  ResolvedDataSpaceEntityWithOrigin,
+  type ResolvedDataSpaceEntityWithOrigin,
   createQueryClassTaggedValue,
   createQueryDataSpaceTaggedValue,
 } from '@finos/legend-extension-dsl-data-space/application';
+import { createQueryDataProductTaggedValue } from '../../components/data-product/QueryDataProductUtil.js';
 import { DataProductSelectorState } from './DataProductSelectorState.js';
 import { LegendQueryUserDataHelper } from '../../__lib__/LegendQueryUserDataHelper.js';
 import {
@@ -397,15 +398,14 @@ export class DataProductQueryCreatorStore extends QueryEditorStore {
       queryableDataProduct.id,
       queryableDataProduct.dataProductType as DataProductAccessType,
       async (dataProductInfo: DepotEntityWithOrigin) => {
-        if (dataProductInfo instanceof ResolvedDataSpaceEntityWithOrigin) {
-          flowResult(this.changeDataSpace(dataProductInfo)).catch(
-            this.applicationStore.alertUnhandledError,
-          );
-        } else {
-          flowResult(this.changeDataProduct(dataProductInfo)).catch(
-            this.applicationStore.alertUnhandledError,
-          );
-        }
+        flowResult(this.changeDataProduct(dataProductInfo)).catch(
+          this.applicationStore.alertUnhandledError,
+        );
+      },
+      (dataSpaceInfo: ResolvedDataSpaceEntityWithOrigin) => {
+        flowResult(this.changeDataSpace(dataSpaceInfo)).catch(
+          this.applicationStore.alertUnhandledError,
+        );
       },
       this.productSelectorState,
     );
@@ -666,7 +666,14 @@ export class DataProductQueryCreatorStore extends QueryEditorStore {
               createQueryClassTaggedValue(this.queryBuilderState.class.path),
             );
           }
-          taggedValues.push(createQueryDataSpaceTaggedValue(element.path));
+          if (element instanceof QueryableDataProduct) {
+            // For new DataProducts, use the 'dataProduct' tagged value so the
+            // query is not mis-classified as a DataSpace when reopened.
+            taggedValues.push(createQueryDataProductTaggedValue(element.path));
+          } else {
+            // Legacy DataSpace-based products use the 'dataSpace' tagged value.
+            taggedValues.push(createQueryDataSpaceTaggedValue(element.path));
+          }
           query.taggedValues = taggedValues;
         },
       };
@@ -755,7 +762,9 @@ export class DataProductQueryCreatorStore extends QueryEditorStore {
         }),
       ];
       val.taggedValues = [
-        createQueryDataSpaceTaggedValue(this.queryableElement.path),
+        this.queryableElement instanceof QueryableDataProduct
+          ? createQueryDataProductTaggedValue(this.queryableElement.path)
+          : createQueryDataSpaceTaggedValue(this.queryableElement.path),
       ];
       val.combineTaggedValuesCondition = true;
     }