@marimo-team/frontend 0.25.2-dev56 → 0.25.2-dev62

package/src/components/data-table/export-actions.tsx

@@ -3,10 +3,15 @@
 import { useAtomValue } from "jotai";
 import {
   AlertCircleIcon,
+  BracesIcon,
+  BrickWallIcon,
   ChevronRightIcon,
   ClipboardCopyIcon,
   DownloadIcon,
+  FileTextIcon,
   InfoIcon,
+  MapIcon,
+  TableIcon,
 } from "lucide-react";
 import React from "react";
 import { downloadSizeLimitAtom } from "./download-policy/atoms";
@@ -14,6 +19,8 @@ import type {
   DownloadAsArgs,
   DownloadAsOptions,
   DownloadFormat,
+  ExportMetadata,
+  GetExportMetadata,
 } from "./schemas";
 import { logNever } from "@/utils/assertNever";
 import { cn } from "@/utils/cn";
@@ -47,6 +54,7 @@ const EXPORT_OPTIONS = [
     label: "CSV",
     format: "csv",
     description: "Comma-separated values",
+    icon: TableIcon,
     canDownload: true,
     canCopy: true,
   },
@@ -54,6 +62,7 @@ const EXPORT_OPTIONS = [
     label: "TSV",
     format: "tsv",
     description: "Best for Excel and Google Sheets",
+    icon: TableIcon,
     canDownload: true,
     canCopy: true,
   },
@@ -61,6 +70,7 @@ const EXPORT_OPTIONS = [
     label: "JSON",
     format: "json",
     description: "Raw JSON data",
+    icon: BracesIcon,
     canDownload: true,
     canCopy: true,
   },
@@ -68,13 +78,23 @@ const EXPORT_OPTIONS = [
     label: "Parquet",
     format: "parquet",
     description: "Columnar binary format",
+    icon: BrickWallIcon,
     canDownload: true,
     canCopy: false,
   },
+  {
+    label: "GeoJSON",
+    format: "geojson",
+    description: "Geometry in longitude and latitude",
+    icon: MapIcon,
+    canDownload: true,
+    canCopy: true,
+  },
   {
     label: "Markdown",
     format: "markdown",
     description: "Preserves hyperlinks and formatting",
+    icon: FileTextIcon,
     canDownload: false,
     canCopy: true,
   },
@@ -82,6 +102,10 @@ const EXPORT_OPTIONS = [
 
 type ExportFormat = (typeof EXPORT_OPTIONS)[number]["format"];
 type CopyFormat = Exclude<ExportFormat, "parquet">;
+type GeographicFormat = "parquet" | "geojson";
+
+const isGeographic = (format: ExportFormat): format is GeographicFormat =>
+  format === "parquet" || format === "geojson";
 type ExportAction =
   | { destination: "download"; format: DownloadFormat }
   | { destination: "copy"; format: CopyFormat };
@@ -95,10 +119,8 @@ type ExportFailure =
     }
   | {
       kind: "missing-packages";
-      title: string;
       description?: string | null;
       packages: string[];
-      featureName: string;
       action: ExportAction;
     };
 
@@ -108,6 +130,7 @@ const COPY_SOURCE_FORMAT: Record<CopyFormat, DownloadFormat> = {
   csv: "csv",
   tsv: "tsv",
   json: "json",
+  geojson: "geojson",
   markdown: "json",
 };
 
