Runtime: Deprecate runtime/server/export.go – migrate frontend to use…

… the new Export API (#3419)

* added export impl

* UI changes to use the new API

* updated limit clause

---------

Co-authored-by: Aditya Hegde <[email protected]>

runtime/queries/table_head.go

@@ -15,6 +15,7 @@ type TableHead struct {
 	TableName string
 	Limit     int
 	Result    []*structpb.Struct
+	Schema    *runtimev1.StructType
 }
 
 var _ runtime.Query = &TableHead{}
@@ -79,9 +80,77 @@ func (q *TableHead) Resolve(ctx context.Context, rt *runtime.Runtime, instanceID
 	}
 
 	q.Result = data
+	q.Schema = rows.Schema
 	return nil
 }
 
 func (q *TableHead) Export(ctx context.Context, rt *runtime.Runtime, instanceID string, w io.Writer, opts *runtime.ExportOptions) error {
-	return ErrExportNotSupported
+	olap, release, err := rt.OLAP(ctx, instanceID)
+	if err != nil {
+		return err
+	}
+	defer release()
+
+	switch olap.Dialect() {
+	case drivers.DialectDuckDB:
+		if opts.Format == runtimev1.ExportFormat_EXPORT_FORMAT_CSV || opts.Format == runtimev1.ExportFormat_EXPORT_FORMAT_PARQUET {
+			filename := q.TableName
+
+			limitClause := ""
+			if q.Limit > 0 {
+				limitClause = fmt.Sprintf(" LIMIT %d", q.Limit)
+			}
+
+			sql := fmt.Sprintf(
+				`SELECT * FROM %s%s`,
+				safeName(q.TableName),
+				limitClause,
+			)
+			args := []interface{}{}
+			if err := duckDBCopyExport(ctx, w, opts, sql, args, filename, olap, opts.Format); err != nil {
+				return err
+			}
+		} else {
+			if err := q.generalExport(ctx, rt, instanceID, w, opts, olap); err != nil {
+				return err
+			}
+		}
+	case drivers.DialectDruid:
+		if err := q.generalExport(ctx, rt, instanceID, w, opts, olap); err != nil {
+			return err
+		}
+	default:
+		return fmt.Errorf("not available for dialect '%s'", olap.Dialect())
+	}
+
+	return nil
+}
+
+func (q *TableHead) generalExport(ctx context.Context, rt *runtime.Runtime, instanceID string, w io.Writer, opts *runtime.ExportOptions, olap drivers.OLAPStore) error {
+	err := q.Resolve(ctx, rt, instanceID, opts.Priority)
+	if err != nil {
+		return err
+	}
+
+	if opts.PreWriteHook != nil {
+		err = opts.PreWriteHook(q.TableName)
+		if err != nil {
+			return err
+		}
+	}
+
+	meta := structTypeToMetricsViewColumn(q.Schema)
+
+	switch opts.Format {
+	case runtimev1.ExportFormat_EXPORT_FORMAT_UNSPECIFIED:
+		return fmt.Errorf("unspecified format")
+	case runtimev1.ExportFormat_EXPORT_FORMAT_CSV:
+		return writeCSV(meta, q.Result, w)
+	case runtimev1.ExportFormat_EXPORT_FORMAT_XLSX:
+		return writeXLSX(meta, q.Result, w)
+	case runtimev1.ExportFormat_EXPORT_FORMAT_PARQUET:
+		return writeParquet(meta, q.Result, w)
+	}
+
+	return nil
 }

runtime/server/downloads.go

@@ -292,6 +292,18 @@ func (s *Server) downloadHandler(w http.ResponseWriter, req *http.Request) {
 			MetricsView:         mv,
 			ResolvedMVSecurity:  security,
 		}
+	case *runtimev1.Query_TableRowsRequest:
+		r := v.TableRowsRequest
+		if !auth.GetClaims(req.Context()).CanInstance(r.InstanceId, auth.ReadOLAP) {
+			http.Error(w, "action not allowed", http.StatusUnauthorized)
+			return
+		}
+
+		q = &queries.TableHead{
+			TableName: r.TableName,
+			Limit:     int(r.Limit),
+			Result:    nil,
+		}
 	default:
 		http.Error(w, fmt.Sprintf("unsupported request type: %s", reflect.TypeOf(v).Name()), http.StatusBadRequest)
 		return

runtime/server/server.go

@@ -190,13 +190,6 @@ func (s *Server) HTTPHandler(ctx context.Context, registerAdditionalHandlers fun
 		panic(err)
 	}
 
