Use DataPicker in JoinTablePicker

frontend/src/metabase-lib/join.ts

@@ -1,4 +1,10 @@
 import * as ML from "cljs/metabase.lib.js";
+import type {
+  CardId,
+  ConcreteTableId,
+  DatabaseId,
+  VirtualTableId,
+} from "metabase-types/api";
 
 import { expressionParts } from "./expression";
 import { isColumnMetadata } from "./internal";
@@ -237,13 +243,22 @@ export function joinedThing(query: Query, join: Join): Joinable {
   return ML.joined_thing(query, join);
 }
 
-export type PickerInfo = {
-  databaseId: number;
-  tableId: string;
-  cardId?: number;
-  isModel?: boolean;
+type CardPickerInfo = {
+  databaseId: DatabaseId;
+  tableId: VirtualTableId;
+  cardId: CardId;
+  isModel: boolean;
+};
+
+type TablePickerInfo = {
+  databaseId: DatabaseId;
+  tableId: ConcreteTableId;
+  cardId?: never;
+  isModel?: never;
 };
 
+export type PickerInfo = TablePickerInfo | CardPickerInfo;
+
 /**
  * Returns `null` when the joined table/card isn't available, e.g. due to sandboxing.
  */

frontend/src/metabase-lib/types.ts

@@ -6,6 +6,7 @@ import type {
   FieldValuesType,
   RowValue,
   TableId,
+  SchemaId,
 } from "metabase-types/api";
 
 import type {
@@ -134,6 +135,8 @@ export type TableDisplayInfo = {
   isSourceTable: boolean;
   isFromJoin: boolean;
   isImplicitlyJoinable: boolean;
+  schema: SchemaId;
+  isModel?: boolean;
 };
 
 export type CardDisplayInfo = TableDisplayInfo;

frontend/src/metabase/common/components/CollectionPicker/components/CollectionItemList.tsx

@@ -10,13 +10,16 @@ export const CollectionItemList = ({
   isFolder,
   isCurrentLevel,
   shouldDisableItem,
+  shouldShowItem,
 }: CollectionItemListProps) => {
   const {
     data: collectionItems,
     error,
     isLoading,
   } = useListCollectionItemsQuery<{
-    data: { data: CollectionPickerItem[] };
+    data: {
+      data: CollectionPickerItem[];
+    };
     error: any;
     isLoading: boolean;
   }>(query ? query : skipToken);
@@ -31,6 +34,7 @@ export const CollectionItemList = ({
       isFolder={isFolder}
       isCurrentLevel={isCurrentLevel}
       shouldDisableItem={shouldDisableItem}
+      shouldShowItem={shouldShowItem}
     />
   );
 };

frontend/src/metabase/common/components/CollectionPicker/components/CollectionItemPickerResolver.tsx

@@ -14,6 +14,7 @@ export const CollectionItemPickerResolver = ({
   isFolder,
   isCurrentLevel,
   shouldDisableItem,
+  shouldShowItem,
 }: CollectionItemListProps) => {
   if (!query) {
     return (
@@ -24,6 +25,7 @@ export const CollectionItemPickerResolver = ({
         isFolder={isFolder}
         isCurrentLevel={isCurrentLevel}
         shouldDisableItem={shouldDisableItem}
+        shouldShowItem={shouldShowItem}
       />
     );
   }
@@ -36,6 +38,7 @@ export const CollectionItemPickerResolver = ({
         isFolder={isFolder}
         isCurrentLevel={isCurrentLevel}
         shouldDisableItem={shouldDisableItem}
+        shouldShowItem={shouldShowItem}
         options={options}
       />
     );
@@ -49,6 +52,7 @@ export const CollectionItemPickerResolver = ({
       isFolder={isFolder}
       isCurrentLevel={isCurrentLevel}
       shouldDisableItem={shouldDisableItem}
+      shouldShowItem={shouldShowItem}
       options={options}
     />
   );

frontend/src/metabase/common/components/CollectionPicker/components/PersonalCollectionItemList.tsx

@@ -12,6 +12,7 @@ export const PersonalCollectionsItemList = ({
   isFolder,
   isCurrentLevel,
   shouldDisableItem,
+  shouldShowItem,
 }: CollectionItemListProps) => {
   const {
     data: collections,
@@ -36,6 +37,7 @@ export const PersonalCollectionsItemList = ({
       isFolder={isFolder}
       isCurrentLevel={isCurrentLevel}
       shouldDisableItem={shouldDisableItem}
+      shouldShowItem={shouldShowItem}
     />
   );
 };

frontend/src/metabase/common/components/CollectionPicker/components/RootItemList.tsx

@@ -36,6 +36,7 @@ export const RootItemList = ({
   isFolder,
   isCurrentLevel,
   shouldDisableItem,
+  shouldShowItem,
 }: CollectionItemListProps) => {
   const isAdmin = useSelector(getUserIsAdmin);
   const currentUser = useSelector(getUser);
@@ -135,6 +136,7 @@ export const RootItemList = ({
       isFolder={isFolder}
       isCurrentLevel={isCurrentLevel}
       shouldDisableItem={shouldDisableItem}
+      shouldShowItem={shouldShowItem}
     />
   );
 };

frontend/src/metabase/common/components/DataPicker/components/DataPickerModal.tsx

@@ -1,9 +1,12 @@
 import { useCallback } from "react";
 import { t } from "ttag";
-import _ from "underscore";
 
 import { getQuestionVirtualTableId } from "metabase-lib/v1/metadata/utils/saved-questions";
-import type { CollectionItemModel, TableId } from "metabase-types/api";
+import type {
+  CollectionItemModel,
+  DatabaseId,
+  TableId,
+} from "metabase-types/api";
 
 import type { EntityTab } from "../../EntityPicker";
 import { EntityPickerModal, defaultOptions } from "../../EntityPicker";
@@ -25,6 +28,11 @@ import {
 import { TablePicker } from "./TablePicker";
 
 interface Props {
+  /**
+   * Limit selection to a particular database
+   */
+  databaseId?: DatabaseId;
+  title: string;
   value: DataPickerValue | undefined;
   onChange: (value: TableId) => void;
   onClose: () => void;
@@ -41,8 +49,14 @@ const options: DataPickerModalOptions = {
   showRootCollection: true,
 };
 
-export const DataPickerModal = ({ value, onChange, onClose }: Props) => {
-  const { hasModels, hasQuestions } = useAvailableData();
+export const DataPickerModal = ({
+  databaseId,
+  title,
+  value,
+  onChange,
+  onClose,
+}: Props) => {
+  const { hasModels, hasQuestions } = useAvailableData({ databaseId });
 
   const handleChange = useCallback(
     (item: NotebookDataPickerValueItem) => {
@@ -78,6 +92,7 @@ export const DataPickerModal = ({ value, onChange, onClose }: Props) => {
           icon: "model",
           element: (
             <QuestionPicker
+              databaseId={databaseId}
               initialValue={isModelItem(value) ? value : undefined}
               models={MODEL_PICKER_MODELS}
               options={options}
@@ -92,6 +107,7 @@ export const DataPickerModal = ({ value, onChange, onClose }: Props) => {
       icon: "table",
       element: (
         <TablePicker
+          databaseId={databaseId}
           value={isTableItem(value) ? value : undefined}
           onChange={handleChange}
         />
@@ -104,6 +120,7 @@ export const DataPickerModal = ({ value, onChange, onClose }: Props) => {
           icon: "folder",
           element: (
             <QuestionPicker
+              databaseId={databaseId}
               initialValue={isQuestionItem(value) ? value : undefined}
               models={QUESTION_PICKER_MODELS}
               options={options}
@@ -120,11 +137,12 @@ export const DataPickerModal = ({ value, onChange, onClose }: Props) => {
   return (
     <EntityPickerModal
       canSelectItem
+      databaseId={databaseId}
       initialValue={value}
       options={options}
       selectedItem={value ?? null}
       tabs={tabs}
-      title={t`Pick your starting data`}
+      title={title}
       onClose={onClose}
       onItemSelect={handleChange}
     />

frontend/src/metabase/common/components/DataPicker/components/TablePicker.tsx

@@ -23,12 +23,18 @@ import { SchemaList } from "./SchemaList";
 import { TableList } from "./TableList";
 
 interface Props {
+  /**
+   * Limit selection to a particular database
+   */
+  databaseId?: DatabaseId;
   value: TablePickerValue | undefined;
   onChange: (value: NotebookDataPickerValueItem) => void;
 }
 
-export const TablePicker = ({ value, onChange }: Props) => {
-  const [dbId, setDbId] = useState<DatabaseId | undefined>(value?.db_id);
+export const TablePicker = ({ databaseId, value, onChange }: Props) => {
+  const [dbId, setDbId] = useState<DatabaseId | undefined>(
+    databaseId ?? value?.db_id,
+  );
   const [schemaName, setSchemaName] = useState<SchemaName | undefined>(
     value?.schema,
   );
@@ -109,14 +115,16 @@ export const TablePicker = ({ value, onChange }: Props) => {
       data-testid="nested-item-picker"
     >
       <Flex h="100%" w="fit-content">
-        <DatabaseList
-          databases={isLoadingDatabases ? undefined : databases?.data}
-          error={errorDatabases}
-          isCurrentLevel={!schemaName || (schemas?.length === 1 && !tableId)}
-          isLoading={isLoadingDatabases}
-          selectedItem={selectedDbItem}
-          onClick={folder => handleFolderSelect({ folder })}
-        />
+        {!databaseId && (
+          <DatabaseList
+            databases={isLoadingDatabases ? undefined : databases?.data}
+            error={errorDatabases}
+            isCurrentLevel={!schemaName || (schemas?.length === 1 && !tableId)}
+            isLoading={isLoadingDatabases}
+            selectedItem={selectedDbItem}
+            onClick={folder => handleFolderSelect({ folder })}
+          />
+        )}
 
         {isNotNull(dbId) && (
           <SchemaList

frontend/src/metabase/common/components/DataPicker/hooks/useAvailableData.ts

@@ -1,7 +1,16 @@
 import { useSearchQuery } from "metabase/api";
+import type { DatabaseId } from "metabase-types/api";
 
-export const useAvailableData = () => {
-  const { data } = useSearchQuery({ models: ["card"], limit: 0 });
+interface Props {
+  databaseId?: DatabaseId;
+}
+
+export const useAvailableData = ({ databaseId }: Props = {}) => {
+  const { data } = useSearchQuery({
+    limit: 0,
+    models: ["card"],
+    table_db_id: databaseId,
+  });
   const availableModels = data?.available_models ?? [];
   const hasModels = availableModels.includes("dataset");
   const hasQuestions = availableModels.includes("card");

frontend/src/metabase/common/components/DataPicker/utils.ts

@@ -1,5 +1,7 @@
 import { humanize, titleize } from "metabase/lib/formatting";
+import * as Lib from "metabase-lib";
 import TableEntity from "metabase-lib/v1/metadata/Table";
+import { getSchemaName } from "metabase-lib/v1/metadata/utils/schema";
 import type {
   Card,
   Database,
@@ -56,6 +58,35 @@ export const dataPickerValueFromTable = (
   };
 };
 
+export const dataPickerValueFromJoinable = (
+  query: Lib.Query,
+  stageIndex: number,
+  joinable: Lib.Joinable,
+): DataPickerValue | undefined => {
+  const pickerInfo = Lib.pickerInfo(query, joinable);
+  const displayInfo = Lib.displayInfo(query, stageIndex, joinable);
+
+  if (!pickerInfo) {
+    return undefined;
+  }
+
+  if (typeof pickerInfo.cardId === "number") {
+    return {
+      id: pickerInfo.cardId,
+      name: displayInfo.displayName,
+      model: displayInfo.isModel ? "dataset" : "card",
+    };
+  }
+
+  return {
+    id: pickerInfo.tableId,
+    name: displayInfo.displayName,
+    model: "table",
+    db_id: pickerInfo.databaseId,
+    schema: getSchemaName(displayInfo.schema),
+  };
+};
+
 const tablePickerValueFromTableEntity = (
   table: TableEntity,
 ): TablePickerValue => {

frontend/src/metabase/common/components/EntityPicker/components/EntityPickerModal/EntityPickerModal.tsx

@@ -7,6 +7,7 @@ import { BULK_ACTIONS_Z_INDEX } from "metabase/components/BulkActionBar";
 import { useModalOpen } from "metabase/hooks/use-modal-open";
 import { Modal } from "metabase/ui";
 import type {
+  DatabaseId,
   SearchModel,
   SearchResult,
   SearchResultId,
@@ -44,6 +45,10 @@ export const defaultOptions: EntityPickerModalOptions = {
 export const ENTITY_PICKER_Z_INDEX = BULK_ACTIONS_Z_INDEX;
 
 export interface EntityPickerModalProps<Model extends string, Item> {
+  /**
+   * Limit selection to a particular database
+   */
+  databaseId?: DatabaseId;
   title?: string;
   selectedItem: Item | null;
   initialValue?: Partial<Item>;
@@ -63,6 +68,7 @@ export function EntityPickerModal<
   Model extends SearchModel,
   Item extends TypeWithModel<Id, Model>,
 >({
+  databaseId,
   title = t`Choose an item`,
   onItemSelect,
   canSelectItem,
@@ -122,6 +128,7 @@ export function EntityPickerModal<
             <Modal.Title lh="2.5rem">{title}</Modal.Title>
             {hydratedOptions.showSearch && (
               <EntityPickerSearchInput
+                databaseId={databaseId}
                 models={tabModels}
                 setSearchResults={setSearchResults}
                 searchQuery={searchQuery}