@@ -140,7 +163,7 @@ const DEFAULT_SETTINGS: ExportSettings = {
 
 type ConfigurableFormat = keyof ExportSettings;
 
-const isConfigurable = (format: ExportFormat): format is ConfigurableFormat =>
+const isConfigurable = (format: ExportFormat): format is keyof ExportSettings =>
   format in DEFAULT_SETTINGS;
 
 /**
@@ -182,6 +205,7 @@ function requestOptions(
       }
       break;
     case "parquet":
+    case "geojson":
       break;
     default:
       logNever(format);
@@ -212,6 +236,7 @@ function settingsSummary(
         ? "Escapes non-ASCII"
         : "Keeps non-ASCII";
     case "parquet":
+    case "geojson":
     case "markdown":
       return null;
     default:
@@ -222,6 +247,8 @@ function settingsSummary(
 
 export interface ExportActionProps {
   downloadAs: DownloadAsArgs;
+  getExportMetadata?: GetExportMetadata;
+  metadataSource?: unknown;
   // JSON-serialized size of the currently-rendered data. Used together with
   // downloadSizeLimitAtom to disable the Export button when a host (e.g.,
   // marimo-lsp inside VS Code) declares a download size cap. Null/undefined
@@ -230,8 +257,14 @@ export interface ExportActionProps {
   sizeBytesIsLoading?: boolean;
 }
 
-const labelForFormat = (format: ExportFormat): string =>
-  EXPORT_OPTIONS.find((option) => option.format === format)?.label ?? format;
+const labelForFormat = (
+  format: ExportFormat | DownloadFormat,
+  hasGeometry = false,
+): string =>
+  format === "parquet" && hasGeometry
+    ? "GeoParquet"
+    : (EXPORT_OPTIONS.find((option) => option.format === format)?.label ??
+      format);
 
 const failureTitleForAction = (action: ExportAction): string =>
   action.destination === "download"
@@ -242,15 +275,62 @@ const failureDescription = (error: unknown): string =>
   typeof error === "string" ? error : prettyError(error);
 
 export const ExportActions: React.FC<ExportActionProps> = (props) => {
+  const { getExportMetadata, metadataSource } = props;
   const [exportDialogOpen, setExportDialogOpen] = React.useState(false);
   const [failure, setFailure] = React.useState<ExportFailure | null>(null);
   const [settings, setSettings] =
     React.useState<ExportSettings>(DEFAULT_SETTINGS);
-  const [expandedFormat, setExpandedFormat] =
-    React.useState<ConfigurableFormat | null>(null);
+  const [metadata, setMetadata] = React.useState<ExportMetadata | null>(null);
+  const [metadataLoading, setMetadataLoading] = React.useState(false);
+  const [metadataError, setMetadataError] = React.useState<string | null>(null);
+  const [metadataRefresh, setMetadataRefresh] = React.useState(0);
+  const [geometryColumn, setGeometryColumn] = React.useState<string | null>(
+    null,
+  );
+  const hasGeometry = !!metadata?.geometry_columns.length;
+  const [expandedFormat, setExpandedFormat] = React.useState<
+    ConfigurableFormat | GeographicFormat | null
+  >(null);
   const triggerRef = React.useRef<HTMLButtonElement>(null);
   const dialogRef = React.useRef<HTMLDivElement>(null);
   const latestActionId = React.useRef(0);
+  React.useEffect(() => {
+    latestActionId.current += 1;
+  }, [metadataSource]);
+  React.useEffect(() => {
+    if (!exportDialogOpen || !getExportMetadata) {
+      return;
+    }
+    let active = true;
+    setMetadata(null);
+    setMetadataLoading(true);
+    setMetadataError(null);
+    void getExportMetadata({}).then(
+      (result) => {
+        if (!active) {
+          return;
+        }
+        setMetadata(result);
+        setGeometryColumn((current) =>
+          current !== null &&
+          result.geometry_columns.some((column) => column.name === current)
+            ? current
+            : result.default_geometry_column,
+        );
+        setMetadataLoading(false);
+      },
+      (error: unknown) => {
+        if (!active) {
+          return;
+        }
+        setMetadataError(failureDescription(error));
+        setMetadataLoading(false);
+      },
+    );
+    return () => {
+      active = false;
+    };
+  }, [exportDialogOpen, getExportMetadata, metadataSource, metadataRefresh]);
   const policy = useAtomValue(downloadSizeLimitAtom);
   const overLimit = !!(
     policy &&
@@ -288,6 +368,9 @@ export const ExportActions: React.FC<ExportActionProps> = (props) => {
       setFailure(null);
       setExpandedFormat(null);
       setSettings(DEFAULT_SETTINGS);
+      setMetadata(null);
+      setMetadataError(null);
+      setGeometryColumn(null);
     }
   };
 
@@ -313,27 +396,33 @@ export const ExportActions: React.FC<ExportActionProps> = (props) => {
     filename: string;
   } | null> => {
     const options = requestOptions(settings, format, action.destination);
-    const response = await props.downloadAs(
-      options ? { format, options } : { format },
-    );
+    const request = {
+      format,
+      ...(options ? { options } : {}),
+      ...(isGeographic(format) && geometryColumn !== null
+        ? { geometry_column: geometryColumn }
+        : {}),
+    };
+    const response = await props.downloadAs(request);
+    if (actionId !== latestActionId.current) {
+      return null;
+    }
 
     if (response.missing_packages && response.missing_packages.length > 0) {
       setActionFailure(actionId, {
         kind: "missing-packages",
-        title: "Export failed",
         packages: response.missing_packages,
-        featureName: `${labelForFormat(action.format)} export`,
         description: response.error,
         action,
       });
       return null;
     }
 
-    if (response.error) {
+    if (response.error || response.code || !response.url) {
       setActionFailure(actionId, {
         kind: "error",
         title: failureTitleForAction(action),
-        description: response.error,
+        description: response.error ?? "The export did not produce a file.",
         action,
       });
       return null;
@@ -348,7 +437,7 @@ export const ExportActions: React.FC<ExportActionProps> = (props) => {
   const handleDownload = async (format: DownloadFormat) => {
     const action: ExportAction = { destination: "download", format };
     const actionId = beginAction();
-    const label = labelForFormat(format);
+    const label = labelForFormat(format, hasGeometry);
     try {
       const ok = await withLoadingToast(
         `Preparing ${label} export...`,
@@ -357,6 +446,9 @@ export const ExportActions: React.FC<ExportActionProps> = (props) => {
           if (!result) {
             return false;
           }
+          if (actionId !== latestActionId.current) {
+            return false;
+          }
           const rawName = (result.filename ?? "").trim();
           const baseName = Filenames.withoutExtension(rawName) || "download";
           const downloadName = `${baseName}.${format}`;
@@ -396,7 +488,7 @@ export const ExportActions: React.FC<ExportActionProps> = (props) => {
     actionId: number,
   ) => {
     await withLoadingToast(
-      `Preparing ${labelForFormat(format)} for clipboard...`,
+      `Preparing ${labelForFormat(format, hasGeometry)} for clipboard...`,
       async () => {
         const sourceFormat = COPY_SOURCE_FORMAT[format];
         const result = await resolveDownloadUrl(sourceFormat, action, actionId);
@@ -408,6 +500,7 @@ export const ExportActions: React.FC<ExportActionProps> = (props) => {
         switch (format) {
           case "tsv":
           case "csv":
+          case "geojson":
             text = await fetchText(result.url);
             break;
           case "json": {
@@ -425,6 +518,9 @@ export const ExportActions: React.FC<ExportActionProps> = (props) => {
             return;
         }
 
+        if (actionId !== latestActionId.current) {
+          return;
+        }
         await copyToClipboard(text);
         toast({
           title: "Copied to clipboard",
@@ -457,8 +553,61 @@ export const ExportActions: React.FC<ExportActionProps> = (props) => {
     }
   };
 
+  const retryAfterInstall = async (action: ExportAction) => {
+    if (isGeographic(action.format) && getExportMetadata) {
+      const actionId = latestActionId.current;
+      let result: ExportMetadata;
+      try {
+        result = await getExportMetadata({});
+      } catch (error) {
+        if (actionId === latestActionId.current) {
+          setMetadataError(failureDescription(error));
+        }
+        return;
+      }
+      if (actionId !== latestActionId.current) {
+        return;
+      }
+      setMetadata(result);
+      setMetadataError(null);
+      if (result.formats[action.format]?.available === false) {
+        setFailure(null);
+        return;
+      }
+      if (
+        geometryColumn !== null &&
+        !result.geometry_columns.some(
+          (column) => column.name === geometryColumn,
+        )
+      ) {
+        setGeometryColumn(result.default_geometry_column);
+        setActionFailure(actionId, {
+          kind: "error",
+          title: failureTitleForAction(action),
+          description:
+            "The selected geometry is no longer available. Choose another geometry to export.",
+          action,
+        });
+        return;
+      }
+      if (result.geometry_columns.length > 0 && geometryColumn === null) {
+        return;
+      }
+      if (
+        action.format === "geojson" &&
+        !result.geometry_columns.find(
+          (column) => column.name === geometryColumn,
+        )?.crs
+      ) {
+        setFailure(null);
+        return;
+      }
+    }
+    retryAction(action);
+  };
+
   const toggleExpanded = (format: ExportFormat) => {
-    if (!isConfigurable(format)) {
+    if (!isConfigurable(format) && !(isGeographic(format) && hasGeometry)) {
       return;
     }
     setExpandedFormat((current) => (current === format ? null : format));
@@ -504,30 +653,80 @@ export const ExportActions: React.FC<ExportActionProps> = (props) => {
           aria-label="Export formats"
           className="list-none overflow-hidden rounded-md border divide-y"
         >
-          {EXPORT_OPTIONS.map((option) => {
-            const configurableFormat = isConfigurable(option.format)
-              ? option.format
-              : null;
+          {EXPORT_OPTIONS.filter(
+            (option) => option.format !== "geojson" || hasGeometry,
+          ).map((option) => {
+            const optionLabel = labelForFormat(option.format, hasGeometry);
+            const configurableFormat =
+              isConfigurable(option.format) ||
+              (isGeographic(option.format) && hasGeometry)
+                ? option.format
+                : null;
             const expanded =
               configurableFormat !== null &&
               expandedFormat === configurableFormat;
             const panelId = `export-options-${option.format}`;
             const rowFailure =
               failure?.action.format === option.format ? failure : null;
-            const summary = settingsSummary(settings, option.format);
+            const summary =
+              isGeographic(option.format) && hasGeometry
+                ? geometryColumn === ""
+                  ? "(unnamed geometry)"
+                  : (geometryColumn ?? "Choose geometry")
+                : settingsSummary(settings, option.format);
+            const selectedGeometryIndex =
+              metadata?.geometry_columns.findIndex(
+                (column) => column.name === geometryColumn,
+              ) ?? -1;
+            const eligibility = isGeographic(option.format)
+              ? metadata?.formats[option.format]
+              : undefined;
+            const selectedGeometry =
+              metadata?.geometry_columns[selectedGeometryIndex];
+            const missingCrs =
+              option.format === "geojson" &&
+              selectedGeometry !== undefined &&
+              !selectedGeometry.crs;
+            const geometryReason = metadataError
+              ? `Could not load geometry options: ${metadataError}`
+              : (eligibility?.reason ??
+                (hasGeometry && geometryColumn === null
+                  ? `Choose a primary geometry column to export ${optionLabel}.`
+                  : missingCrs
+                    ? `Geometry column ${geometryColumn === "" ? "(unnamed geometry)" : geometryColumn} needs a CRS for GeoJSON export. Declare its source CRS in Python before exporting.`
+                    : null));
+            const missingPackageFailure =
+              rowFailure?.kind === "missing-packages" ? rowFailure : null;
+            const missingPackages =
+              missingPackageFailure?.packages ??
+              (eligibility?.available === false
+                ? eligibility.missing_packages
+                : []);
+            const geographicUnavailable =
+              isGeographic(option.format) &&
+              ((!!getExportMetadata && !metadata) ||
+                metadataLoading ||
+                !!metadataError ||
+                eligibility?.available === false ||
+                missingCrs ||
+                (!!metadata?.geometry_columns.length &&
+                  geometryColumn === null));
+            const copyUnavailable =
+              isGeographic(option.format) &&
+              (geographicUnavailable || missingPackages.length > 0);
             return (
               <li
                 key={option.format}
                 data-testid={`export-row-${option.format}`}
                 className={cn(expanded && "bg-muted/40")}
               >
-                <div className="grid min-h-[60px] grid-cols-[28px_1fr_auto] items-center gap-x-1 px-2.5 py-2 hover:bg-accent/50">
+                <div className="grid min-h-[60px] grid-cols-[28px_minmax(0,1fr)_auto] items-center gap-x-1 px-2.5 py-2 hover:bg-accent/50">
                   {configurableFormat ? (
                     <Button
                       type="button"
                       variant="ghost"
                       size="icon"
-                      aria-label={`${option.label} options`}
+                      aria-label={`${optionLabel} options`}
                       aria-expanded={expanded}
                       aria-controls={panelId}
                       onClick={() => toggleExpanded(option.format)}
@@ -542,30 +741,39 @@ export const ExportActions: React.FC<ExportActionProps> = (props) => {
                   ) : (
                     <span aria-hidden={true} className="h-6 w-6" />
                   )}
-                  <div className="grid gap-1 pl-1">
-                    <span className="flex items-baseline gap-x-2">
-                      <span className="text-sm font-medium leading-tight">
-                        {option.label}
-                      </span>
-                      {summary && (
-                        <span
-                          data-testid={`export-summary-${option.format}`}
-                          className="text-xs text-muted-foreground"
-                        >
-                          {summary}
+                  <div className="flex min-w-0 items-center gap-3 pl-1">
+                    <option.icon
+                      aria-hidden={true}
+                      data-testid={`export-format-icon-${option.format}`}
+                      className="h-4 w-4 shrink-0 text-muted-foreground"
+                    />
+                    <div className="grid min-w-0 gap-1">
+                      <span className="flex min-w-0 items-baseline gap-x-2">
+                        <span className="text-sm font-medium leading-tight">
+                          {optionLabel}
                         </span>
-                      )}
-                    </span>
-                    <span className="text-xs text-muted-foreground leading-tight">
-                      {option.description}
-                    </span>
+                        {summary && (
+                          <span
+                            data-testid={`export-summary-${option.format}`}
+                            className="min-w-0 truncate text-xs text-muted-foreground"
+                          >
+                            {summary}
+                          </span>
+                        )}
+                      </span>
+                      <span className="text-xs text-muted-foreground leading-tight">
+                        {option.format === "parquet" && hasGeometry
+                          ? "Parquet with geometry and CRS metadata"
+                          : option.description}
+                      </span>
+                    </div>
                   </div>
                   <div className="flex items-center gap-0.5">
                     <Tooltip
                       content={
                         option.canDownload
-                          ? `Download ${option.label}`
-                          : `${option.label} is available for copy only.`
+                          ? `Download ${optionLabel}`
+                          : `${optionLabel} is available for copy only.`
                       }
                     >
                       <span className="inline-flex">
@@ -573,10 +781,16 @@ export const ExportActions: React.FC<ExportActionProps> = (props) => {
                           type="button"
                           variant="ghost"
                           size="icon"
-                          aria-label={`Download ${option.label}`}
-                          disabled={!option.canDownload}
+                          aria-label={`Download ${optionLabel}`}
+                          disabled={
+                            !option.canDownload ||
+                            geographicUnavailable ||
+                            missingPackages.length > 0
+                          }
                           onClick={
-                            option.canDownload
+                            option.canDownload &&
+                            !geographicUnavailable &&
+                            missingPackages.length === 0
                               ? () => {
                                   void handleDownload(option.format);
                                 }
@@ -590,8 +804,8 @@ export const ExportActions: React.FC<ExportActionProps> = (props) => {
                     <Tooltip
                       content={
                         option.canCopy
-                          ? `Copy ${option.label}`
-                          : `${option.label} is available for download only.`
+                          ? `Copy ${optionLabel}`
+                          : `${optionLabel} is available for download only.`
                       }
                     >
                       <span className="inline-flex">
@@ -599,10 +813,10 @@ export const ExportActions: React.FC<ExportActionProps> = (props) => {
                           type="button"
                           variant="ghost"
                           size="icon"
-                          aria-label={`Copy ${option.label}`}
-                          disabled={!option.canCopy}
+                          aria-label={`Copy ${optionLabel}`}
+                          disabled={!option.canCopy || copyUnavailable}
                           onClick={
-                            option.canCopy
+                            option.canCopy && !copyUnavailable
                               ? () => {
                                   void handleCopyAction(option.format);
                                 }
@@ -620,42 +834,112 @@ export const ExportActions: React.FC<ExportActionProps> = (props) => {
                     id={panelId}
                     className="grid grid-cols-2 gap-3 border-t bg-muted/50 px-3.5 py-3.5 pl-[52px]"
                   >
-                    <legend className="sr-only">{option.label} options</legend>
-                    <ExportSettingsFields
-                      format={configurableFormat}
-                      settings={settings}
-                      onChange={setSettings}
-                    />
+                    <legend className="sr-only">{optionLabel} options</legend>
+                    {isGeographic(configurableFormat) && (
+                      <p className="col-span-2 text-xs text-muted-foreground">
+                        {configurableFormat === "geojson"
+                          ? "Reproject the selected geometry to longitude and latitude (WGS 84). Keep other geometry columns as complete WKT in their source coordinates."
+                          : "Use the selected column as the primary geometry. Preserve all geometry columns and their coordinate reference systems (CRS)."}
+                      </p>
+                    )}
+                    {isGeographic(configurableFormat) ? (
+                      <SettingField
+                        id="export-primary-geometry"
+                        label="Primary geometry"
+                        help={`Geometry column to use as the primary geometry in the ${optionLabel} file.`}
+                      >
+                        <NativeSelect
+                          id="export-primary-geometry"
+                          className="mb-0 w-full"
+                          value={String(selectedGeometryIndex)}
+                          onChange={(event) => {
+                            const index = Number(event.target.value);
+                            setGeometryColumn(
+                              metadata?.geometry_columns[index]?.name ?? null,
+                            );
+                          }}
+                        >
+                          <option value="-1" disabled={true}>
+                            Choose geometry
+                          </option>
+                          {metadata?.geometry_columns.map((column, index) => (
+                            <option key={column.name} value={String(index)}>
+                              {column.name || "(unnamed geometry)"}
+                            </option>
+                          ))}
+                        </NativeSelect>
+                      </SettingField>
+                    ) : (
+                      <ExportSettingsFields
+                        format={configurableFormat}
+                        settings={settings}
+                        onChange={setSettings}
+                      />
+                    )}
                   </fieldset>
                 )}
-                {rowFailure && (
+                {isGeographic(option.format) &&
+                  missingPackages.length === 0 &&
+                  geometryReason && (
+                    <div className="flex items-center justify-between gap-2 border-t px-3 py-2 pl-[42px] text-xs text-muted-foreground">
+                      <span>{geometryReason}</span>
+                      {metadataError && (
+                        <Button
+                          type="button"
+                          size="xs"
+                          variant="outline"
+                          onClick={() =>
+                            setMetadataRefresh((current) => current + 1)
+                          }
+                        >
+                          Retry
+                        </Button>
+                      )}
+                    </div>
+                  )}
+                {missingPackages.length > 0 && (
                   <div className="border-t px-3 py-2 pl-[42px]">
-                    <Alert
-                      variant="destructive"
-                      className="p-3 has-[svg]:pl-9 [&>svg]:left-3 [&>svg]:top-3"
-                    >
-                      <AlertCircleIcon className="h-4 w-4" />
-                      <div>
-                        <AlertTitle className="text-sm">
-                          {rowFailure.title}
-                        </AlertTitle>
-                        <AlertDescription className="text-xs">
-                          {rowFailure.kind === "missing-packages" ? (
-                            <MissingPackagePrompt
-                              packages={rowFailure.packages}
-                              featureName={rowFailure.featureName}
-                              description={rowFailure.description}
-                              onInstall={() => retryAction(rowFailure.action)}
-                              className="items-start"
-                            />
-                          ) : (
-                            rowFailure.description
-                          )}
-                        </AlertDescription>
-                      </div>
-                    </Alert>
+                    <MissingPackagePrompt
+                      packages={missingPackages}
+                      featureName={`${optionLabel} export`}
+                      description={
+                        missingPackageFailure?.description ??
+                        eligibility?.reason
+                      }
+                      onInstall={() => {
+                        void retryAfterInstall(
+                          missingPackageFailure?.action ?? {
+                            destination: "download",
+                            format:
+                              option.format === "geojson"
+                                ? "geojson"
+                                : "parquet",
+                          },
+                        );
+                      }}
+                      className="items-start"
+                    />
                   </div>
                 )}
+                {rowFailure?.kind === "error" &&
+                  missingPackages.length === 0 && (
+                    <div className="border-t px-3 py-2 pl-[42px]">
+                      <Alert
+                        variant="destructive"
+                        className="p-3 has-[svg]:pl-9 [&>svg]:left-3 [&>svg]:top-3"
+                      >
+                        <AlertCircleIcon className="h-4 w-4" />
+                        <div>
+                          <AlertTitle className="text-sm">
+                            {rowFailure.title}
+                          </AlertTitle>
+                          <AlertDescription className="text-xs">
+                            {rowFailure.description}
+                          </AlertDescription>
+                        </div>
+                      </Alert>
+                    </div>
+                  )}
               </li>
             );
           })}