-	// One-off REST-only path for file export
-	// NOTE: It's local only and we should deprecate it in favor of a cloud-friendly alternative.
-	err = gwMux.HandlePath("GET", "/v1/instances/{instance_id}/table/{table_name}/export/{format}", auth.GatewayMiddleware(s.aud, s.ExportTable))
-	if err != nil {
-		panic(err)
-	}
-
 	// Call callback to register additional paths
 	// NOTE: This is so ugly, but not worth refactoring it properly right now.
 	httpMux := http.NewServeMux()

web-common/src/features/models/workspace/ModelWorkspaceCTAs.svelte

@@ -8,8 +8,9 @@
   import CaretDownIcon from "@rilldata/web-common/components/icons/CaretDownIcon.svelte";
   import Forward from "@rilldata/web-common/components/icons/Forward.svelte";
   import { Menu, MenuItem } from "@rilldata/web-common/components/menu";
+  import { createExportTableMutation } from "@rilldata/web-common/features/models/workspace/export-table";
   import type { V1Resource } from "@rilldata/web-common/runtime-client";
-  import { RuntimeUrl } from "@rilldata/web-local/lib/application-state-stores/initialize-node-store-contexts";
+  import { V1ExportFormat } from "@rilldata/web-common/runtime-client";
 
   import ResponsiveButtonText from "@rilldata/web-common/components/panel/ResponsiveButtonText.svelte";
   import Tooltip from "@rilldata/web-common/components/tooltip/Tooltip.svelte";
@@ -24,11 +25,16 @@
 
   export let collapse = false;
 
-  const onExport = async (exportExtension: "csv" | "parquet") => {
-    // TODO: how do we handle errors ?
-    window.open(
-      `${RuntimeUrl}/v1/instances/${$runtime.instanceId}/table/${modelName}/export/${exportExtension}`
-    );
+  const exportModelMutation = createExportTableMutation();
+
+  const onExport = async (format: V1ExportFormat) => {
+    return $exportModelMutation.mutateAsync({
+      data: {
+        instanceId: $runtime.instanceId,
+        format,
+        tableName: modelName,
+      },
+    });
   };
 </script>
 
@@ -66,19 +72,27 @@
       <MenuItem
         on:select={() => {
           toggleFloatingElement();
-          onExport("parquet");
+          onExport(V1ExportFormat.EXPORT_FORMAT_PARQUET);
         }}
       >
         Export as Parquet
       </MenuItem>
       <MenuItem
         on:select={() => {
           toggleFloatingElement();
-          onExport("csv");
+          onExport(V1ExportFormat.EXPORT_FORMAT_CSV);
         }}
       >
         Export as CSV
       </MenuItem>
+      <MenuItem
+        on:select={() => {
+          toggleFloatingElement();
+          onExport(V1ExportFormat.EXPORT_FORMAT_XLSX);
+        }}
+      >
+        Export as XLSX
+      </MenuItem>
     </Menu>
   </WithTogglableFloatingElement>
   <TooltipContent slot="tooltip-content">

web-common/src/features/models/workspace/export-table.ts

@@ -0,0 +1,60 @@
+import {
+  createQueryServiceExport,
+  RpcStatus,
+  V1ExportFormat,
+} from "@rilldata/web-common/runtime-client";
+import { runtime } from "@rilldata/web-common/runtime-client/runtime-store";
+import { createMutation, CreateMutationOptions } from "@tanstack/svelte-query";
+import type { MutationFunction } from "@tanstack/svelte-query";
+import { get } from "svelte/store";
+
+export type ExportTableRequest = {
+  instanceId: string;
+  format: V1ExportFormat;
+  tableName: string;
+};
+
+export function createExportTableMutation<
+  TError = { response: { data: RpcStatus } },
+  TContext = unknown
+>(options?: {
+  mutation?: CreateMutationOptions<
+    Awaited<Promise<void>>,
+    TError,
+    { data: ExportTableRequest },
+    TContext
+  >;
+}) {
+  const { mutation: mutationOptions } = options ?? {};
+  const exporter = createQueryServiceExport();
+
+  const mutationFn: MutationFunction<
+    Awaited<Promise<void>>,
+    { data: ExportTableRequest }
+  > = async (props) => {
+    const { data } = props ?? {};
+    if (!data.instanceId) throw new Error("Missing instanceId");
+
+    const exportResp = await get(exporter).mutateAsync({
+      instanceId: data.instanceId,
+      data: {
+        format: data.format,
+        query: {
+          tableRowsRequest: {
+            instanceId: data.instanceId,
+            tableName: data.tableName,
+          },
+        },
+      },
+    });
+    const downloadUrl = `${get(runtime).host}${exportResp.downloadUrlPath}`;
+    window.open(downloadUrl, "_self");
+  };
+
+  return createMutation<
+    Awaited<Promise<void>>,
+    TError,
+    { data: ExportTableRequest },
+    TContext
+  >(mutationFn, mutationOptions);
+}