From 1cafb94a62be084e976e9c3f6755a5def53e36dc Mon Sep 17 00:00:00 2001 From: xiongyuanwen-byted Date: Mon, 8 Jun 2026 12:58:11 +0800 Subject: [PATCH] refactor(sheets): reuse the drive export core in +workbook-export Replace +workbook-export's parallel export-task implementation with the shared drive ExportParams/RunExport core (pinned to type=sheet). Drops ~90 lines of duplicated poll/download code; +workbook-export now inherits drive's ctx cancellation, resume-on-timeout, filename sanitize/overwrite, and the full set of export status labels. The output contract aligns with drive's (adds ready/downloaded/doc_type; saved_path preserved). Also normalize an empty drive --output-dir to "." so drive +export behavior is unchanged, and fix the sheets export e2e to call +workbook-export instead of a nonexistent +export. --- shortcuts/drive/drive_export.go | 460 ++++++++++-------- shortcuts/drive/drive_export_test.go | 66 +++ shortcuts/sheets/lark_sheet_workbook.go | 198 ++------ .../sheets/lark_sheet_workbook_export_test.go | 72 +++ shortcuts/sheets/lark_sheet_workbook_test.go | 28 +- .../sheets/sheets_crud_workflow_test.go | 4 +- 6 files changed, 453 insertions(+), 375 deletions(-) create mode 100644 shortcuts/sheets/lark_sheet_workbook_export_test.go diff --git a/shortcuts/drive/drive_export.go b/shortcuts/drive/drive_export.go index 86e7c99a5..56f3a6e03 100644 --- a/shortcuts/drive/drive_export.go +++ b/shortcuts/drive/drive_export.go @@ -39,230 +39,296 @@ var DriveExport = common.Shortcut{ {Name: "overwrite", Type: "bool", Desc: "overwrite existing output file"}, }, Validate: func(ctx context.Context, runtime *common.RuntimeContext) error { - return validateDriveExportSpec(driveExportSpec{ - Token: runtime.Str("token"), - DocType: runtime.Str("doc-type"), - FileExtension: runtime.Str("file-extension"), - SubID: runtime.Str("sub-id"), - }) + return ValidateExport(exportParamsFromFlags(runtime)) }, DryRun: func(ctx context.Context, runtime *common.RuntimeContext) *common.DryRunAPI { - spec := driveExportSpec{ - Token: runtime.Str("token"), - DocType: runtime.Str("doc-type"), - FileExtension: runtime.Str("file-extension"), - SubID: runtime.Str("sub-id"), - } - // Markdown export is a special case: docx markdown comes from the V2 - // docs_ai fetch API directly instead of the Drive export task API. - if spec.FileExtension == "markdown" { - apiPath := fmt.Sprintf("/open-apis/docs_ai/v1/documents/%s/fetch", validate.EncodePathSegment(spec.Token)) - dr := common.NewDryRunAPI(). - Desc("2-step orchestration: fetch docx markdown -> write local file"). - POST(apiPath). - Body(map[string]interface{}{ - "format": "markdown", - }). - Set("output_dir", runtime.Str("output-dir")) - if name := strings.TrimSpace(runtime.Str("file-name")); name != "" { - dr.Set("file_name", ensureExportFileExtension(sanitizeExportFileName(name, spec.Token), spec.FileExtension)) - } - return dr - } + return PlanExportDryRun(runtime, exportParamsFromFlags(runtime)) + }, + Execute: func(ctx context.Context, runtime *common.RuntimeContext) error { + return RunExport(ctx, runtime, exportParamsFromFlags(runtime)) + }, +} - body := map[string]interface{}{ - "token": spec.Token, - "type": spec.DocType, - "file_extension": spec.FileExtension, - } - if strings.TrimSpace(spec.SubID) != "" { - body["sub_id"] = spec.SubID - } +// ExportParams holds the user-facing inputs for an export flow, decoupled from +// cobra flags so other command groups (e.g. sheets +workbook-export) can reuse +// the drive export implementation. An empty OutputDir means "create the export +// task and poll, but do not download" — callers that only need the ready file +// token / status get it back without writing a local file. +type ExportParams struct { + Token string + DocType string + FileExtension string + SubID string + OutputDir string + FileName string + Overwrite bool +} +func (p ExportParams) spec() driveExportSpec { + return driveExportSpec{ + Token: p.Token, + DocType: p.DocType, + FileExtension: p.FileExtension, + SubID: p.SubID, + } +} + +// exportParamsFromFlags reads the standard drive +export flag set. +func exportParamsFromFlags(runtime *common.RuntimeContext) ExportParams { + // drive +export always downloads; an empty --output-dir historically means + // the current directory (saveContentToOutputDir maps "" -> "."), so normalize + // it here to keep behavior identical and stay off the export-only ("" => skip + // download) path that only sheets +workbook-export uses. + outputDir := runtime.Str("output-dir") + if outputDir == "" { + outputDir = "." + } + return ExportParams{ + Token: runtime.Str("token"), + DocType: runtime.Str("doc-type"), + FileExtension: runtime.Str("file-extension"), + SubID: runtime.Str("sub-id"), + OutputDir: outputDir, + FileName: strings.TrimSpace(runtime.Str("file-name")), + Overwrite: runtime.Bool("overwrite"), + } +} + +// ValidateExport runs the CLI-level export constraint checks. +func ValidateExport(p ExportParams) error { + return validateDriveExportSpec(p.spec()) +} + +// PlanExportDryRun builds the dry-run plan for an export without performing I/O. +func PlanExportDryRun(runtime *common.RuntimeContext, p ExportParams) *common.DryRunAPI { + spec := p.spec() + // Markdown export is a special case: docx markdown comes from the V2 + // docs_ai fetch API directly instead of the Drive export task API. + if spec.FileExtension == "markdown" { + apiPath := fmt.Sprintf("/open-apis/docs_ai/v1/documents/%s/fetch", validate.EncodePathSegment(spec.Token)) dr := common.NewDryRunAPI(). - Desc("3-step orchestration: create export task -> limited polling -> download file"). - POST("/open-apis/drive/v1/export_tasks"). - Body(body). - Set("output_dir", runtime.Str("output-dir")) - if name := strings.TrimSpace(runtime.Str("file-name")); name != "" { + Desc("2-step orchestration: fetch docx markdown -> write local file"). + POST(apiPath). + Body(map[string]interface{}{ + "format": "markdown", + }). + Set("output_dir", p.OutputDir) + if name := strings.TrimSpace(p.FileName); name != "" { dr.Set("file_name", ensureExportFileExtension(sanitizeExportFileName(name, spec.Token), spec.FileExtension)) } return dr - }, - Execute: func(ctx context.Context, runtime *common.RuntimeContext) error { - spec := driveExportSpec{ - Token: runtime.Str("token"), - DocType: runtime.Str("doc-type"), - FileExtension: runtime.Str("file-extension"), - SubID: runtime.Str("sub-id"), + } + + body := map[string]interface{}{ + "token": spec.Token, + "type": spec.DocType, + "file_extension": spec.FileExtension, + } + if strings.TrimSpace(spec.SubID) != "" { + body["sub_id"] = spec.SubID + } + + dr := common.NewDryRunAPI(). + Desc("3-step orchestration: create export task -> limited polling -> download file"). + POST("/open-apis/drive/v1/export_tasks"). + Body(body). + Set("output_dir", p.OutputDir) + if name := strings.TrimSpace(p.FileName); name != "" { + dr.Set("file_name", ensureExportFileExtension(sanitizeExportFileName(name, spec.Token), spec.FileExtension)) + } + return dr +} + +// RunExport drives create export task -> bounded poll -> optional download. It +// is the shared core behind both drive +export and sheets +workbook-export. An +// empty p.OutputDir skips the download step and returns the ready file token. +func RunExport(ctx context.Context, runtime *common.RuntimeContext, p ExportParams) error { + spec := p.spec() + outputDir := p.OutputDir + preferredFileName := strings.TrimSpace(p.FileName) + overwrite := p.Overwrite + + // Markdown export bypasses the async export task and writes the fetched + // markdown content directly to disk. Uses the V2 docs_ai fetch API for + // higher-quality Lark-flavored Markdown output. + if spec.FileExtension == "markdown" { + fmt.Fprintf(runtime.IO().ErrOut, "Exporting docx as markdown: %s\n", common.MaskToken(spec.Token)) + apiPath := fmt.Sprintf("/open-apis/docs_ai/v1/documents/%s/fetch", validate.EncodePathSegment(spec.Token)) + data, err := runtime.CallAPITyped( + "POST", + apiPath, + nil, + map[string]interface{}{ + "format": "markdown", + }, + ) + if err != nil { + return err } - outputDir := runtime.Str("output-dir") - preferredFileName := strings.TrimSpace(runtime.Str("file-name")) - overwrite := runtime.Bool("overwrite") - // Markdown export bypasses the async export task and writes the fetched - // markdown content directly to disk. Uses the V2 docs_ai fetch API for - // higher-quality Lark-flavored Markdown output. - if spec.FileExtension == "markdown" { - fmt.Fprintf(runtime.IO().ErrOut, "Exporting docx as markdown: %s\n", common.MaskToken(spec.Token)) - apiPath := fmt.Sprintf("/open-apis/docs_ai/v1/documents/%s/fetch", validate.EncodePathSegment(spec.Token)) - data, err := runtime.CallAPITyped( - "POST", - apiPath, - nil, - map[string]interface{}{ - "format": "markdown", - }, - ) + // Extract content from the V2 response: data.document.content + doc, ok := data["document"].(map[string]interface{}) + if !ok { + return errs.NewInternalError(errs.SubtypeInvalidResponse, "invalid markdown fetch response: missing document object") + } + content, ok := doc["content"].(string) + if !ok { + return errs.NewInternalError(errs.SubtypeInvalidResponse, "invalid markdown fetch response: missing document.content") + } + + fileName := preferredFileName + if fileName == "" { + // Prefer the remote title for the exported file name, but still fall + // back to the token if metadata is empty. + title, err := common.FetchDriveMetaTitle(runtime, spec.Token, spec.DocType) if err != nil { - return err + fmt.Fprintf(runtime.IO().ErrOut, "Title lookup failed, using token as filename: %v\n", err) + title = spec.Token } + fileName = title + } + fileName = ensureExportFileExtension(sanitizeExportFileName(fileName, spec.Token), spec.FileExtension) + savedPath, err := saveContentToOutputDir(runtime.FileIO(), outputDir, fileName, []byte(content), overwrite) + if err != nil { + return err + } - // Extract content from the V2 response: data.document.content - doc, ok := data["document"].(map[string]interface{}) - if !ok { - return errs.NewInternalError(errs.SubtypeInvalidResponse, "invalid markdown fetch response: missing document object") + runtime.Out(map[string]interface{}{ + "token": spec.Token, + "doc_type": spec.DocType, + "file_extension": spec.FileExtension, + "file_name": filepath.Base(savedPath), + "saved_path": savedPath, + "size_bytes": len(content), + }, nil) + return nil + } + + ticket, err := createDriveExportTask(runtime, spec) + if err != nil { + return err + } + fmt.Fprintf(runtime.IO().ErrOut, "Created export task: %s\n", ticket) + + var lastStatus driveExportStatus + var lastPollErr error + hasObservedStatus := false + // Keep the command responsive by polling for a bounded window. If the task + // is still running after that, return a resume command instead of blocking. + for attempt := 1; attempt <= driveExportPollAttempts; attempt++ { + if attempt > 1 { + select { + case <-ctx.Done(): + return ctx.Err() + case <-time.After(driveExportPollInterval): } - content, ok := doc["content"].(string) - if !ok { - return errs.NewInternalError(errs.SubtypeInvalidResponse, "invalid markdown fetch response: missing document.content") + } + if err := ctx.Err(); err != nil { + return err + } + + status, err := getDriveExportStatus(runtime, spec.Token, ticket) + if err != nil { + // Treat polling failures as transient so short-lived backend hiccups + // do not immediately fail an otherwise healthy export task. + lastPollErr = err + fmt.Fprintf(runtime.IO().ErrOut, "Export status attempt %d/%d failed: %v\n", attempt, driveExportPollAttempts, err) + continue + } + lastStatus = status + hasObservedStatus = true + + if status.Ready() { + fmt.Fprintf(runtime.IO().ErrOut, "Export task completed: %s\n", common.MaskToken(status.FileToken)) + + // Export-only mode: caller wants the ready file token / metadata but + // no local download (e.g. sheets +workbook-export without an output + // path). Skip the download and return the status envelope. + if strings.TrimSpace(outputDir) == "" { + runtime.Out(map[string]interface{}{ + "ticket": ticket, + "token": spec.Token, + "doc_type": spec.DocType, + "file_extension": spec.FileExtension, + "file_token": status.FileToken, + "file_name": status.FileName, + "file_size": status.FileSize, + "ready": true, + "downloaded": false, + }, nil) + return nil } fileName := preferredFileName if fileName == "" { - // Prefer the remote title for the exported file name, but still fall - // back to the token if metadata is empty. - title, err := common.FetchDriveMetaTitle(runtime, spec.Token, spec.DocType) - if err != nil { - fmt.Fprintf(runtime.IO().ErrOut, "Title lookup failed, using token as filename: %v\n", err) - title = spec.Token - } - fileName = title + fileName = status.FileName } fileName = ensureExportFileExtension(sanitizeExportFileName(fileName, spec.Token), spec.FileExtension) - savedPath, err := saveContentToOutputDir(runtime.FileIO(), outputDir, fileName, []byte(content), overwrite) + out, err := downloadDriveExportFile(ctx, runtime, status.FileToken, outputDir, fileName, overwrite) if err != nil { - return err + recoveryCommand := driveExportDownloadCommand(status.FileToken, fileName, outputDir, overwrite) + hint := fmt.Sprintf( + "the export artifact is already ready (ticket=%s, file_token=%s)\nretry download with: %s", + ticket, + status.FileToken, + recoveryCommand, + ) + return appendDriveExportRecoveryHint(err, hint) } - - runtime.Out(map[string]interface{}{ - "token": spec.Token, - "doc_type": spec.DocType, - "file_extension": spec.FileExtension, - "file_name": filepath.Base(savedPath), - "saved_path": savedPath, - "size_bytes": len(content), - }, nil) + out["ticket"] = ticket + out["doc_type"] = spec.DocType + out["file_extension"] = spec.FileExtension + runtime.Out(out, nil) return nil } - ticket, err := createDriveExportTask(runtime, spec) - if err != nil { - return err - } - fmt.Fprintf(runtime.IO().ErrOut, "Created export task: %s\n", ticket) - - var lastStatus driveExportStatus - var lastPollErr error - hasObservedStatus := false - // Keep the command responsive by polling for a bounded window. If the task - // is still running after that, return a resume command instead of blocking. - for attempt := 1; attempt <= driveExportPollAttempts; attempt++ { - if attempt > 1 { - select { - case <-ctx.Done(): - return ctx.Err() - case <-time.After(driveExportPollInterval): - } + if status.Failed() { + msg := strings.TrimSpace(status.JobErrorMsg) + if msg == "" { + msg = status.StatusLabel() } - if err := ctx.Err(); err != nil { - return err - } - - status, err := getDriveExportStatus(runtime, spec.Token, ticket) - if err != nil { - // Treat polling failures as transient so short-lived backend hiccups - // do not immediately fail an otherwise healthy export task. - lastPollErr = err - fmt.Fprintf(runtime.IO().ErrOut, "Export status attempt %d/%d failed: %v\n", attempt, driveExportPollAttempts, err) - continue - } - lastStatus = status - hasObservedStatus = true - - if status.Ready() { - fmt.Fprintf(runtime.IO().ErrOut, "Export task completed: %s\n", common.MaskToken(status.FileToken)) - fileName := preferredFileName - if fileName == "" { - fileName = status.FileName - } - fileName = ensureExportFileExtension(sanitizeExportFileName(fileName, spec.Token), spec.FileExtension) - out, err := downloadDriveExportFile(ctx, runtime, status.FileToken, outputDir, fileName, overwrite) - if err != nil { - recoveryCommand := driveExportDownloadCommand(status.FileToken, fileName, outputDir, overwrite) - hint := fmt.Sprintf( - "the export artifact is already ready (ticket=%s, file_token=%s)\nretry download with: %s", - ticket, - status.FileToken, - recoveryCommand, - ) - return appendDriveExportRecoveryHint(err, hint) - } - out["ticket"] = ticket - out["doc_type"] = spec.DocType - out["file_extension"] = spec.FileExtension - runtime.Out(out, nil) - return nil - } - - if status.Failed() { - msg := strings.TrimSpace(status.JobErrorMsg) - if msg == "" { - msg = status.StatusLabel() - } - return errs.NewAPIError(errs.SubtypeServerError, "export task failed: %s (ticket=%s)", msg, ticket) - } - - fmt.Fprintf(runtime.IO().ErrOut, "Export status %d/%d: %s\n", attempt, driveExportPollAttempts, status.StatusLabel()) + return errs.NewAPIError(errs.SubtypeServerError, "export task failed: %s (ticket=%s)", msg, ticket) } - nextCommand := driveExportTaskResultCommand(ticket, spec.Token) - if !hasObservedStatus && lastPollErr != nil { - hint := fmt.Sprintf( - "the export task was created but every status poll failed (ticket=%s)\nretry status lookup with: %s", - ticket, - nextCommand, - ) - return appendDriveExportRecoveryHint(lastPollErr, hint) - } + fmt.Fprintf(runtime.IO().ErrOut, "Export status %d/%d: %s\n", attempt, driveExportPollAttempts, status.StatusLabel()) + } - failed := false - var jobStatus interface{} - jobStatusLabel := "unknown" - if hasObservedStatus { - failed = lastStatus.Failed() - jobStatus = lastStatus.JobStatus - jobStatusLabel = lastStatus.StatusLabel() - } - // Return the last observed status so callers can resume from a known task - // state instead of losing all progress information on timeout. - result := map[string]interface{}{ - "ticket": ticket, - "token": spec.Token, - "doc_type": spec.DocType, - "file_extension": spec.FileExtension, - "ready": false, - "failed": failed, - "job_status": jobStatus, - "job_status_label": jobStatusLabel, - "timed_out": true, - "next_command": nextCommand, - } - if preferredFileName != "" { - result["file_name"] = ensureExportFileExtension(sanitizeExportFileName(preferredFileName, spec.Token), spec.FileExtension) - } - runtime.Out(result, nil) - fmt.Fprintf(runtime.IO().ErrOut, "Export task is still in progress. Continue with: %s\n", nextCommand) - return nil - }, + nextCommand := driveExportTaskResultCommand(ticket, spec.Token) + if !hasObservedStatus && lastPollErr != nil { + hint := fmt.Sprintf( + "the export task was created but every status poll failed (ticket=%s)\nretry status lookup with: %s", + ticket, + nextCommand, + ) + return appendDriveExportRecoveryHint(lastPollErr, hint) + } + + failed := false + var jobStatus interface{} + jobStatusLabel := "unknown" + if hasObservedStatus { + failed = lastStatus.Failed() + jobStatus = lastStatus.JobStatus + jobStatusLabel = lastStatus.StatusLabel() + } + // Return the last observed status so callers can resume from a known task + // state instead of losing all progress information on timeout. + result := map[string]interface{}{ + "ticket": ticket, + "token": spec.Token, + "doc_type": spec.DocType, + "file_extension": spec.FileExtension, + "ready": false, + "failed": failed, + "job_status": jobStatus, + "job_status_label": jobStatusLabel, + "timed_out": true, + "next_command": nextCommand, + } + if preferredFileName != "" { + result["file_name"] = ensureExportFileExtension(sanitizeExportFileName(preferredFileName, spec.Token), spec.FileExtension) + } + runtime.Out(result, nil) + fmt.Fprintf(runtime.IO().ErrOut, "Export task is still in progress. Continue with: %s\n", nextCommand) + return nil } diff --git a/shortcuts/drive/drive_export_test.go b/shortcuts/drive/drive_export_test.go index 619c9d43d..f35dbfab3 100644 --- a/shortcuts/drive/drive_export_test.go +++ b/shortcuts/drive/drive_export_test.go @@ -488,6 +488,72 @@ func TestDriveExportAsyncSuccess(t *testing.T) { } } +// TestDriveExportEmptyOutputDirDownloadsToCwd guards the export refactor: an +// explicit empty --output-dir must still download to the current directory +// (normalized to "."), not trigger the export-only no-download path that the +// shared RunExport core uses for sheets +workbook-export. +func TestDriveExportEmptyOutputDirDownloadsToCwd(t *testing.T) { + f, stdout, _, reg := cmdutil.TestFactory(t, driveTestConfig()) + reg.Register(&httpmock.Stub{ + Method: "POST", + URL: "/open-apis/drive/v1/export_tasks", + Body: map[string]interface{}{"code": 0, "data": map[string]interface{}{"ticket": "tk_e"}}, + }) + reg.Register(&httpmock.Stub{ + Method: "GET", + URL: "/open-apis/drive/v1/export_tasks/tk_e", + Body: map[string]interface{}{"code": 0, "data": map[string]interface{}{ + "result": map[string]interface{}{ + "job_status": 0, "file_token": "box_e", "file_name": "report", + "file_extension": "pdf", "type": "docx", "file_size": 3, + }, + }}, + }) + reg.Register(&httpmock.Stub{ + Method: "GET", + URL: "/open-apis/drive/v1/export_tasks/file/box_e/download", + Status: 200, + RawBody: []byte("pdf"), + Headers: http.Header{ + "Content-Type": []string{"application/pdf"}, + "Content-Disposition": []string{`attachment; filename="report.pdf"`}, + }, + }) + + tmpDir := t.TempDir() + withDriveWorkingDir(t, tmpDir) + + prevAttempts, prevInterval := driveExportPollAttempts, driveExportPollInterval + driveExportPollAttempts, driveExportPollInterval = 1, 0 + t.Cleanup(func() { + driveExportPollAttempts, driveExportPollInterval = prevAttempts, prevInterval + }) + + err := mountAndRunDrive(t, DriveExport, []string{ + "+export", + "--token", "docx123", + "--doc-type", "docx", + "--file-extension", "pdf", + "--output-dir", "", + "--as", "bot", + }, f, stdout) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + // Empty --output-dir must still write to cwd, not skip the download. + data, err := os.ReadFile(filepath.Join(tmpDir, "report.pdf")) + if err != nil { + t.Fatalf("empty --output-dir should still download to cwd: %v", err) + } + if string(data) != "pdf" { + t.Fatalf("downloaded content = %q", string(data)) + } + if strings.Contains(stdout.String(), `"downloaded": false`) { + t.Fatalf("export-only path must not trigger for drive +export: %s", stdout.String()) + } +} + func TestDriveExportAsyncUsesProvidedFileName(t *testing.T) { f, stdout, _, reg := cmdutil.TestFactory(t, driveTestConfig()) reg.Register(&httpmock.Stub{ diff --git a/shortcuts/sheets/lark_sheet_workbook.go b/shortcuts/sheets/lark_sheet_workbook.go index e853be1a7..509f93021 100644 --- a/shortcuts/sheets/lark_sheet_workbook.go +++ b/shortcuts/sheets/lark_sheet_workbook.go @@ -6,18 +6,12 @@ package sheets import ( "context" "fmt" - "net/http" "path/filepath" "strings" - "time" - - larkcore "github.com/larksuite/oapi-sdk-go/v3/core" "github.com/larksuite/cli/extension/fileio" - "github.com/larksuite/cli/internal/client" "github.com/larksuite/cli/internal/output" "github.com/larksuite/cli/internal/util" - "github.com/larksuite/cli/internal/validate" "github.com/larksuite/cli/shortcuts/common" "github.com/larksuite/cli/shortcuts/drive" ) @@ -845,178 +839,62 @@ var WorkbookExport = common.Shortcut{ return nil }, DryRun: func(ctx context.Context, runtime *common.RuntimeContext) *common.DryRunAPI { - token, _ := resolveSpreadsheetToken(runtime) - ext := runtime.Str("file-extension") - if ext == "" { - ext = "xlsx" - } - body := map[string]interface{}{ - "token": token, - "type": "sheet", - "file_extension": ext, - } - if sid := strings.TrimSpace(runtime.Str("sheet-id")); sid != "" { - body["sub_id"] = sid - } - dry := common.NewDryRunAPI(). - POST("/open-apis/drive/v1/export_tasks"). - Desc("create export task"). - Body(body). - GET("/open-apis/drive/v1/export_tasks/"). - Desc("poll task status"). - Params(map[string]interface{}{"token": token}) - if strings.TrimSpace(runtime.Str("output-path")) != "" { - dry.GET("/open-apis/drive/v1/export_tasks/file//download"). - Desc("download exported file") - } - return dry + p, _ := workbookExportParams(runtime) + p.OutputDir = strings.TrimSpace(runtime.Str("output-path")) + return drive.PlanExportDryRun(runtime, p) }, Execute: func(ctx context.Context, runtime *common.RuntimeContext) error { - token, err := resolveSpreadsheetToken(runtime) + p, err := workbookExportParams(runtime) if err != nil { return err } - ext := runtime.Str("file-extension") - if ext == "" { - ext = "xlsx" - } - body := map[string]interface{}{ - "token": token, - "type": "sheet", - "file_extension": ext, - } - if sid := strings.TrimSpace(runtime.Str("sheet-id")); sid != "" { - body["sub_id"] = sid - } - taskData, err := runtime.CallAPI("POST", "/open-apis/drive/v1/export_tasks", nil, body) - if err != nil { - return err - } - ticket := common.GetString(taskData, "ticket") - if ticket == "" { - return output.Errorf(output.ExitAPI, "api_error", "export task created but ticket missing") - } - - result := map[string]interface{}{ - "ticket": ticket, - "file_extension": ext, - } - - // Poll up to ~30s for completion. - var fileToken, fileName string - for attempt := 0; attempt < 15; attempt++ { - status, err := pollExportTask(runtime, token, ticket) - if err != nil { - return err - } - switch status.JobStatus { - case 0: // success - fileToken = status.FileToken - fileName = status.FileName - result["file_token"] = fileToken - result["file_name"] = fileName - result["file_size"] = status.FileSize - attempt = 999 // break outer loop - case 1, 2: // pending / in progress - time.Sleep(2 * time.Second) - continue - default: // any non-zero status outside the in-progress window is a failure - if status.JobErrorMsg != "" { - return output.Errorf(output.ExitAPI, "api_error", "export task %s failed: %s", ticket, status.JobErrorMsg) - } - return output.Errorf(output.ExitAPI, "api_error", "export task %s failed with job_status=%d", ticket, status.JobStatus) - } - } - if fileToken == "" { - result["status"] = "polling_timeout" - runtime.Out(result, nil) - return nil - } - - outPath := strings.TrimSpace(runtime.Str("output-path")) - if outPath == "" { - runtime.Out(result, nil) - return nil - } - - saved, err := downloadExportFile(ctx, runtime, fileToken, outPath, fileName) - if err != nil { - return err - } - result["saved_path"] = saved - runtime.Out(result, nil) - return nil + applyWorkbookOutputPath(&p, runtime.FileIO(), runtime.Str("output-path")) + return drive.RunExport(ctx, runtime, p) }, Tips: []string{ - "Polls up to ~30s (15 × 2s). For very large workbooks rerun and pass --output-path to capture the file once status flips to success.", + "Polls for a bounded window; if the export is still running it returns a resume reference instead of blocking. Pass --output-path to download the file once ready (omit it to only create the export task and get the file token back).", }, } -type exportTaskStatus struct { - JobStatus int - JobErrorMsg string - FileToken string - FileName string - FileSize int64 - FileExtension string -} - -func pollExportTask(runtime *common.RuntimeContext, token, ticket string) (exportTaskStatus, error) { - data, err := runtime.CallAPI( - "GET", - fmt.Sprintf("/open-apis/drive/v1/export_tasks/%s", validate.EncodePathSegment(ticket)), - map[string]interface{}{"token": token}, - nil, - ) +// workbookExportParams builds the shared drive export request for +// +workbook-export: spreadsheet token + sheet locator, pinned to type=sheet. +// workbook-export has always overwritten the target, so Overwrite is set. The +// --output-path → OutputDir/FileName split (which needs a Stat) is applied +// separately by applyWorkbookOutputPath so Validate/DryRun stay I/O-free. +func workbookExportParams(runtime *common.RuntimeContext) (drive.ExportParams, error) { + token, err := resolveSpreadsheetToken(runtime) if err != nil { - return exportTaskStatus{}, err + return drive.ExportParams{}, err } - result := common.GetMap(data, "result") - if result == nil { - return exportTaskStatus{}, output.Errorf(output.ExitAPI, "api_error", "export task %s: empty result", ticket) + ext := runtime.Str("file-extension") + if ext == "" { + ext = "xlsx" } - js, _ := util.ToFloat64(result["job_status"]) - fs, _ := util.ToFloat64(result["file_size"]) - return exportTaskStatus{ - JobStatus: int(js), - JobErrorMsg: common.GetString(result, "job_error_msg"), - FileToken: common.GetString(result, "file_token"), - FileName: common.GetString(result, "file_name"), - FileSize: int64(fs), - FileExtension: common.GetString(result, "file_extension"), + return drive.ExportParams{ + Token: token, + DocType: "sheet", + FileExtension: ext, + SubID: strings.TrimSpace(runtime.Str("sheet-id")), + Overwrite: true, }, nil } -func downloadExportFile(ctx context.Context, runtime *common.RuntimeContext, fileToken, outPath, preferredName string) (string, error) { - apiResp, err := runtime.DoAPI(&larkcore.ApiReq{ - HttpMethod: http.MethodGet, - ApiPath: fmt.Sprintf("/open-apis/drive/v1/export_tasks/file/%s/download", validate.EncodePathSegment(fileToken)), - }, larkcore.WithFileDownload()) - if err != nil { - return "", output.ErrNetwork("download failed: %s", err) +// applyWorkbookOutputPath maps the single --output-path flag onto the drive +// export OutputDir/FileName pair, preserving the legacy behavior: empty = no +// download (return the ready file token only); an existing directory = download +// into it under the server-provided name; otherwise treat it as a file path and +// split into dir + base name. +func applyWorkbookOutputPath(p *drive.ExportParams, fio fileio.FileIO, outputPath string) { + outputPath = strings.TrimSpace(outputPath) + if outputPath == "" { + return } - if apiResp.StatusCode >= 400 { - return "", output.ErrNetwork("download failed: HTTP %d: %s", apiResp.StatusCode, string(apiResp.RawBody)) + if info, err := fio.Stat(outputPath); err == nil && info.IsDir() { + p.OutputDir = outputPath + return } - target := outPath - if info, statErr := runtime.FileIO().Stat(outPath); statErr == nil && info.IsDir() { - name := strings.TrimSpace(preferredName) - if name == "" { - name = client.ResolveFilename(apiResp) - } - target = filepath.Join(outPath, name) - } - if _, err := runtime.FileIO().Save(target, fileio.SaveOptions{ - ContentType: apiResp.Header.Get("Content-Type"), - ContentLength: int64(len(apiResp.RawBody)), - }, strings.NewReader(string(apiResp.RawBody))); err != nil { - return "", common.WrapSaveErrorByCategory(err, "io") - } - resolved, _ := runtime.FileIO().ResolvePath(target) - if resolved == "" { - resolved = target - } - return resolved, nil + p.OutputDir = filepath.Dir(outputPath) + p.FileName = filepath.Base(outputPath) } // lookupSheetIndex finds a sub-sheet by id or name and returns its canonical diff --git a/shortcuts/sheets/lark_sheet_workbook_export_test.go b/shortcuts/sheets/lark_sheet_workbook_export_test.go new file mode 100644 index 000000000..d192d8406 --- /dev/null +++ b/shortcuts/sheets/lark_sheet_workbook_export_test.go @@ -0,0 +1,72 @@ +// Copyright (c) 2026 Lark Technologies Pte. Ltd. +// SPDX-License-Identifier: MIT + +package sheets + +import ( + "encoding/json" + "strings" + "testing" + + "github.com/larksuite/cli/internal/httpmock" +) + +// TestWorkbookExport_ExecuteExportOnly covers the no-download path: without +// --output-path, +workbook-export delegates to the shared drive export core +// with OutputDir="" so it creates + polls the export task and returns the ready +// file token without writing a local file (downloaded=false). +func TestWorkbookExport_ExecuteExportOnly(t *testing.T) { + stubs := []*httpmock.Stub{ + { + Method: "POST", + URL: "/open-apis/drive/v1/export_tasks", + Body: map[string]interface{}{ + "code": 0, "msg": "ok", + "data": map[string]interface{}{"ticket": "tk_export"}, + }, + }, + { + Method: "GET", + URL: "/open-apis/drive/v1/export_tasks/tk_export", + Body: map[string]interface{}{ + "code": 0, "msg": "ok", + "data": map[string]interface{}{"result": map[string]interface{}{ + "job_status": float64(0), + "file_token": "ftk_xlsx", + "file_name": "report.xlsx", + "file_size": float64(2048), + }}, + }, + }, + } + + out, err := runShortcutWithStubs(t, WorkbookExport, []string{ + "--url", testURL, "--file-extension", "xlsx", "--as", "user", + }, stubs...) + if err != nil { + t.Fatalf("export-only execute failed: %v\n%s", err, out) + } + + idx := strings.Index(out, "{") + if idx < 0 { + t.Fatalf("no JSON envelope:\n%s", out) + } + var env struct { + Data map[string]interface{} `json:"data"` + } + if err := json.Unmarshal([]byte(out[idx:]), &env); err != nil { + t.Fatalf("decode envelope: %v\nraw=%s", err, out) + } + if env.Data["ready"] != true { + t.Errorf("ready = %v, want true", env.Data["ready"]) + } + if env.Data["downloaded"] != false { + t.Errorf("downloaded = %v, want false (no --output-path)", env.Data["downloaded"]) + } + if env.Data["file_token"] != "ftk_xlsx" { + t.Errorf("file_token = %v, want ftk_xlsx", env.Data["file_token"]) + } + if env.Data["doc_type"] != "sheet" { + t.Errorf("doc_type = %v, want sheet", env.Data["doc_type"]) + } +} diff --git a/shortcuts/sheets/lark_sheet_workbook_test.go b/shortcuts/sheets/lark_sheet_workbook_test.go index de089c11b..d5c081cac 100644 --- a/shortcuts/sheets/lark_sheet_workbook_test.go +++ b/shortcuts/sheets/lark_sheet_workbook_test.go @@ -361,21 +361,21 @@ func TestWorkbookCreate_DataValidation(t *testing.T) { } } -// TestWorkbookExport_DryRun checks the 2-or-3 step plan depending on -// --output-path. The order should be: POST → GET (poll) → optional GET -// (download). +// TestWorkbookExport_DryRun verifies the export dry-run now delegates to the +// shared drive export core: a single create-task POST (poll + download are +// described inline rather than as separate api entries). func TestWorkbookExport_DryRun(t *testing.T) { t.Parallel() - t.Run("xlsx without --output-path → 2 steps", func(t *testing.T) { + t.Run("xlsx create-task body pins type=sheet", func(t *testing.T) { t.Parallel() calls := parseDryRunAPI(t, WorkbookExport, []string{"--url", testURL, "--file-extension", "xlsx"}) - if len(calls) != 2 { - t.Fatalf("api calls = %d, want 2 (create + poll)", len(calls)) + if len(calls) != 1 { + t.Fatalf("api calls = %d, want 1 (create export task)", len(calls)) } create := calls[0].(map[string]interface{}) if create["url"] != "/open-apis/drive/v1/export_tasks" { - t.Errorf("first url = %v", create["url"]) + t.Errorf("url = %v", create["url"]) } body, _ := create["body"].(map[string]interface{}) if body["type"] != "sheet" || body["file_extension"] != "xlsx" || body["token"] != testToken { @@ -383,22 +383,18 @@ func TestWorkbookExport_DryRun(t *testing.T) { } }) - t.Run("csv → 3 steps, with sub_id", func(t *testing.T) { + t.Run("csv includes sub_id from --sheet-id", func(t *testing.T) { t.Parallel() calls := parseDryRunAPI(t, WorkbookExport, []string{ "--url", testURL, "--file-extension", "csv", "--sheet-id", "sh1", "--output-path", "/tmp/out.csv", }) - if len(calls) != 3 { - t.Fatalf("api calls = %d, want 3", len(calls)) + if len(calls) != 1 { + t.Fatalf("api calls = %d, want 1", len(calls)) } body, _ := calls[0].(map[string]interface{})["body"].(map[string]interface{}) - if body["sub_id"] != "sh1" { - t.Errorf("csv export missing sub_id: %#v", body) - } - dl := calls[2].(map[string]interface{}) - if !strings.Contains(dl["url"].(string), "/export_tasks/file/") { - t.Errorf("download url = %v", dl["url"]) + if body["type"] != "sheet" || body["sub_id"] != "sh1" { + t.Errorf("csv export body = %#v (want type=sheet, sub_id=sh1)", body) } }) diff --git a/tests/cli_e2e/sheets/sheets_crud_workflow_test.go b/tests/cli_e2e/sheets/sheets_crud_workflow_test.go index 5c8fff1d4..beaf4d799 100644 --- a/tests/cli_e2e/sheets/sheets_crud_workflow_test.go +++ b/tests/cli_e2e/sheets/sheets_crud_workflow_test.go @@ -143,14 +143,14 @@ func TestSheets_CRUDE2EWorkflow(t *testing.T) { assert.True(t, len(matchedCells.Array()) > 0, "should find at least one cell containing 'Alice'") }) - t.Run("export spreadsheet with +export as bot", func(t *testing.T) { + t.Run("export spreadsheet with +workbook-export as bot", func(t *testing.T) { require.NotEmpty(t, spreadsheetToken, "spreadsheet token is required") outputDir := t.TempDir() outputPath := filepath.Join(outputDir, "export.xlsx") result, err := clie2e.RunCmd(ctx, clie2e.Request{ Args: []string{ - "sheets", "+export", + "sheets", "+workbook-export", "--spreadsheet-token", spreadsheetToken, "--file-extension", "xlsx", "--output-path", "./export.xlsx",