diff --git a/runtime/queries/table_head.go b/runtime/queries/table_head.go index 79b8006a622..bb1453f1ad3 100644 --- a/runtime/queries/table_head.go +++ b/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 } diff --git a/runtime/server/downloads.go b/runtime/server/downloads.go index c7cd48fc485..c6abecc3b4b 100644 --- a/runtime/server/downloads.go +++ b/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 diff --git a/runtime/server/export.go b/runtime/server/export.go deleted file mode 100644 index b8cb09e55c5..00000000000 --- a/runtime/server/export.go +++ /dev/null @@ -1,74 +0,0 @@ -package server - -import ( - "fmt" - "io" - "net/http" - "os" - "path" - - "github.com/rilldata/rill/runtime/drivers" - "github.com/rilldata/rill/runtime/server/auth" -) - -// ExportTable exports a table or view as a flat file and triggers a HTTP download of it. -// It's mounted as a REST API only, and is not available over gRPC. -// -// TODO: This is a temporary hack that only supports DuckDB. -// We should add a generic workflow for data export that also supports possibly very large tables. -func (s *Server) ExportTable(w http.ResponseWriter, req *http.Request, pathParams map[string]string) { - if !auth.GetClaims(req.Context()).CanInstance(pathParams["instance_id"], auth.ReadOLAP) { - http.Error(w, "action not allowed", http.StatusUnauthorized) - return - } - - var exportString string - switch pathParams["format"] { - case "csv": - exportString = "FORMAT CSV, HEADER" - case "parquet": - exportString = "FORMAT PARQUET" - default: - http.Error(w, fmt.Sprintf("unknown format: %s", pathParams), http.StatusBadRequest) - } - - if pathParams["instance_id"] == "" || pathParams["table_name"] == "" { - http.Error(w, "missing params", http.StatusBadRequest) - return - } - - fileName := fmt.Sprintf("%s.%s", pathParams["table_name"], pathParams["format"]) - filePath := path.Join(os.TempDir(), fileName) - defer os.Remove(filePath) - - // select * from the table and write to the temp file (DuckDB only) - olap, release, err := s.runtime.OLAP(req.Context(), pathParams["instance_id"]) - if err != nil { - http.Error(w, err.Error(), http.StatusBadRequest) - return - } - defer release() - err = olap.Exec(req.Context(), &drivers.Statement{ - Query: fmt.Sprintf("COPY (SELECT * FROM %q) TO '%s' (%s)", pathParams["table_name"], filePath, exportString), - }) - if err != nil { - http.Error(w, err.Error(), http.StatusBadRequest) - return - } - - // set the header to trigger download - w.Header().Set("Content-Disposition", fmt.Sprintf("attachment; filename=%s", fileName)) - w.Header().Set("Content-Type", req.Header.Get("Content-Type")) - - // read and stream the file - file, err := os.Open(filePath) - if err != nil { - http.Error(w, err.Error(), http.StatusBadRequest) - return - } - _, err = io.Copy(w, file) - if err != nil { - http.Error(w, err.Error(), http.StatusBadRequest) - return - } -} diff --git a/runtime/server/server.go b/runtime/server/server.go index 7a3ffb0961f..05257885074 100644 --- a/runtime/server/server.go +++ b/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() diff --git a/web-common/src/features/models/workspace/ModelWorkspaceCTAs.svelte b/web-common/src/features/models/workspace/ModelWorkspaceCTAs.svelte index 2653aa8b429..0f467b4ac71 100644 --- a/web-common/src/features/models/workspace/ModelWorkspaceCTAs.svelte +++ b/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, + }, + }); }; @@ -66,7 +72,7 @@ { toggleFloatingElement(); - onExport("parquet"); + onExport(V1ExportFormat.EXPORT_FORMAT_PARQUET); }} > Export as Parquet @@ -74,11 +80,19 @@ { toggleFloatingElement(); - onExport("csv"); + onExport(V1ExportFormat.EXPORT_FORMAT_CSV); }} > Export as CSV + { + toggleFloatingElement(); + onExport(V1ExportFormat.EXPORT_FORMAT_XLSX); + }} + > + Export as XLSX + diff --git a/web-common/src/features/models/workspace/export-table.ts b/web-common/src/features/models/workspace/export-table.ts new file mode 100644 index 00000000000..363bba9e6d4 --- /dev/null +++ b/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>, + TError, + { data: ExportTableRequest }, + TContext + >; +}) { + const { mutation: mutationOptions } = options ?? {}; + const exporter = createQueryServiceExport(); + + const mutationFn: MutationFunction< + Awaited>, + { 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>, + TError, + { data: ExportTableRequest }, + TContext + >(mutationFn, mutationOptions); +}