github/larksuite-cli
1
0
Fork 0
mirror of https://github.com/larksuite/cli.git synced 2026-07-03 14:02:43 +08:00
Code Issues Packages Projects Releases Wiki Activity

support multipart doc media uploads (#294)

Browse Source 
Change-Id: I9d9fb00079dacfc96b5781e12e6ce79945baa2ed
This commit is contained in:
liujinkun2025
2026-04-08 15:43:15 +08:00
committed by GitHub
parent 7158dc2f3c
commit 6ac5b4d566
16 changed files with 1084 additions and 798 deletions
Download Patch File Download Diff File Expand all files Collapse all files

245
shortcuts/common/drive_media_upload.go Normal file
View File

@@ -0,0 +1,245 @@
// Copyright (c) 2026 Lark Technologies Pte. Ltd.
// SPDX-License-Identifier: MIT
package common
import (
"bytes"
"encoding/json"
"errors"
"fmt"
"io"
"net/http"
"github.com/larksuite/cli/internal/validate"
"github.com/larksuite/cli/internal/vfs"
larkcore "github.com/larksuite/oapi-sdk-go/v3/core"
"github.com/larksuite/cli/internal/output"
)
const MaxDriveMediaUploadSinglePartSize int64 = 20 * 1024 * 1024 // 20MB
const (
driveMediaUploadAllAction = "upload media failed"
driveMediaUploadPartAction = "upload media part failed"
driveMediaUploadFinishAction = "upload media finish failed"
)
type DriveMediaMultipartUploadSession struct {
UploadID string
BlockSize int64
BlockNum int
}
type DriveMediaUploadAllConfig struct {
FilePath string
FileName string
FileSize int64
ParentType string
ParentNode *string
Extra string
}
type DriveMediaMultipartUploadConfig struct {
FilePath string
FileName string
FileSize int64
ParentType string
ParentNode string
Extra string
}
func UploadDriveMediaAll(runtime *RuntimeContext, cfg DriveMediaUploadAllConfig) (string, error) {
safeFilePath, err := validate.SafeInputPath(cfg.FilePath)
if err != nil {
return "", output.ErrValidation("invalid file path: %s", err)
}
f, err := vfs.Open(safeFilePath)
if err != nil {
return "", output.ErrValidation("cannot read file: %s", err)
}
defer f.Close()
fd := larkcore.NewFormdata()
fd.AddField("file_name", cfg.FileName)
fd.AddField("parent_type", cfg.ParentType)
fd.AddField("size", fmt.Sprintf("%d", cfg.FileSize))
if cfg.ParentNode != nil {
fd.AddField("parent_node", *cfg.ParentNode)
}
if cfg.Extra != "" {
fd.AddField("extra", cfg.Extra)
}
fd.AddFile("file", f)
apiResp, err := runtime.DoAPI(&larkcore.ApiReq{
HttpMethod: http.MethodPost,
ApiPath: "/open-apis/drive/v1/medias/upload_all",
Body: fd,
}, larkcore.WithFileUpload())
if err != nil {
return "", WrapDriveMediaUploadRequestError(err, driveMediaUploadAllAction)
}
data, err := ParseDriveMediaUploadResponse(apiResp, driveMediaUploadAllAction)
if err != nil {
return "", err
}
return ExtractDriveMediaUploadFileToken(data, driveMediaUploadAllAction)
}
func UploadDriveMediaMultipart(runtime *RuntimeContext, cfg DriveMediaMultipartUploadConfig) (string, error) {
// upload_prepare expects parent_node to be present even when the caller wants
// the service default/root behavior, so multipart callers pass an explicit
// string instead of relying on field omission like upload_all does.
prepareBody := map[string]interface{}{
"file_name": cfg.FileName,
"parent_type": cfg.ParentType,
"parent_node": cfg.ParentNode,
"size": cfg.FileSize,
}
if cfg.Extra != "" {
prepareBody["extra"] = cfg.Extra
}
data, err := runtime.CallAPI("POST", "/open-apis/drive/v1/medias/upload_prepare", nil, prepareBody)
if err != nil {
return "", err
}
session, err := ParseDriveMediaMultipartUploadSession(data)
if err != nil {
return "", err
}
fmt.Fprintf(runtime.IO().ErrOut, "Multipart upload initialized: %d chunks x %s\n", session.BlockNum, FormatSize(session.BlockSize))
if err = uploadDriveMediaMultipartParts(runtime, cfg.FilePath, cfg.FileSize, session); err != nil {
return "", err
}
return finishDriveMediaMultipartUpload(runtime, session.UploadID, session.BlockNum)
}
func ParseDriveMediaMultipartUploadSession(data map[string]interface{}) (DriveMediaMultipartUploadSession, error) {
// The backend chooses both chunk size and chunk count. Validate them once so
// the streaming loop can follow the returned plan without re-checking shape.
session := DriveMediaMultipartUploadSession{
UploadID: GetString(data, "upload_id"),
BlockSize: int64(GetFloat(data, "block_size")),
BlockNum: int(GetFloat(data, "block_num")),
}
if session.UploadID == "" {
return DriveMediaMultipartUploadSession{}, output.Errorf(output.ExitAPI, "api_error", "upload prepare failed: no upload_id returned")
}
if session.BlockSize <= 0 {
return DriveMediaMultipartUploadSession{}, output.Errorf(output.ExitAPI, "api_error", "upload prepare failed: invalid block_size returned")
}
if session.BlockNum <= 0 {
return DriveMediaMultipartUploadSession{}, output.Errorf(output.ExitAPI, "api_error", "upload prepare failed: invalid block_num returned")
}
return session, nil
}
func WrapDriveMediaUploadRequestError(err error, action string) error {
var exitErr *output.ExitError
if errors.As(err, &exitErr) {
return err
}
return output.ErrNetwork("%s: %v", action, err)
}
func ParseDriveMediaUploadResponse(apiResp *larkcore.ApiResp, action string) (map[string]interface{}, error) {
var result map[string]interface{}
if err := json.Unmarshal(apiResp.RawBody, &result); err != nil {
return nil, output.Errorf(output.ExitAPI, "api_error", "%s: invalid response JSON: %v", action, err)
}
if larkCode := int(GetFloat(result, "code")); larkCode != 0 {
msg, _ := result["msg"].(string)
return nil, output.ErrAPI(larkCode, fmt.Sprintf("%s: [%d] %s", action, larkCode, msg), result["error"])
}
data, _ := result["data"].(map[string]interface{})
return data, nil
}
func ExtractDriveMediaUploadFileToken(data map[string]interface{}, action string) (string, error) {
fileToken := GetString(data, "file_token")
if fileToken == "" {
return "", output.Errorf(output.ExitAPI, "api_error", "%s: no file_token returned", action)
}
return fileToken, nil
}
func uploadDriveMediaMultipartParts(runtime *RuntimeContext, filePath string, fileSize int64, session DriveMediaMultipartUploadSession) error {
safeFilePath, err := validate.SafeInputPath(filePath)
if err != nil {
return output.ErrValidation("invalid file path: %s", err)
}
f, err := vfs.Open(safeFilePath)
if err != nil {
return output.ErrValidation("cannot read file: %s", err)
}
defer f.Close()
maxInt := int64(^uint(0) >> 1)
bufferSize := session.BlockSize
if bufferSize <= 0 || bufferSize > maxInt {
return output.Errorf(output.ExitAPI, "api_error", "upload prepare failed: invalid block_size returned")
}
buffer := make([]byte, int(bufferSize))
remaining := fileSize
// Follow the server-declared block plan exactly; upload_finish expects the
// same block count returned by upload_prepare.
for seq := 0; seq < session.BlockNum; seq++ {
chunkSize := session.BlockSize
if remaining > 0 && chunkSize > remaining {
chunkSize = remaining
}
n, readErr := io.ReadFull(f, buffer[:int(chunkSize)])
if readErr != nil {
return output.ErrValidation("cannot read file: %s", readErr)
}
if err = uploadDriveMediaMultipartPart(runtime, session.UploadID, seq, buffer[:n]); err != nil {
return err
}
fmt.Fprintf(runtime.IO().ErrOut, " Block %d/%d uploaded (%s)\n", seq+1, session.BlockNum, FormatSize(int64(n)))
remaining -= int64(n)
}
return nil
}
func uploadDriveMediaMultipartPart(runtime *RuntimeContext, uploadID string, seq int, chunk []byte) error {
fd := larkcore.NewFormdata()
fd.AddField("upload_id", uploadID)
fd.AddField("seq", fmt.Sprintf("%d", seq))
fd.AddField("size", fmt.Sprintf("%d", len(chunk)))
fd.AddFile("file", bytes.NewReader(chunk))
apiResp, err := runtime.DoAPI(&larkcore.ApiReq{
HttpMethod: http.MethodPost,
ApiPath: "/open-apis/drive/v1/medias/upload_part",
Body: fd,
}, larkcore.WithFileUpload())
if err != nil {
return WrapDriveMediaUploadRequestError(err, driveMediaUploadPartAction)
}
_, err = ParseDriveMediaUploadResponse(apiResp, driveMediaUploadPartAction)
return err
}
func finishDriveMediaMultipartUpload(runtime *RuntimeContext, uploadID string, blockNum int) (string, error) {
data, err := runtime.CallAPI("POST", "/open-apis/drive/v1/medias/upload_finish", nil, map[string]interface{}{
"upload_id": uploadID,
"block_num": blockNum,
})
if err != nil {
return "", err
}
return ExtractDriveMediaUploadFileToken(data, driveMediaUploadFinishAction)
}

528
shortcuts/common/drive_media_upload_test.go Normal file
View File

@@ -0,0 +1,528 @@
// Copyright (c) 2026 Lark Technologies Pte. Ltd.
// SPDX-License-Identifier: MIT
package common
import (
"bytes"
"context"
"encoding/json"
"errors"
"fmt"
"io"
"mime"
"mime/multipart"
"os"
"strings"
"sync/atomic"
"testing"
larkcore "github.com/larksuite/oapi-sdk-go/v3/core"
"github.com/larksuite/cli/internal/cmdutil"
"github.com/larksuite/cli/internal/core"
"github.com/larksuite/cli/internal/httpmock"
"github.com/larksuite/cli/internal/output"
)
var commonDriveMediaUploadTestSeq atomic.Int64
func TestUploadDriveMediaAllBuildsMultipartBody(t *testing.T) {
tests := []struct {
name string
parentNode *string
wantParentNode string
wantParentSet bool
}{
{
name: "includes parent_node when provided",
parentNode: strPtr("blk_parent"),
wantParentNode: "blk_parent",
wantParentSet: true,
},
{
name: "omits parent_node when not provided",
parentNode: nil,
wantParentSet: false,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
runtime, reg := newDriveMediaUploadTestRuntime(t)
withDriveMediaUploadWorkingDir(t, t.TempDir())
uploadStub := &httpmock.Stub{
Method: "POST",
URL: "/open-apis/drive/v1/medias/upload_all",
Body: map[string]interface{}{
"code": 0,
"data": map[string]interface{}{"file_token": "file_all_123"},
},
}
reg.Register(uploadStub)
filePath := writeDriveMediaUploadTestFile(t, "small.bin", 3)
fileToken, err := UploadDriveMediaAll(runtime, DriveMediaUploadAllConfig{
FilePath: filePath,
FileName: "small.bin",
FileSize: 3,
ParentType: "docx_file",
ParentNode: tt.parentNode,
Extra: `{"drive_route_token":"doxcn123"}`,
})
if err != nil {
t.Fatalf("UploadDriveMediaAll() error: %v", err)
}
if fileToken != "file_all_123" {
t.Fatalf("fileToken = %q, want %q", fileToken, "file_all_123")
}
body := decodeCapturedDriveMediaMultipartBody(t, uploadStub)
if got := body.Fields["file_name"]; got != "small.bin" {
t.Fatalf("file_name = %q, want %q", got, "small.bin")
}
if got := body.Fields["parent_type"]; got != "docx_file" {
t.Fatalf("parent_type = %q, want %q", got, "docx_file")
}
if got := body.Fields["size"]; got != "3" {
t.Fatalf("size = %q, want %q", got, "3")
}
if got := body.Fields["extra"]; got != `{"drive_route_token":"doxcn123"}` {
t.Fatalf("extra = %q, want drive route token payload", got)
}
if got := len(body.Files["file"]); got != 3 {
t.Fatalf("file size = %d, want %d", got, 3)
}
gotParentNode, hasParentNode := body.Fields["parent_node"]
if hasParentNode != tt.wantParentSet {
t.Fatalf("parent_node present = %v, want %v", hasParentNode, tt.wantParentSet)
}
if hasParentNode && gotParentNode != tt.wantParentNode {
t.Fatalf("parent_node = %q, want %q", gotParentNode, tt.wantParentNode)
}
})
}
}
func TestUploadDriveMediaMultipartBuildsPreparePartsAndFinish(t *testing.T) {
runtime, reg := newDriveMediaUploadTestRuntime(t)
withDriveMediaUploadWorkingDir(t, t.TempDir())
prepareStub := &httpmock.Stub{
Method: "POST",
URL: "/open-apis/drive/v1/medias/upload_prepare",
Body: map[string]interface{}{
"code": 0,
"data": map[string]interface{}{
"upload_id": "upload_123",
"block_size": float64(4 * 1024 * 1024),
"block_num": float64(6),
},
},
}
reg.Register(prepareStub)
partStubs := make([]*httpmock.Stub, 0, 6)
for i := 0; i < 6; i++ {
stub := &httpmock.Stub{
Method: "POST",
URL: "/open-apis/drive/v1/medias/upload_part",
Body: map[string]interface{}{
"code": 0,
"msg": "ok",
},
}
partStubs = append(partStubs, stub)
reg.Register(stub)
}
finishStub := &httpmock.Stub{
Method: "POST",
URL: "/open-apis/drive/v1/medias/upload_finish",
Body: map[string]interface{}{
"code": 0,
"data": map[string]interface{}{"file_token": "file_multi_123"},
},
}
reg.Register(finishStub)
filePath := writeDriveMediaUploadSizedFile(t, "large.bin", MaxDriveMediaUploadSinglePartSize+1)
fileToken, err := UploadDriveMediaMultipart(runtime, DriveMediaMultipartUploadConfig{
FilePath: filePath,
FileName: "large.bin",
FileSize: MaxDriveMediaUploadSinglePartSize + 1,
ParentType: "ccm_import_open",
ParentNode: "",
Extra: `{"obj_type":"sheet","file_extension":"xlsx"}`,
})
if err != nil {
t.Fatalf("UploadDriveMediaMultipart() error: %v", err)
}
if fileToken != "file_multi_123" {
t.Fatalf("fileToken = %q, want %q", fileToken, "file_multi_123")
}
prepareBody := decodeCapturedDriveMediaJSONBody(t, prepareStub)
if got, _ := prepareBody["parent_type"].(string); got != "ccm_import_open" {
t.Fatalf("prepare parent_type = %q, want %q", got, "ccm_import_open")
}
rawParentNode, ok := prepareBody["parent_node"]
if !ok {
t.Fatal("prepare body missing parent_node")
}
if got, ok := rawParentNode.(string); !ok || got != "" {
t.Fatalf("prepare parent_node = %#v, want empty string", rawParentNode)
}
if got, _ := prepareBody["extra"].(string); got != `{"obj_type":"sheet","file_extension":"xlsx"}` {
t.Fatalf("prepare extra = %q, want import payload", got)
}
if got, _ := prepareBody["size"].(float64); got != float64(MaxDriveMediaUploadSinglePartSize+1) {
t.Fatalf("prepare size = %v, want %d", got, MaxDriveMediaUploadSinglePartSize+1)
}
firstPart := decodeCapturedDriveMediaMultipartBody(t, partStubs[0])
if got := firstPart.Fields["upload_id"]; got != "upload_123" {
t.Fatalf("first part upload_id = %q, want %q", got, "upload_123")
}
if got := firstPart.Fields["seq"]; got != "0" {
t.Fatalf("first part seq = %q, want %q", got, "0")
}
if got := firstPart.Fields["size"]; got != "4194304" {
t.Fatalf("first part size = %q, want %q", got, "4194304")
}
lastPart := decodeCapturedDriveMediaMultipartBody(t, partStubs[len(partStubs)-1])
if got := lastPart.Fields["seq"]; got != "5" {
t.Fatalf("last part seq = %q, want %q", got, "5")
}
if got := lastPart.Fields["size"]; got != "1" {
t.Fatalf("last part size = %q, want %q", got, "1")
}
if got := len(lastPart.Files["file"]); got != 1 {
t.Fatalf("last part file size = %d, want %d", got, 1)
}
finishBody := decodeCapturedDriveMediaJSONBody(t, finishStub)
if got, _ := finishBody["upload_id"].(string); got != "upload_123" {
t.Fatalf("finish upload_id = %q, want %q", got, "upload_123")
}
if got, _ := finishBody["block_num"].(float64); got != 6 {
t.Fatalf("finish block_num = %v, want %d", got, 6)
}
}
func TestParseDriveMediaMultipartUploadSessionValidatesResponseFields(t *testing.T) {
t.Parallel()
tests := []struct {
name string
data map[string]interface{}
wantText string
}{
{
name: "missing upload id",
data: map[string]interface{}{
"block_size": 4 * 1024 * 1024,
"block_num": 6,
},
wantText: "upload prepare failed: no upload_id returned",
},
{
name: "missing block size",
data: map[string]interface{}{
"upload_id": "upload_123",
"block_num": 6,
},
wantText: "upload prepare failed: invalid block_size returned",
},
{
name: "missing block num",
data: map[string]interface{}{
"upload_id": "upload_123",
"block_size": 4 * 1024 * 1024,
},
wantText: "upload prepare failed: invalid block_num returned",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
t.Parallel()
_, err := ParseDriveMediaMultipartUploadSession(tt.data)
if err == nil || !strings.Contains(err.Error(), tt.wantText) {
t.Fatalf("err = %v, want substring %q", err, tt.wantText)
}
})
}
}
func TestUploadDriveMediaMultipartPartAPIFailure(t *testing.T) {
runtime, reg := newDriveMediaUploadTestRuntime(t)
withDriveMediaUploadWorkingDir(t, t.TempDir())
reg.Register(&httpmock.Stub{
Method: "POST",
URL: "/open-apis/drive/v1/medias/upload_prepare",
Body: map[string]interface{}{
"code": 0,
"data": map[string]interface{}{
"upload_id": "upload_123",
"block_size": float64(4 * 1024 * 1024),
"block_num": float64(6),
},
},
})
reg.Register(&httpmock.Stub{
Method: "POST",
URL: "/open-apis/drive/v1/medias/upload_part",
Body: map[string]interface{}{
"code": 999,
"msg": "chunk rejected",
},
})
filePath := writeDriveMediaUploadSizedFile(t, "large.bin", MaxDriveMediaUploadSinglePartSize+1)
_, err := UploadDriveMediaMultipart(runtime, DriveMediaMultipartUploadConfig{
FilePath: filePath,
FileName: "large.bin",
FileSize: MaxDriveMediaUploadSinglePartSize + 1,
ParentType: "ccm_import_open",
ParentNode: "",
})
if err == nil {
t.Fatal("expected error, got nil")
}
if !strings.Contains(err.Error(), "upload media part failed: [999] chunk rejected") {
t.Fatalf("unexpected error: %v", err)
}
}
func TestUploadDriveMediaMultipartFinishRequiresFileToken(t *testing.T) {
runtime, reg := newDriveMediaUploadTestRuntime(t)
withDriveMediaUploadWorkingDir(t, t.TempDir())
reg.Register(&httpmock.Stub{
Method: "POST",
URL: "/open-apis/drive/v1/medias/upload_prepare",
Body: map[string]interface{}{
"code": 0,
"data": map[string]interface{}{
"upload_id": "upload_123",
"block_size": float64(4 * 1024 * 1024),
"block_num": float64(6),
},
},
})
for i := 0; i < 6; i++ {
reg.Register(&httpmock.Stub{
Method: "POST",
URL: "/open-apis/drive/v1/medias/upload_part",
Body: map[string]interface{}{
"code": 0,
"msg": "ok",
},
})
}
reg.Register(&httpmock.Stub{
Method: "POST",
URL: "/open-apis/drive/v1/medias/upload_finish",
Body: map[string]interface{}{
"code": 0,
"data": map[string]interface{}{},
},
})
filePath := writeDriveMediaUploadSizedFile(t, "large.bin", MaxDriveMediaUploadSinglePartSize+1)
_, err := UploadDriveMediaMultipart(runtime, DriveMediaMultipartUploadConfig{
FilePath: filePath,
FileName: "large.bin",
FileSize: MaxDriveMediaUploadSinglePartSize + 1,
ParentType: "ccm_import_open",
ParentNode: "",
})
if err == nil {
t.Fatal("expected error, got nil")
}
if !strings.Contains(err.Error(), "upload media finish failed: no file_token returned") {
t.Fatalf("unexpected error: %v", err)
}
}
func TestParseDriveMediaUploadResponseErrors(t *testing.T) {
t.Parallel()
t.Run("invalid json", func(t *testing.T) {
t.Parallel()
_, err := ParseDriveMediaUploadResponse(&larkcore.ApiResp{RawBody: []byte("{")}, "upload media failed")
if err == nil || !strings.Contains(err.Error(), "invalid response JSON") {
t.Fatalf("expected invalid JSON error, got %v", err)
}
})
t.Run("api code error", func(t *testing.T) {
t.Parallel()
_, err := ParseDriveMediaUploadResponse(&larkcore.ApiResp{RawBody: []byte(`{"code":999,"msg":"boom","error":{"detail":"x"}}`)}, "upload media failed")
if err == nil || !strings.Contains(err.Error(), "upload media failed: [999] boom") {
t.Fatalf("expected API error, got %v", err)
}
})
}
func TestExtractDriveMediaUploadFileTokenRequiresToken(t *testing.T) {
t.Parallel()
_, err := ExtractDriveMediaUploadFileToken(map[string]interface{}{}, "upload media failed")
if err == nil || !strings.Contains(err.Error(), "upload media failed: no file_token returned") {
t.Fatalf("err = %v, want missing file_token error", err)
}
}
func TestWrapDriveMediaUploadRequestError(t *testing.T) {
t.Parallel()
t.Run("preserves exit error", func(t *testing.T) {
t.Parallel()
original := output.ErrValidation("bad input")
got := WrapDriveMediaUploadRequestError(original, "upload media failed")
if got != original {
t.Fatalf("expected same exit error pointer, got %v", got)
}
})
t.Run("wraps generic error as network", func(t *testing.T) {
t.Parallel()
got := WrapDriveMediaUploadRequestError(io.EOF, "upload media failed")
var exitErr *output.ExitError
if !errors.As(got, &exitErr) {
t.Fatalf("expected ExitError, got %T", got)
}
if exitErr.Code != output.ExitNetwork {
t.Fatalf("exit code = %d, want %d", exitErr.Code, output.ExitNetwork)
}
if !strings.Contains(got.Error(), "upload media failed") {
t.Fatalf("unexpected error: %v", got)
}
})
}
type capturedDriveMediaMultipartBody struct {
Fields map[string]string
Files map[string][]byte
}
func newDriveMediaUploadTestRuntime(t *testing.T) (*RuntimeContext, *httpmock.Registry) {
t.Helper()
t.Setenv("LARKSUITE_CLI_CONFIG_DIR", t.TempDir())
cfg := &core.CliConfig{
AppID: fmt.Sprintf("common-drive-media-test-%d", commonDriveMediaUploadTestSeq.Add(1)), AppSecret: "test-secret", Brand: core.BrandFeishu,
}
f, _, _, reg := cmdutil.TestFactory(t, cfg)
runtime := &RuntimeContext{
ctx: context.Background(),
Config: cfg,
Factory: f,
resolvedAs: core.AsBot,
}
return runtime, reg
}
func withDriveMediaUploadWorkingDir(t *testing.T, dir string) {
t.Helper()
cwd, err := os.Getwd()
if err != nil {
t.Fatalf("Getwd() error: %v", err)
}
if err := os.Chdir(dir); err != nil {
t.Fatalf("Chdir(%q) error: %v", dir, err)
}
t.Cleanup(func() {
if err := os.Chdir(cwd); err != nil {
t.Fatalf("restore cwd error: %v", err)
}
})
}
func writeDriveMediaUploadTestFile(t *testing.T, name string, size int) string {
t.Helper()
if err := os.WriteFile(name, bytes.Repeat([]byte("a"), size), 0644); err != nil {
t.Fatalf("WriteFile(%q) error: %v", name, err)
}
return name
}
func writeDriveMediaUploadSizedFile(t *testing.T, name string, size int64) string {
t.Helper()
fh, err := os.Create(name)
if err != nil {
t.Fatalf("Create(%q) error: %v", name, err)
}
if err := fh.Truncate(size); err != nil {
t.Fatalf("Truncate(%q) error: %v", name, err)
}
if err := fh.Close(); err != nil {
t.Fatalf("Close(%q) error: %v", name, err)
}
return name
}
func decodeCapturedDriveMediaJSONBody(t *testing.T, stub *httpmock.Stub) map[string]interface{} {
t.Helper()
var body map[string]interface{}
if err := json.Unmarshal(stub.CapturedBody, &body); err != nil {
t.Fatalf("decode captured JSON body: %v", err)
}
return body
}
func decodeCapturedDriveMediaMultipartBody(t *testing.T, stub *httpmock.Stub) capturedDriveMediaMultipartBody {
t.Helper()
contentType := stub.CapturedHeaders.Get("Content-Type")
mediaType, params, err := mime.ParseMediaType(contentType)
if err != nil {
t.Fatalf("parse multipart content type: %v", err)
}
if mediaType != "multipart/form-data" {
t.Fatalf("content type = %q, want multipart/form-data", mediaType)
}
reader := multipart.NewReader(bytes.NewReader(stub.CapturedBody), params["boundary"])
body := capturedDriveMediaMultipartBody{
Fields: map[string]string{},
Files: map[string][]byte{},
}
for {
part, err := reader.NextPart()
if err == io.EOF {
break
}
if err != nil {
t.Fatalf("read multipart part: %v", err)
}
data, err := io.ReadAll(part)
if err != nil {
t.Fatalf("read multipart data: %v", err)
}
if part.FileName() != "" {
body.Files[part.FormName()] = data
continue
}
body.Fields[part.FormName()] = string(data)
}
return body
}
func strPtr(s string) *string {
return &s
}

138
shortcuts/doc/doc_media_insert.go
View File

@@ -5,23 +5,15 @@ package doc
import (
"context"
"encoding/json"
"errors"
"fmt"
"net/http"
"path/filepath"
larkcore "github.com/larksuite/oapi-sdk-go/v3/core"
"github.com/larksuite/cli/internal/output"
"github.com/larksuite/cli/internal/util"
"github.com/larksuite/cli/internal/validate"
"github.com/larksuite/cli/internal/vfs"
"github.com/larksuite/cli/shortcuts/common"
)
const maxFileSize = 20 * 1024 * 1024 // 20MB
var alignMap = map[string]int{
"left": 1,
"center": 2,
@@ -36,7 +28,7 @@ var DocMediaInsert = common.Shortcut{
Scopes: []string{"docs:document.media:upload", "docx:document:write_only", "docx:document:readonly"},
AuthTypes: []string{"user", "bot"},
Flags: []common.Flag{
{Name: "file", Desc: "local file path (max 20MB)", Required: true},
{Name: "file", Desc: "local file path (files > 20MB use multipart upload automatically)", Required: true},
{Name: "doc", Desc: "document URL or document_id", Required: true},
{Name: "type", Default: "image", Desc: "type: image | file"},
{Name: "align", Desc: "alignment: left | center | right"},
@@ -86,16 +78,9 @@ var DocMediaInsert = common.Shortcut{
Desc(fmt.Sprintf("[%d] Get document root block", stepBase)).
POST("/open-apis/docx/v1/documents/:document_id/blocks/:document_id/children").
Desc(fmt.Sprintf("[%d] Create empty block at document end", stepBase+1)).
Body(createBlockData).
POST("/open-apis/drive/v1/medias/upload_all").
Desc(fmt.Sprintf("[%d] Upload local file (multipart/form-data)", stepBase+2)).
Body(map[string]interface{}{
"file_name": filepath.Base(filePath),
"parent_type": parentType,
"parent_node": "<new_block_id>",
"file": "@" + filePath,
}).
PATCH("/open-apis/docx/v1/documents/:document_id/blocks/batch_update").
Body(createBlockData)
appendDocMediaInsertUploadDryRun(d, filePath, parentType, stepBase+2)
d.PATCH("/open-apis/docx/v1/documents/:document_id/blocks/batch_update").
Desc(fmt.Sprintf("[%d] Bind uploaded file token to the new block", stepBase+3)).
Body(batchUpdateData)
@@ -112,7 +97,6 @@ var DocMediaInsert = common.Shortcut{
if pathErr != nil {
return output.ErrValidation("unsafe file path: %s", pathErr)
}
filePath = safeFilePath
documentID, err := resolveDocxDocumentID(runtime, docInput)
if err != nil {
@@ -120,16 +104,19 @@ var DocMediaInsert = common.Shortcut{
}
// Validate file
stat, err := vfs.Stat(filePath)
stat, err := vfs.Stat(safeFilePath)
if err != nil {
return output.ErrValidation("file not found: %s", filePath)
}
if stat.Size() > maxFileSize {
return output.ErrValidation("file %.1fMB exceeds 20MB limit", float64(stat.Size())/1024/1024)
if !stat.Mode().IsRegular() {
return output.ErrValidation("file must be a regular file: %s", filePath)
}
fileName := filepath.Base(filePath)
fmt.Fprintf(runtime.IO().ErrOut, "Inserting: %s -> document %s\n", fileName, common.MaskToken(documentID))
if stat.Size() > common.MaxDriveMediaUploadSinglePartSize {
fmt.Fprintf(runtime.IO().ErrOut, "File exceeds 20MB, using multipart upload\n")
}
// Step 1: Get document root block to find where to insert
rootData, err := runtime.CallAPI("GET",
@@ -166,7 +153,8 @@ var DocMediaInsert = common.Shortcut{
fmt.Fprintf(runtime.IO().ErrOut, "Resolved file block targets: upload=%s replace=%s\n", uploadParentNode, replaceBlockID)
}
// Rollback helper
// The placeholder block is created before any upload starts, so failures in
// later steps should try to remove it instead of leaving an empty artifact.
rollback := func() error {
fmt.Fprintf(runtime.IO().ErrOut, "Rolling back: deleting block %s\n", blockId)
_, err := runtime.CallAPI("DELETE",
@@ -185,7 +173,7 @@ var DocMediaInsert = common.Shortcut{
}
// Step 3: Upload media file
fileToken, err := uploadMediaFile(ctx, runtime, filePath, fileName, mediaType, uploadParentNode, documentID)
fileToken, err := uploadDocMediaFile(runtime, filePath, fileName, stat.Size(), parentTypeForMediaType(mediaType), uploadParentNode, documentID)
if err != nil {
return withRollbackWarning(err)
}
@@ -346,6 +334,8 @@ func extractCreatedBlockTargets(createData map[string]interface{}, mediaType str
return blockID, uploadParentNode, replaceBlockID
}
// File blocks are wrapped: the created top-level block owns a nested child
// that is both the upload target and the replace_file target.
nestedChildren, _ := child["children"].([]interface{})
if len(nestedChildren) == 0 {
return blockID, uploadParentNode, replaceBlockID
@@ -357,66 +347,44 @@ func extractCreatedBlockTargets(createData map[string]interface{}, mediaType str
return blockID, uploadParentNode, replaceBlockID
}
// uploadMediaFile uploads a file to Feishu drive as media.
func uploadMediaFile(ctx context.Context, runtime *common.RuntimeContext, filePath, fileName, mediaType, parentNode, docId string) (string, error) {
f, err := vfs.Open(filePath)
if err != nil {
return "", err
}
defer f.Close()
stat, err := f.Stat()
if err != nil {
return "", output.Errorf(output.ExitInternal, "internal_error", "failed to stat file: %v", err)
}
fileSize := stat.Size()
parentType := parentTypeForMediaType(mediaType)
// Build SDK Formdata
fd := larkcore.NewFormdata()
fd.AddField("file_name", fileName)
fd.AddField("parent_type", parentType)
fd.AddField("parent_node", parentNode)
fd.AddField("size", fmt.Sprintf("%d", fileSize))
if docId != "" {
extra, err := buildDriveRouteExtra(docId)
if err != nil {
return "", err
}
fd.AddField("extra", extra)
}
fd.AddFile("file", f)
apiResp, err := runtime.DoAPI(&larkcore.ApiReq{
HttpMethod: http.MethodPost,
ApiPath: "/open-apis/drive/v1/medias/upload_all",
Body: fd,
}, larkcore.WithFileUpload())
if err != nil {
var exitErr *output.ExitError
if errors.As(err, &exitErr) {
return "", err
}
return "", output.ErrNetwork("file upload failed: %v", err)
func appendDocMediaInsertUploadDryRun(d *common.DryRunAPI, filePath, parentType string, step int) {
// The upload step runs only after the empty placeholder block is created, so
// dry-run can refer to that future block ID only symbolically. For large
// files, keep multipart internals as substeps of the single user-facing
// "upload file" step.
if docMediaShouldUseMultipart(filePath) {
d.POST("/open-apis/drive/v1/medias/upload_prepare").
Desc(fmt.Sprintf("[%da] Initialize multipart upload", step)).
Body(map[string]interface{}{
"file_name": filepath.Base(filePath),
"parent_type": parentType,
"parent_node": "<new_block_id>",
"size": "<file_size>",
}).
POST("/open-apis/drive/v1/medias/upload_part").
Desc(fmt.Sprintf("[%db] Upload file parts (repeated)", step)).
Body(map[string]interface{}{
"upload_id": "<upload_id>",
"seq": "<chunk_index>",
"size": "<chunk_size>",
"file": "<chunk_binary>",
}).
POST("/open-apis/drive/v1/medias/upload_finish").
Desc(fmt.Sprintf("[%dc] Finalize multipart upload and get file_token", step)).
Body(map[string]interface{}{
"upload_id": "<upload_id>",
"block_num": "<block_num>",
})
return
}
var result map[string]interface{}
if err := json.Unmarshal(apiResp.RawBody, &result); err != nil {
return "", output.Errorf(output.ExitAPI, "api_error", "file upload failed: invalid response JSON: %v", err)
}
code, _ := util.ToFloat64(result["code"])
if code != 0 {
msg, _ := result["msg"].(string)
return "", output.ErrAPI(int(code), fmt.Sprintf("file upload failed: [%d] %s", int(code), msg), result["error"])
}
data, _ := result["data"].(map[string]interface{})
fileToken, _ := data["file_token"].(string)
if fileToken == "" {
return "", output.Errorf(output.ExitAPI, "api_error", "file upload failed: no file_token returned")
}
return fileToken, nil
d.POST("/open-apis/drive/v1/medias/upload_all").
Desc(fmt.Sprintf("[%d] Upload local file (multipart/form-data)", step)).
Body(map[string]interface{}{
"file_name": filepath.Base(filePath),
"parent_type": parentType,
"parent_node": "<new_block_id>",
"size": "<file_size>",
"file": "@" + filePath,
})
}

141
shortcuts/doc/doc_media_test.go
View File

@@ -5,6 +5,8 @@ package doc
import (
"bytes"
"context"
"encoding/json"
"net/http"
"os"
"path/filepath"
@@ -19,10 +21,6 @@ import (
"github.com/larksuite/cli/shortcuts/common"
)
func docsTestConfig() *core.CliConfig {
return docsTestConfigWithAppID("docs-test-app")
}
func docsTestConfigWithAppID(appID string) *core.CliConfig {
return &core.CliConfig{
AppID: appID, AppSecret: "test-secret", Brand: core.BrandFeishu,
@@ -59,7 +57,7 @@ func withDocsWorkingDir(t *testing.T, dir string) {
}
func TestDocMediaInsertRejectsOldDocURL(t *testing.T) {
f, _, _, _ := cmdutil.TestFactory(t, docsTestConfig())
f, _, _, _ := cmdutil.TestFactory(t, docsTestConfigWithAppID("docs-test-app"))
err := mountAndRunDocs(t, DocMediaInsert, []string{
"+media-insert",
@@ -77,7 +75,7 @@ func TestDocMediaInsertRejectsOldDocURL(t *testing.T) {
}
func TestDocMediaInsertDryRunWikiAddsResolveStep(t *testing.T) {
f, stdout, _, _ := cmdutil.TestFactory(t, docsTestConfig())
f, stdout, _, _ := cmdutil.TestFactory(t, docsTestConfigWithAppID("docs-test-app"))
err := mountAndRunDocs(t, DocMediaInsert, []string{
"+media-insert",
@@ -99,6 +97,98 @@ func TestDocMediaInsertDryRunWikiAddsResolveStep(t *testing.T) {
}
}
func TestDocMediaUploadDryRunUsesMultipartForLargeFile(t *testing.T) {
tmpDir := t.TempDir()
withDocsWorkingDir(t, tmpDir)
writeSizedDocTestFile(t, "large.bin", common.MaxDriveMediaUploadSinglePartSize+1)
cmd := &cobra.Command{Use: "docs +media-upload"}
cmd.Flags().String("file", "", "")
cmd.Flags().String("parent-type", "", "")
cmd.Flags().String("parent-node", "", "")
cmd.Flags().String("doc-id", "", "")
if err := cmd.Flags().Set("file", "./large.bin"); err != nil {
t.Fatalf("set --file: %v", err)
}
if err := cmd.Flags().Set("parent-type", "docx_file"); err != nil {
t.Fatalf("set --parent-type: %v", err)
}
if err := cmd.Flags().Set("parent-node", "blk_parent"); err != nil {
t.Fatalf("set --parent-node: %v", err)
}
dry := decodeDocDryRun(t, MediaUpload.DryRun(context.Background(), common.TestNewRuntimeContext(cmd, nil)))
if dry.Description != "chunked media upload (files > 20MB)" {
t.Fatalf("dry-run description = %q", dry.Description)
}
if len(dry.API) != 3 {
t.Fatalf("expected 3 API calls, got %d", len(dry.API))
}
if dry.API[0].URL != "/open-apis/drive/v1/medias/upload_prepare" {
t.Fatalf("first URL = %q, want upload_prepare", dry.API[0].URL)
}
if dry.API[1].URL != "/open-apis/drive/v1/medias/upload_part" {
t.Fatalf("second URL = %q, want upload_part", dry.API[1].URL)
}
if dry.API[2].URL != "/open-apis/drive/v1/medias/upload_finish" {
t.Fatalf("third URL = %q, want upload_finish", dry.API[2].URL)
}
if got, _ := dry.API[0].Body["parent_node"].(string); got != "blk_parent" {
t.Fatalf("prepare parent_node = %q, want %q", got, "blk_parent")
}
}
func TestDocMediaInsertDryRunUsesMultipartForLargeFile(t *testing.T) {
tmpDir := t.TempDir()
withDocsWorkingDir(t, tmpDir)
writeSizedDocTestFile(t, "large.bin", common.MaxDriveMediaUploadSinglePartSize+1)
cmd := &cobra.Command{Use: "docs +media-insert"}
cmd.Flags().String("file", "", "")
cmd.Flags().String("doc", "", "")
cmd.Flags().String("type", "", "")
cmd.Flags().String("align", "", "")
cmd.Flags().String("caption", "", "")
if err := cmd.Flags().Set("doc", "doxcnDryRunLarge"); err != nil {
t.Fatalf("set --doc: %v", err)
}
if err := cmd.Flags().Set("file", "./large.bin"); err != nil {
t.Fatalf("set --file: %v", err)
}
dry := decodeDocDryRun(t, DocMediaInsert.DryRun(context.Background(), common.TestNewRuntimeContext(cmd, nil)))
if dry.Description != "4-step orchestration: query root → create block → upload file → bind to block (auto-rollback on failure)" {
t.Fatalf("dry-run description = %q", dry.Description)
}
if len(dry.API) != 6 {
t.Fatalf("expected 6 API calls, got %d", len(dry.API))
}
if dry.API[2].URL != "/open-apis/drive/v1/medias/upload_prepare" {
t.Fatalf("third URL = %q, want upload_prepare", dry.API[2].URL)
}
if dry.API[3].URL != "/open-apis/drive/v1/medias/upload_part" {
t.Fatalf("fourth URL = %q, want upload_part", dry.API[3].URL)
}
if dry.API[4].URL != "/open-apis/drive/v1/medias/upload_finish" {
t.Fatalf("fifth URL = %q, want upload_finish", dry.API[4].URL)
}
if dry.API[5].URL != "/open-apis/docx/v1/documents/doxcnDryRunLarge/blocks/batch_update" {
t.Fatalf("last URL = %q, want batch_update", dry.API[5].URL)
}
if !strings.Contains(dry.API[2].Desc, "[3a]") {
t.Fatalf("upload_prepare desc = %q, want [3a] step marker", dry.API[2].Desc)
}
if !strings.Contains(dry.API[3].Desc, "[3b]") {
t.Fatalf("upload_part desc = %q, want [3b] step marker", dry.API[3].Desc)
}
if !strings.Contains(dry.API[4].Desc, "[3c]") {
t.Fatalf("upload_finish desc = %q, want [3c] step marker", dry.API[4].Desc)
}
if !strings.Contains(dry.API[5].Desc, "[4]") {
t.Fatalf("batch_update desc = %q, want [4] step marker", dry.API[5].Desc)
}
}
func TestDocMediaInsertExecuteResolvesWikiBeforeFileCheck(t *testing.T) {
f, _, stderr, reg := cmdutil.TestFactory(t, docsTestConfigWithAppID("docs-insert-exec-app"))
reg.Register(&httpmock.Stub{
@@ -194,3 +284,42 @@ func TestDocMediaDownloadRejectsHTTPErrorBeforeWrite(t *testing.T) {
t.Fatalf("download target should not be created, statErr=%v", statErr)
}
}
type docDryRunOutput struct {
Description string `json:"description"`
API []struct {
Desc string `json:"desc"`
URL string `json:"url"`
Body map[string]interface{} `json:"body"`
} `json:"api"`
}
func writeSizedDocTestFile(t *testing.T, name string, size int64) {
t.Helper()
fh, err := os.Create(name)
if err != nil {
t.Fatalf("Create(%q) error: %v", name, err)
}
if err := fh.Truncate(size); err != nil {
t.Fatalf("Truncate(%q) error: %v", name, err)
}
if err := fh.Close(); err != nil {
t.Fatalf("Close(%q) error: %v", name, err)
}
}
func decodeDocDryRun(t *testing.T, dryAPI *common.DryRunAPI) docDryRunOutput {
t.Helper()
raw, err := json.Marshal(dryAPI)
if err != nil {
t.Fatalf("marshal dry-run output: %v", err)
}
var dry docDryRunOutput
if err := json.Unmarshal(raw, &dry); err != nil {
t.Fatalf("decode dry-run output: %v", err)
}
return dry
}

149
shortcuts/doc/doc_media_upload.go
View File

@@ -5,16 +5,10 @@ package doc
import (
"context"
"encoding/json"
"errors"
"fmt"
"net/http"
"path/filepath"
larkcore "github.com/larksuite/oapi-sdk-go/v3/core"
"github.com/larksuite/cli/internal/output"
"github.com/larksuite/cli/internal/util"
"github.com/larksuite/cli/internal/validate"
"github.com/larksuite/cli/internal/vfs"
"github.com/larksuite/cli/shortcuts/common"
@@ -28,7 +22,7 @@ var MediaUpload = common.Shortcut{
Scopes: []string{"docs:document.media:upload"},
AuthTypes: []string{"user", "bot"},
Flags: []common.Flag{
{Name: "file", Desc: "local file path (max 20MB)", Required: true},
{Name: "file", Desc: "local file path (files > 20MB use multipart upload automatically)", Required: true},
{Name: "parent-type", Desc: "parent type: docx_image | docx_file", Required: true},
{Name: "parent-node", Desc: "parent node ID (block_id)", Required: true},
{Name: "doc-id", Desc: "document ID (for drive_route_token)"},
@@ -42,13 +36,42 @@ var MediaUpload = common.Shortcut{
"file_name": filepath.Base(filePath),
"parent_type": parentType,
"parent_node": parentNode,
"file": "@" + filePath,
}
if docId != "" {
body["extra"] = fmt.Sprintf(`{"drive_route_token":"%s"}`, docId)
}
return common.NewDryRunAPI().
Desc("multipart/form-data upload").
dry := common.NewDryRunAPI()
if docMediaShouldUseMultipart(filePath) {
prepareBody := map[string]interface{}{
"file_name": filepath.Base(filePath),
"parent_type": parentType,
"parent_node": parentNode,
"size": "<file_size>",
}
if extra, ok := body["extra"]; ok {
prepareBody["extra"] = extra
}
dry.Desc("chunked media upload (files > 20MB)").
POST("/open-apis/drive/v1/medias/upload_prepare").
Body(prepareBody).
POST("/open-apis/drive/v1/medias/upload_part").
Body(map[string]interface{}{
"upload_id": "<upload_id>",
"seq": "<chunk_index>",
"size": "<chunk_size>",
"file": "<chunk_binary>",
}).
POST("/open-apis/drive/v1/medias/upload_finish").
Body(map[string]interface{}{
"upload_id": "<upload_id>",
"block_num": "<block_num>",
})
return dry
}
body["file"] = "@" + filePath
body["size"] = "<file_size>"
return dry.Desc("multipart/form-data upload").
POST("/open-apis/drive/v1/medias/upload_all").
Body(body)
},
@@ -62,69 +85,25 @@ var MediaUpload = common.Shortcut{
if pathErr != nil {
return output.ErrValidation("unsafe file path: %s", pathErr)
}
filePath = safeFilePath
// Validate file
stat, err := vfs.Stat(filePath)
stat, err := vfs.Stat(safeFilePath)
if err != nil {
return output.ErrValidation("file not found: %s", filePath)
}
if stat.Size() > maxFileSize {
return output.ErrValidation("file %.1fMB exceeds 20MB limit", float64(stat.Size())/1024/1024)
if !stat.Mode().IsRegular() {
return output.ErrValidation("file must be a regular file: %s", filePath)
}
fileName := filepath.Base(filePath)
fmt.Fprintf(runtime.IO().ErrOut, "Uploading: %s (%d bytes)\n", fileName, stat.Size())
if stat.Size() > common.MaxDriveMediaUploadSinglePartSize {
fmt.Fprintf(runtime.IO().ErrOut, "File exceeds 20MB, using multipart upload\n")
}
f, err := vfs.Open(filePath)
fileToken, err := uploadDocMediaFile(runtime, filePath, fileName, stat.Size(), parentType, parentNode, docId)
if err != nil {
return output.ErrValidation("cannot open file: %v", err)
}
defer f.Close()
// Build SDK Formdata
fd := larkcore.NewFormdata()
fd.AddField("file_name", fileName)
fd.AddField("parent_type", parentType)
fd.AddField("parent_node", parentNode)
fd.AddField("size", fmt.Sprintf("%d", stat.Size()))
if docId != "" {
extra, err := buildDriveRouteExtra(docId)
if err != nil {
return err
}
fd.AddField("extra", extra)
}
fd.AddFile("file", f)
apiResp, err := runtime.DoAPI(&larkcore.ApiReq{
HttpMethod: http.MethodPost,
ApiPath: "/open-apis/drive/v1/medias/upload_all",
Body: fd,
}, larkcore.WithFileUpload())
if err != nil {
var exitErr *output.ExitError
if errors.As(err, &exitErr) {
return err
}
return output.ErrNetwork("upload failed: %v", err)
}
var result map[string]interface{}
if err := json.Unmarshal(apiResp.RawBody, &result); err != nil {
return output.Errorf(output.ExitAPI, "api_error", "upload failed: invalid response JSON: %v", err)
}
code, _ := util.ToFloat64(result["code"])
if code != 0 {
msg, _ := result["msg"].(string)
return output.ErrAPI(int(code), fmt.Sprintf("upload failed: [%d] %s", int(code), msg), result["error"])
}
data, _ := result["data"].(map[string]interface{})
fileToken, _ := data["file_token"].(string)
if fileToken == "" {
return output.Errorf(output.ExitAPI, "api_error", "upload failed: no file_token returned")
return err
}
runtime.Out(map[string]interface{}{
@@ -135,3 +114,49 @@ var MediaUpload = common.Shortcut{
return nil
},
}
func uploadDocMediaFile(runtime *common.RuntimeContext, filePath, fileName string, fileSize int64, parentType, parentNode, docID string) (string, error) {
var extra string
if docID != "" {
var err error
extra, err = buildDriveRouteExtra(docID)
if err != nil {
return "", err
}
}
// Doc media uploads share the generic Drive media transport. The doc-specific
// routing only shows up in parent_type/parent_node and optional route extra.
if fileSize <= common.MaxDriveMediaUploadSinglePartSize {
return common.UploadDriveMediaAll(runtime, common.DriveMediaUploadAllConfig{
FilePath: filePath,
FileName: fileName,
FileSize: fileSize,
ParentType: parentType,
ParentNode: &parentNode,
Extra: extra,
})
}
return common.UploadDriveMediaMultipart(runtime, common.DriveMediaMultipartUploadConfig{
FilePath: filePath,
FileName: fileName,
FileSize: fileSize,
ParentType: parentType,
ParentNode: parentNode,
Extra: extra,
})
}
func docMediaShouldUseMultipart(filePath string) bool {
// Dry-run uses local stat as a best-effort planning hint. Execute re-validates
// the file before choosing the actual upload path.
safeFilePath, err := validate.SafeInputPath(filePath)
if err != nil {
return false
}
info, err := vfs.Stat(safeFilePath)
if err != nil {
return false
}
return info.Mode().IsRegular() && info.Size() > common.MaxDriveMediaUploadSinglePartSize
}

8
shortcuts/drive/drive_export_test.go
View File

@@ -65,7 +65,6 @@ func TestValidateDriveExportSpec(t *testing.T) {
func TestDriveExportMarkdownWritesFile(t *testing.T) {
f, stdout, _, reg := cmdutil.TestFactory(t, driveTestConfig())
registerDriveBotTokenStub(reg)
reg.Register(&httpmock.Stub{
Method: "GET",
URL: "/open-apis/docs/v1/content",
@@ -117,7 +116,6 @@ func TestDriveExportMarkdownWritesFile(t *testing.T) {
func TestDriveExportAsyncSuccess(t *testing.T) {
f, stdout, _, reg := cmdutil.TestFactory(t, driveTestConfig())
registerDriveBotTokenStub(reg)
reg.Register(&httpmock.Stub{
Method: "POST",
URL: "/open-apis/drive/v1/export_tasks",
@@ -188,7 +186,6 @@ func TestDriveExportAsyncSuccess(t *testing.T) {
func TestDriveExportReadyDownloadFailureIncludesRecoveryHint(t *testing.T) {
f, _, _, reg := cmdutil.TestFactory(t, driveTestConfig())
registerDriveBotTokenStub(reg)
reg.Register(&httpmock.Stub{
Method: "POST",
URL: "/open-apis/drive/v1/export_tasks",
@@ -267,7 +264,6 @@ func TestDriveExportReadyDownloadFailureIncludesRecoveryHint(t *testing.T) {
func TestDriveExportTimeoutReturnsFollowUpCommand(t *testing.T) {
f, stdout, _, reg := cmdutil.TestFactory(t, driveTestConfig())
registerDriveBotTokenStub(reg)
reg.Register(&httpmock.Stub{
Method: "POST",
URL: "/open-apis/drive/v1/export_tasks",
@@ -333,7 +329,6 @@ func TestDriveExportTimeoutReturnsFollowUpCommand(t *testing.T) {
func TestDriveExportPollErrorsReturnLastErrorWithRecoveryHint(t *testing.T) {
f, stdout, _, reg := cmdutil.TestFactory(t, driveTestConfig())
registerDriveBotTokenStub(reg)
reg.Register(&httpmock.Stub{
Method: "POST",
URL: "/open-apis/drive/v1/export_tasks",
@@ -389,7 +384,6 @@ func TestDriveExportPollErrorsReturnLastErrorWithRecoveryHint(t *testing.T) {
func TestDriveExportDownloadUsesProvidedFileName(t *testing.T) {
f, stdout, _, reg := cmdutil.TestFactory(t, driveTestConfig())
registerDriveBotTokenStub(reg)
reg.Register(&httpmock.Stub{
Method: "GET",
URL: "/open-apis/drive/v1/export_tasks/file/box_789/download",
@@ -424,7 +418,6 @@ func TestDriveExportDownloadUsesProvidedFileName(t *testing.T) {
func TestDriveExportDownloadRejectsOverwriteWithoutFlag(t *testing.T) {
f, _, _, reg := cmdutil.TestFactory(t, driveTestConfig())
registerDriveBotTokenStub(reg)
reg.Register(&httpmock.Stub{
Method: "GET",
URL: "/open-apis/drive/v1/export_tasks/file/box_dup/download",
@@ -480,7 +473,6 @@ func TestSaveContentToOutputDirRejectsOverwriteWithoutFlag(t *testing.T) {
func TestDriveTaskResultExportIncludesReadyFlags(t *testing.T) {
f, stdout, _, reg := cmdutil.TestFactory(t, driveTestConfig())
registerDriveBotTokenStub(reg)
reg.Register(&httpmock.Stub{
Method: "GET",
URL: "/open-apis/drive/v1/export_tasks/tk_export",

6
shortcuts/drive/drive_import.go
View File

@@ -7,7 +7,6 @@ import (
"context"
"fmt"
"path/filepath"
"strings"
"github.com/larksuite/cli/internal/vfs"
@@ -147,9 +146,8 @@ func preflightDriveImportFile(spec *driveImportSpec) (int64, error) {
if err != nil {
return 0, output.ErrValidation("unsafe file path: %s", err)
}
spec.FilePath = safeFilePath
info, err := vfs.Stat(spec.FilePath)
info, err := vfs.Stat(safeFilePath)
if err != nil {
return 0, output.ErrValidation("cannot read file: %s", err)
}
@@ -168,7 +166,7 @@ func appendDriveImportUploadDryRun(dry *common.DryRunAPI, spec driveImportSpec,
extra = fmt.Sprintf(`{"obj_type":"%s","file_extension":"%s"}`, spec.DocType, spec.FileExtension())
}
if fileSize > maxDriveUploadFileSize {
if fileSize > common.MaxDriveMediaUploadSinglePartSize {
dry.POST("/open-apis/drive/v1/medias/upload_prepare").
Desc("[1a] Initialize multipart upload").
Body(map[string]interface{}{

215
shortcuts/drive/drive_import_common.go
View File

@@ -4,21 +4,15 @@
package drive
import (
"bytes"
"context"
"encoding/json"
"errors"
"fmt"
"io"
"net/http"
"path/filepath"
"strings"
"time"
"github.com/larksuite/cli/internal/vfs"
larkcore "github.com/larksuite/oapi-sdk-go/v3/core"
"github.com/larksuite/cli/internal/output"
"github.com/larksuite/cli/internal/validate"
"github.com/larksuite/cli/shortcuts/common"
@@ -37,12 +31,6 @@ const (
driveImport800MBFileSizeLimit int64 = 800 * 1024 * 1024
)
type driveMultipartUploadSession struct {
UploadID string
BlockSize int
BlockNum int
}
// driveImportExtToDocTypes defines which source file extensions can be imported
// into which Drive-native document types.
var driveImportExtToDocTypes = map[string][]string{
@@ -106,163 +94,41 @@ func uploadMediaForImport(ctx context.Context, runtime *common.RuntimeContext, f
if err = validateDriveImportFileSize(filePath, docType, fileSize); err != nil {
return "", err
}
fileSizeValue, err := driveUploadSizeValue(fileSize)
if err != nil {
return "", err
}
extra, err := buildImportMediaExtra(filePath, docType)
if err != nil {
return "", err
}
if fileSize <= maxDriveUploadFileSize {
if fileSize <= common.MaxDriveMediaUploadSinglePartSize {
fmt.Fprintf(runtime.IO().ErrOut, "Uploading media for import: %s (%s)\n", fileName, common.FormatSize(fileSize))
return uploadMediaForImportAll(runtime, filePath, fileName, fileSizeValue, extra)
// upload_all for import works without parent_node; omitting it preserves
// the existing root-level import staging behavior.
return common.UploadDriveMediaAll(runtime, common.DriveMediaUploadAllConfig{
FilePath: filePath,
FileName: fileName,
FileSize: fileSize,
ParentType: "ccm_import_open",
Extra: extra,
})
}
fmt.Fprintf(runtime.IO().ErrOut, "Uploading media for import via multipart upload: %s (%s)\n", fileName, common.FormatSize(fileSize))
return uploadMediaForImportMultipart(runtime, filePath, fileName, fileSizeValue, extra)
}
func uploadMediaForImportAll(runtime *common.RuntimeContext, filePath, fileName string, fileSize int, extra string) (string, error) {
f, err := vfs.Open(filePath)
if err != nil {
return "", output.ErrValidation("cannot read file: %s", err)
}
defer f.Close()
fd := larkcore.NewFormdata()
fd.AddField("file_name", fileName)
fd.AddField("parent_type", "ccm_import_open")
fd.AddField("size", fmt.Sprintf("%d", fileSize))
fd.AddField("extra", extra)
fd.AddFile("file", f)
apiResp, err := runtime.DoAPI(&larkcore.ApiReq{
HttpMethod: http.MethodPost,
ApiPath: "/open-apis/drive/v1/medias/upload_all",
Body: fd,
}, larkcore.WithFileUpload())
if err != nil {
return "", wrapDriveUploadRequestError(err, "upload media failed")
}
data, err := parseDriveUploadResponse(apiResp, "upload media failed")
if err != nil {
return "", err
}
return extractDriveUploadFileToken(data, "upload media failed")
}
func uploadMediaForImportMultipart(runtime *common.RuntimeContext, filePath, fileName string, fileSize int, extra string) (string, error) {
session, err := prepareMediaImportUpload(runtime, fileName, fileSize, extra)
if err != nil {
fmt.Fprintf(runtime.IO().ErrOut, "Multipart upload prepare failed: %s\n", err)
return "", err
}
totalBlocks := session.BlockNum
fmt.Fprintf(runtime.IO().ErrOut, "Multipart upload initialized: %d chunks x %s\n", totalBlocks, common.FormatSize(int64(session.BlockSize)))
f, err := vfs.Open(filePath)
if err != nil {
return "", output.ErrValidation("cannot read file: %s", err)
}
defer f.Close()
buffer := make([]byte, session.BlockSize)
remaining := fileSize
uploadedBlocks := 0
for remaining > 0 {
chunkSize := session.BlockSize
if chunkSize > remaining {
chunkSize = remaining
}
n, readErr := io.ReadFull(f, buffer[:chunkSize])
if readErr != nil {
return "", output.ErrValidation("cannot read file: %s", readErr)
}
if err = uploadMediaImportPart(runtime, session.UploadID, uploadedBlocks, buffer[:n]); err != nil {
fmt.Fprintf(runtime.IO().ErrOut, "Multipart upload part failed: %s\n", err)
return "", err
}
remaining -= n
uploadedBlocks++
}
if session.BlockNum > 0 && session.BlockNum != uploadedBlocks {
return "", output.Errorf(output.ExitAPI, "api_error", "upload prepare mismatch: expected %d blocks, uploaded %d", session.BlockNum, uploadedBlocks)
}
return finishMediaImportUpload(runtime, session.UploadID, uploadedBlocks)
}
func prepareMediaImportUpload(runtime *common.RuntimeContext, fileName string, fileSize int, extra string) (driveMultipartUploadSession, error) {
data, err := runtime.CallAPI("POST", "/open-apis/drive/v1/medias/upload_prepare", nil, map[string]interface{}{
"file_name": fileName,
"parent_type": "ccm_import_open", // For media import uploads, parent_type must be ccm_import_open.
"size": fileSize,
"extra": extra,
"parent_node": "", // For media import uploads, parent_node must be an explicit empty string; unlike medias/upload_all, this field cannot be omitted.
// upload_prepare is stricter than upload_all here and expects parent_node to
// be sent explicitly, even when import uses the implicit root staging area.
return common.UploadDriveMediaMultipart(runtime, common.DriveMediaMultipartUploadConfig{
FilePath: filePath,
FileName: fileName,
FileSize: fileSize,
ParentType: "ccm_import_open",
ParentNode: "",
Extra: extra,
})
if err != nil {
return driveMultipartUploadSession{}, err
}
session := driveMultipartUploadSession{
UploadID: common.GetString(data, "upload_id"),
BlockSize: int(common.GetFloat(data, "block_size")),
BlockNum: int(common.GetFloat(data, "block_num")),
}
if session.UploadID == "" {
return driveMultipartUploadSession{}, output.Errorf(output.ExitAPI, "api_error", "upload prepare failed: no upload_id returned")
}
if session.BlockSize <= 0 {
return driveMultipartUploadSession{}, output.Errorf(output.ExitAPI, "api_error", "upload prepare failed: invalid block_size returned")
}
if session.BlockNum <= 0 {
return driveMultipartUploadSession{}, output.Errorf(output.ExitAPI, "api_error", "upload prepare failed: invalid block_num returned")
}
return session, nil
}
func uploadMediaImportPart(runtime *common.RuntimeContext, uploadID string, seq int, chunk []byte) error {
fd := larkcore.NewFormdata()
fd.AddField("upload_id", uploadID)
fd.AddField("seq", seq)
fd.AddField("size", len(chunk))
fd.AddFile("file", bytes.NewReader(chunk))
apiResp, err := runtime.DoAPI(&larkcore.ApiReq{
HttpMethod: http.MethodPost,
ApiPath: "/open-apis/drive/v1/medias/upload_part",
Body: fd,
}, larkcore.WithFileUpload())
if err != nil {
return wrapDriveUploadRequestError(err, "upload media part failed")
}
_, err = parseDriveUploadResponse(apiResp, "upload media part failed")
return err
}
func finishMediaImportUpload(runtime *common.RuntimeContext, uploadID string, blockNum int) (string, error) {
data, err := runtime.CallAPI("POST", "/open-apis/drive/v1/medias/upload_finish", nil, map[string]interface{}{
"upload_id": uploadID,
"block_num": blockNum,
})
if err != nil {
fmt.Fprintf(runtime.IO().ErrOut, "Multipart upload finish failed: %s\n", err)
return "", err
}
return extractDriveUploadFileToken(data, "upload media finish failed")
}
func buildImportMediaExtra(filePath, docType string) (string, error) {
// The import media endpoint uses extra to decide both the target native type
// and how to interpret the uploaded source file.
extraBytes, err := json.Marshal(map[string]string{
"obj_type": docType,
"file_extension": strings.TrimPrefix(strings.ToLower(filepath.Ext(filePath)), "."),
@@ -318,45 +184,6 @@ func validateDriveImportFileSize(filePath, docType string, fileSize int64) error
)
}
func driveUploadSizeValue(fileSize int64) (int, error) {
maxInt := int64(^uint(0) >> 1)
if fileSize > maxInt {
return 0, output.ErrValidation("file %s is too large to upload", common.FormatSize(fileSize))
}
return int(fileSize), nil
}
func wrapDriveUploadRequestError(err error, action string) error {
var exitErr *output.ExitError
if errors.As(err, &exitErr) {
return err
}
return output.ErrNetwork("%s: %v", action, err)
}
func parseDriveUploadResponse(apiResp *larkcore.ApiResp, action string) (map[string]interface{}, error) {
var result map[string]interface{}
if err := json.Unmarshal(apiResp.RawBody, &result); err != nil {
return nil, output.Errorf(output.ExitAPI, "api_error", "%s: invalid response JSON: %v", action, err)
}
if larkCode := int(common.GetFloat(result, "code")); larkCode != 0 {
msg, _ := result["msg"].(string)
return nil, output.ErrAPI(larkCode, fmt.Sprintf("%s: [%d] %s", action, larkCode, msg), result["error"])
}
data, _ := result["data"].(map[string]interface{})
return data, nil
}
func extractDriveUploadFileToken(data map[string]interface{}, action string) (string, error) {
fileToken := common.GetString(data, "file_token")
if fileToken == "" {
return "", output.Errorf(output.ExitAPI, "api_error", "%s: no file_token returned", action)
}
return fileToken, nil
}
// validateDriveImportSpec enforces the CLI-level compatibility rules before any
// upload or import request is sent to the backend.
func validateDriveImportSpec(spec driveImportSpec) error {

404
shortcuts/drive/drive_import_common_test.go
View File

@@ -5,20 +5,12 @@ package drive
import (
"bytes"
"encoding/json"
"errors"
"io"
"mime"
"mime/multipart"
"os"
"strings"
"testing"
larkcore "github.com/larksuite/oapi-sdk-go/v3/core"
"github.com/larksuite/cli/internal/cmdutil"
"github.com/larksuite/cli/internal/httpmock"
"github.com/larksuite/cli/internal/output"
)
func TestValidateDriveImportSpecRejectsMismatchedType(t *testing.T) {
@@ -144,7 +136,6 @@ func TestDriveImportStatusPendingWithoutToken(t *testing.T) {
func TestDriveImportTimeoutReturnsFollowUpCommand(t *testing.T) {
f, stdout, _, reg := cmdutil.TestFactory(t, driveTestConfig())
registerDriveBotTokenStub(reg)
reg.Register(&httpmock.Stub{
Method: "POST",
URL: "/open-apis/drive/v1/medias/upload_all",
@@ -207,295 +198,6 @@ func TestDriveImportTimeoutReturnsFollowUpCommand(t *testing.T) {
}
}
func TestDriveImportUsesMultipartUploadForLargeFile(t *testing.T) {
f, stdout, _, reg := cmdutil.TestFactory(t, driveTestConfig())
registerDriveBotTokenStub(reg)
prepareStub := &httpmock.Stub{
Method: "POST",
URL: "/open-apis/drive/v1/medias/upload_prepare",
Body: map[string]interface{}{
"code": 0,
"data": map[string]interface{}{
"upload_id": "upload_123",
"block_size": 4 * 1024 * 1024,
"block_num": 6,
},
},
}
reg.Register(prepareStub)
partStubs := make([]*httpmock.Stub, 0, 6)
for i := 0; i < 6; i++ {
stub := &httpmock.Stub{
Method: "POST",
URL: "/open-apis/drive/v1/medias/upload_part",
Body: map[string]interface{}{
"code": 0,
"msg": "ok",
},
}
partStubs = append(partStubs, stub)
reg.Register(stub)
}
finishStub := &httpmock.Stub{
Method: "POST",
URL: "/open-apis/drive/v1/medias/upload_finish",
Body: map[string]interface{}{
"code": 0,
"data": map[string]interface{}{
"file_token": "file_123",
},
},
}
reg.Register(finishStub)
reg.Register(&httpmock.Stub{
Method: "POST",
URL: "/open-apis/drive/v1/import_tasks",
Body: map[string]interface{}{
"code": 0,
"data": map[string]interface{}{"ticket": "tk_import"},
},
})
reg.Register(&httpmock.Stub{
Method: "GET",
URL: "/open-apis/drive/v1/import_tasks/tk_import",
Body: map[string]interface{}{
"code": 0,
"data": map[string]interface{}{
"result": map[string]interface{}{
"type": "sheet",
"job_status": 0,
"token": "sheet_123",
"url": "https://example.com/sheets/sheet_123",
},
},
},
})
tmpDir := t.TempDir()
withDriveWorkingDir(t, tmpDir)
writeSizedDriveImportFile(t, "large.xlsx", int64(maxDriveUploadFileSize)+1)
err := mountAndRunDrive(t, DriveImport, []string{
"+import",
"--file", "large.xlsx",
"--type", "sheet",
"--as", "bot",
}, f, stdout)
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if !bytes.Contains(stdout.Bytes(), []byte(`"token": "sheet_123"`)) {
t.Fatalf("stdout missing imported token: %s", stdout.String())
}
prepareBody := decodeCapturedJSONBody(t, prepareStub)
if got, _ := prepareBody["parent_type"].(string); got != "ccm_import_open" {
t.Fatalf("prepare parent_type = %q, want %q", got, "ccm_import_open")
}
if got, _ := prepareBody["file_name"].(string); got != "large.xlsx" {
t.Fatalf("prepare file_name = %q, want %q", got, "large.xlsx")
}
if got, _ := prepareBody["size"].(float64); got != float64(maxDriveUploadFileSize+1) {
t.Fatalf("prepare size = %v, want %d", got, maxDriveUploadFileSize+1)
}
firstPart := decodeCapturedMultipartBody(t, partStubs[0])
if got := firstPart.Fields["upload_id"]; got != "upload_123" {
t.Fatalf("first part upload_id = %q, want %q", got, "upload_123")
}
if got := firstPart.Fields["seq"]; got != "0" {
t.Fatalf("first part seq = %q, want %q", got, "0")
}
if got := firstPart.Fields["size"]; got != "4194304" {
t.Fatalf("first part size = %q, want %q", got, "4194304")
}
if got := len(firstPart.Files["file"]); got != 4*1024*1024 {
t.Fatalf("first part file size = %d, want %d", got, 4*1024*1024)
}
lastPart := decodeCapturedMultipartBody(t, partStubs[len(partStubs)-1])
if got := lastPart.Fields["seq"]; got != "5" {
t.Fatalf("last part seq = %q, want %q", got, "5")
}
if got := lastPart.Fields["size"]; got != "1" {
t.Fatalf("last part size = %q, want %q", got, "1")
}
if got := len(lastPart.Files["file"]); got != 1 {
t.Fatalf("last part file size = %d, want %d", got, 1)
}
finishBody := decodeCapturedJSONBody(t, finishStub)
if got, _ := finishBody["upload_id"].(string); got != "upload_123" {
t.Fatalf("finish upload_id = %q, want %q", got, "upload_123")
}
if got, _ := finishBody["block_num"].(float64); got != 6 {
t.Fatalf("finish block_num = %v, want %d", got, 6)
}
}
func TestDriveImportMultipartPrepareValidatesResponseFields(t *testing.T) {
tests := []struct {
name string
data map[string]interface{}
wantText string
}{
{
name: "missing upload id",
data: map[string]interface{}{
"block_size": 4 * 1024 * 1024,
"block_num": 6,
},
wantText: "upload prepare failed: no upload_id returned",
},
{
name: "missing block size",
data: map[string]interface{}{
"upload_id": "upload_123",
"block_num": 6,
},
wantText: "upload prepare failed: invalid block_size returned",
},
{
name: "missing block num",
data: map[string]interface{}{
"upload_id": "upload_123",
"block_size": 4 * 1024 * 1024,
},
wantText: "upload prepare failed: invalid block_num returned",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
f, _, _, reg := cmdutil.TestFactory(t, driveTestConfig())
registerDriveBotTokenStub(reg)
reg.Register(&httpmock.Stub{
Method: "POST",
URL: "/open-apis/drive/v1/medias/upload_prepare",
Body: map[string]interface{}{
"code": 0,
"data": tt.data,
},
})
tmpDir := t.TempDir()
withDriveWorkingDir(t, tmpDir)
writeSizedDriveImportFile(t, "large.xlsx", int64(maxDriveUploadFileSize)+1)
err := mountAndRunDrive(t, DriveImport, []string{
"+import",
"--file", "large.xlsx",
"--type", "sheet",
"--as", "bot",
}, f, nil)
if err == nil {
t.Fatal("expected error, got nil")
}
if !strings.Contains(err.Error(), tt.wantText) {
t.Fatalf("error = %v, want substring %q", err, tt.wantText)
}
})
}
}
func TestDriveImportMultipartUploadPartAPIFailure(t *testing.T) {
f, _, _, reg := cmdutil.TestFactory(t, driveTestConfig())
registerDriveBotTokenStub(reg)
reg.Register(&httpmock.Stub{
Method: "POST",
URL: "/open-apis/drive/v1/medias/upload_prepare",
Body: map[string]interface{}{
"code": 0,
"data": map[string]interface{}{
"upload_id": "upload_123",
"block_size": 4 * 1024 * 1024,
"block_num": 6,
},
},
})
reg.Register(&httpmock.Stub{
Method: "POST",
URL: "/open-apis/drive/v1/medias/upload_part",
Body: map[string]interface{}{
"code": 999,
"msg": "chunk rejected",
},
})
tmpDir := t.TempDir()
withDriveWorkingDir(t, tmpDir)
writeSizedDriveImportFile(t, "large.xlsx", int64(maxDriveUploadFileSize)+1)
err := mountAndRunDrive(t, DriveImport, []string{
"+import",
"--file", "large.xlsx",
"--type", "sheet",
"--as", "bot",
}, f, nil)
if err == nil {
t.Fatal("expected error, got nil")
}
if !strings.Contains(err.Error(), "upload media part failed: [999] chunk rejected") {
t.Fatalf("unexpected error: %v", err)
}
}
func TestDriveImportMultipartFinishRequiresFileToken(t *testing.T) {
f, _, _, reg := cmdutil.TestFactory(t, driveTestConfig())
registerDriveBotTokenStub(reg)
reg.Register(&httpmock.Stub{
Method: "POST",
URL: "/open-apis/drive/v1/medias/upload_prepare",
Body: map[string]interface{}{
"code": 0,
"data": map[string]interface{}{
"upload_id": "upload_123",
"block_size": 4 * 1024 * 1024,
"block_num": 6,
},
},
})
for i := 0; i < 6; i++ {
reg.Register(&httpmock.Stub{
Method: "POST",
URL: "/open-apis/drive/v1/medias/upload_part",
Body: map[string]interface{}{
"code": 0,
"msg": "ok",
},
})
}
reg.Register(&httpmock.Stub{
Method: "POST",
URL: "/open-apis/drive/v1/medias/upload_finish",
Body: map[string]interface{}{
"code": 0,
"data": map[string]interface{}{},
},
})
tmpDir := t.TempDir()
withDriveWorkingDir(t, tmpDir)
writeSizedDriveImportFile(t, "large.xlsx", int64(maxDriveUploadFileSize)+1)
err := mountAndRunDrive(t, DriveImport, []string{
"+import",
"--file", "large.xlsx",
"--type", "sheet",
"--as", "bot",
}, f, nil)
if err == nil {
t.Fatal("expected error, got nil")
}
if !strings.Contains(err.Error(), "upload media finish failed: no file_token returned") {
t.Fatalf("unexpected error: %v", err)
}
}
func TestDriveImportRejectsOversizedFileByImportLimit(t *testing.T) {
f, _, _, _ := cmdutil.TestFactory(t, driveTestConfig())
@@ -517,73 +219,6 @@ func TestDriveImportRejectsOversizedFileByImportLimit(t *testing.T) {
}
}
func TestParseDriveUploadResponseErrors(t *testing.T) {
t.Parallel()
t.Run("invalid json", func(t *testing.T) {
t.Parallel()
_, err := parseDriveUploadResponse(&larkcore.ApiResp{RawBody: []byte("{")}, "upload media failed")
if err == nil || !strings.Contains(err.Error(), "invalid response JSON") {
t.Fatalf("expected invalid JSON error, got %v", err)
}
})
t.Run("api code error", func(t *testing.T) {
t.Parallel()
_, err := parseDriveUploadResponse(&larkcore.ApiResp{RawBody: []byte(`{"code":999,"msg":"boom","error":{"detail":"x"}}`)}, "upload media failed")
if err == nil || !strings.Contains(err.Error(), "upload media failed: [999] boom") {
t.Fatalf("expected API error, got %v", err)
}
})
}
func TestWrapDriveUploadRequestError(t *testing.T) {
t.Parallel()
t.Run("preserves exit error", func(t *testing.T) {
t.Parallel()
original := output.ErrValidation("bad input")
got := wrapDriveUploadRequestError(original, "upload media failed")
if got != original {
t.Fatalf("expected same exit error pointer, got %v", got)
}
})
t.Run("wraps generic error as network", func(t *testing.T) {
t.Parallel()
got := wrapDriveUploadRequestError(io.EOF, "upload media failed")
var exitErr *output.ExitError
if !errors.As(got, &exitErr) {
t.Fatalf("expected ExitError, got %T", got)
}
if exitErr.Code != output.ExitNetwork {
t.Fatalf("exit code = %d, want %d", exitErr.Code, output.ExitNetwork)
}
if !strings.Contains(got.Error(), "upload media failed") {
t.Fatalf("unexpected error: %v", got)
}
})
}
type capturedMultipartBody struct {
Fields map[string]string
Files map[string][]byte
}
func decodeCapturedJSONBody(t *testing.T, stub *httpmock.Stub) map[string]interface{} {
t.Helper()
var body map[string]interface{}
if err := json.Unmarshal(stub.CapturedBody, &body); err != nil {
t.Fatalf("decode captured JSON body: %v", err)
}
return body
}
func writeSizedDriveImportFile(t *testing.T, name string, size int64) {
t.Helper()
@@ -598,42 +233,3 @@ func writeSizedDriveImportFile(t *testing.T, name string, size int64) {
t.Fatalf("Close(%q) error: %v", name, err)
}
}
func decodeCapturedMultipartBody(t *testing.T, stub *httpmock.Stub) capturedMultipartBody {
t.Helper()
contentType := stub.CapturedHeaders.Get("Content-Type")
mediaType, params, err := mime.ParseMediaType(contentType)
if err != nil {
t.Fatalf("parse multipart content type: %v", err)
}
if mediaType != "multipart/form-data" {
t.Fatalf("content type = %q, want multipart/form-data", mediaType)
}
reader := multipart.NewReader(bytes.NewReader(stub.CapturedBody), params["boundary"])
body := capturedMultipartBody{
Fields: map[string]string{},
Files: map[string][]byte{},
}
for {
part, err := reader.NextPart()
if err == io.EOF {
break
}
if err != nil {
t.Fatalf("read multipart part: %v", err)
}
data, err := io.ReadAll(part)
if err != nil {
t.Fatalf("read multipart data: %v", err)
}
if part.FileName() != "" {
body.Files[part.FormName()] = data
continue
}
body.Fields[part.FormName()] = string(data)
}
return body
}

2
shortcuts/drive/drive_import_test.go
View File

@@ -132,7 +132,7 @@ func TestDriveImportDryRunShowsMultipartUploadForLargeFile(t *testing.T) {
if err != nil {
t.Fatalf("Create() error: %v", err)
}
if err := fh.Truncate(int64(maxDriveUploadFileSize) + 1); err != nil {
if err := fh.Truncate(common.MaxDriveMediaUploadSinglePartSize + 1); err != nil {
t.Fatalf("Truncate() error: %v", err)
}
if err := fh.Close(); err != nil {

31
shortcuts/drive/drive_io_test.go
View File

@@ -18,9 +18,6 @@ import (
"github.com/larksuite/cli/shortcuts/common"
)
// registerDriveBotTokenStub is a no-op. TAT is now managed by CredentialProvider, not SDK.
func registerDriveBotTokenStub(_ *httpmock.Registry) {}
func driveTestConfig() *core.CliConfig {
return &core.CliConfig{
AppID: "drive-test-app", AppSecret: "test-secret", Brand: core.BrandFeishu,
@@ -62,7 +59,6 @@ func TestDriveUploadLargeFileUsesMultipart(t *testing.T) {
AppID: "drive-upload-test-app", AppSecret: "test-secret", Brand: core.BrandFeishu,
}
f, stdout, _, reg := cmdutil.TestFactory(t, uploadTestConfig)
registerDriveBotTokenStub(reg)
// Step 1: upload_prepare
reg.Register(&httpmock.Stub{
@@ -72,7 +68,7 @@ func TestDriveUploadLargeFileUsesMultipart(t *testing.T) {
"code": 0, "msg": "ok",
"data": map[string]interface{}{
"upload_id": "test-upload-id",
"block_size": float64(maxDriveUploadFileSize),
"block_size": float64(common.MaxDriveMediaUploadSinglePartSize),
"block_num": float64(2),
},
},
@@ -116,7 +112,7 @@ func TestDriveUploadLargeFileUsesMultipart(t *testing.T) {
if err != nil {
t.Fatalf("Create() error: %v", err)
}
if err := fh.Truncate(maxDriveUploadFileSize + 1); err != nil {
if err := fh.Truncate(common.MaxDriveMediaUploadSinglePartSize + 1); err != nil {
t.Fatalf("Truncate() error: %v", err)
}
if err := fh.Close(); err != nil {
@@ -141,7 +137,6 @@ func TestDriveUploadSmallFile(t *testing.T) {
AppID: "drive-upload-small-test", AppSecret: "test-secret", Brand: core.BrandFeishu,
}
f, stdout, _, reg := cmdutil.TestFactory(t, uploadTestConfig)
registerDriveBotTokenStub(reg)
reg.Register(&httpmock.Stub{
Method: "POST",
@@ -181,7 +176,6 @@ func TestDriveUploadSmallFileAPIError(t *testing.T) {
AppID: "drive-upload-small-err", AppSecret: "test-secret", Brand: core.BrandFeishu,
}
f, stdout, _, reg := cmdutil.TestFactory(t, uploadTestConfig)
registerDriveBotTokenStub(reg)
reg.Register(&httpmock.Stub{
Method: "POST",
@@ -218,7 +212,6 @@ func TestDriveUploadSmallFileNoToken(t *testing.T) {
AppID: "drive-upload-small-notoken", AppSecret: "test-secret", Brand: core.BrandFeishu,
}
f, stdout, _, reg := cmdutil.TestFactory(t, uploadTestConfig)
registerDriveBotTokenStub(reg)
reg.Register(&httpmock.Stub{
Method: "POST",
@@ -256,7 +249,6 @@ func TestDriveUploadSmallFileInvalidJSON(t *testing.T) {
AppID: "drive-upload-small-json", AppSecret: "test-secret", Brand: core.BrandFeishu,
}
f, stdout, _, reg := cmdutil.TestFactory(t, uploadTestConfig)
registerDriveBotTokenStub(reg)
reg.Register(&httpmock.Stub{
Method: "POST",
@@ -291,7 +283,6 @@ func TestDriveUploadPrepareInvalidResponse(t *testing.T) {
AppID: "drive-upload-prepare-bad", AppSecret: "test-secret", Brand: core.BrandFeishu,
}
f, stdout, _, reg := cmdutil.TestFactory(t, uploadTestConfig)
registerDriveBotTokenStub(reg)
reg.Register(&httpmock.Stub{
Method: "POST",
@@ -317,7 +308,7 @@ func TestDriveUploadPrepareInvalidResponse(t *testing.T) {
if err != nil {
t.Fatalf("Create() error: %v", err)
}
if err := fh.Truncate(maxDriveUploadFileSize + 1); err != nil {
if err := fh.Truncate(common.MaxDriveMediaUploadSinglePartSize + 1); err != nil {
t.Fatalf("Truncate() error: %v", err)
}
fh.Close()
@@ -338,7 +329,6 @@ func TestDriveUploadPartAPIError(t *testing.T) {
AppID: "drive-upload-part-err", AppSecret: "test-secret", Brand: core.BrandFeishu,
}
f, stdout, _, reg := cmdutil.TestFactory(t, uploadTestConfig)
registerDriveBotTokenStub(reg)
reg.Register(&httpmock.Stub{
Method: "POST",
@@ -347,7 +337,7 @@ func TestDriveUploadPartAPIError(t *testing.T) {
"code": 0, "msg": "ok",
"data": map[string]interface{}{
"upload_id": "test-upload-id",
"block_size": float64(maxDriveUploadFileSize),
"block_size": float64(common.MaxDriveMediaUploadSinglePartSize),
"block_num": float64(2),
},
},
@@ -380,7 +370,7 @@ func TestDriveUploadPartAPIError(t *testing.T) {
if err != nil {
t.Fatalf("Create() error: %v", err)
}
if err := fh.Truncate(maxDriveUploadFileSize + 1); err != nil {
if err := fh.Truncate(common.MaxDriveMediaUploadSinglePartSize + 1); err != nil {
t.Fatalf("Truncate() error: %v", err)
}
fh.Close()
@@ -401,7 +391,6 @@ func TestDriveUploadPartInvalidJSON(t *testing.T) {
AppID: "drive-upload-part-json", AppSecret: "test-secret", Brand: core.BrandFeishu,
}
f, stdout, _, reg := cmdutil.TestFactory(t, uploadTestConfig)
registerDriveBotTokenStub(reg)
reg.Register(&httpmock.Stub{
Method: "POST",
@@ -410,7 +399,7 @@ func TestDriveUploadPartInvalidJSON(t *testing.T) {
"code": 0, "msg": "ok",
"data": map[string]interface{}{
"upload_id": "test-upload-id",
"block_size": float64(maxDriveUploadFileSize + 1),
"block_size": float64(common.MaxDriveMediaUploadSinglePartSize + 1),
"block_num": float64(1),
},
},
@@ -433,7 +422,7 @@ func TestDriveUploadPartInvalidJSON(t *testing.T) {
if err != nil {
t.Fatalf("Create() error: %v", err)
}
if err := fh.Truncate(maxDriveUploadFileSize + 1); err != nil {
if err := fh.Truncate(common.MaxDriveMediaUploadSinglePartSize + 1); err != nil {
t.Fatalf("Truncate() error: %v", err)
}
fh.Close()
@@ -454,7 +443,6 @@ func TestDriveUploadFinishNoToken(t *testing.T) {
AppID: "drive-upload-finish-notoken", AppSecret: "test-secret", Brand: core.BrandFeishu,
}
f, stdout, _, reg := cmdutil.TestFactory(t, uploadTestConfig)
registerDriveBotTokenStub(reg)
reg.Register(&httpmock.Stub{
Method: "POST",
@@ -463,7 +451,7 @@ func TestDriveUploadFinishNoToken(t *testing.T) {
"code": 0, "msg": "ok",
"data": map[string]interface{}{
"upload_id": "test-upload-id",
"block_size": float64(maxDriveUploadFileSize + 1),
"block_size": float64(common.MaxDriveMediaUploadSinglePartSize + 1),
"block_num": float64(1),
},
},
@@ -495,7 +483,7 @@ func TestDriveUploadFinishNoToken(t *testing.T) {
if err != nil {
t.Fatalf("Create() error: %v", err)
}
if err := fh.Truncate(maxDriveUploadFileSize + 1); err != nil {
if err := fh.Truncate(common.MaxDriveMediaUploadSinglePartSize + 1); err != nil {
t.Fatalf("Truncate() error: %v", err)
}
fh.Close()
@@ -516,7 +504,6 @@ func TestDriveUploadWithCustomName(t *testing.T) {
AppID: "drive-upload-name-test", AppSecret: "test-secret", Brand: core.BrandFeishu,
}
f, stdout, _, reg := cmdutil.TestFactory(t, uploadTestConfig)
registerDriveBotTokenStub(reg)
reg.Register(&httpmock.Stub{
Method: "POST",

2
shortcuts/drive/drive_move_common_test.go
View File

@@ -104,7 +104,6 @@ func TestDriveMoveDryRunFolderIncludesTaskCheckParams(t *testing.T) {
func TestDriveMoveFolderSuccessUsesTaskCheckHelper(t *testing.T) {
f, stdout, _, reg := cmdutil.TestFactory(t, driveTestConfig())
registerDriveBotTokenStub(reg)
reg.Register(&httpmock.Stub{
Method: "POST",
URL: "/open-apis/drive/v1/files/fld_src/move",
@@ -148,7 +147,6 @@ func TestDriveMoveFolderSuccessUsesTaskCheckHelper(t *testing.T) {
func TestDriveMoveFolderTimeoutReturnsFollowUpCommand(t *testing.T) {
f, stdout, _, reg := cmdutil.TestFactory(t, driveTestConfig())
registerDriveBotTokenStub(reg)
reg.Register(&httpmock.Stub{
Method: "POST",
URL: "/open-apis/drive/v1/files/fld_src/move",

2
shortcuts/drive/drive_move_test.go
View File

@@ -13,7 +13,6 @@ import (
func TestDriveMoveUsesRootFolderWhenFolderTokenMissing(t *testing.T) {
f, stdout, _, reg := cmdutil.TestFactory(t, driveTestConfig())
registerDriveBotTokenStub(reg)
reg.Register(&httpmock.Stub{
Method: "GET",
URL: "/open-apis/drive/explorer/v2/root_folder/meta",
@@ -52,7 +51,6 @@ func TestDriveMoveUsesRootFolderWhenFolderTokenMissing(t *testing.T) {
func TestDriveMoveRootFolderLookupRequiresToken(t *testing.T) {
f, _, _, reg := cmdutil.TestFactory(t, driveTestConfig())
registerDriveBotTokenStub(reg)
reg.Register(&httpmock.Stub{
Method: "GET",
URL: "/open-apis/drive/explorer/v2/root_folder/meta",

2
shortcuts/drive/drive_task_result_test.go
View File

@@ -127,7 +127,6 @@ func TestDriveTaskResultDryRunExportIncludesTokenParam(t *testing.T) {
func TestDriveTaskResultImportIncludesReadyFlags(t *testing.T) {
f, stdout, _, reg := cmdutil.TestFactory(t, driveTestConfig())
registerDriveBotTokenStub(reg)
reg.Register(&httpmock.Stub{
Method: "GET",
URL: "/open-apis/drive/v1/import_tasks/tk_import",
@@ -161,7 +160,6 @@ func TestDriveTaskResultImportIncludesReadyFlags(t *testing.T) {
func TestDriveTaskResultTaskCheckIncludesReadyFlags(t *testing.T) {
f, stdout, _, reg := cmdutil.TestFactory(t, driveTestConfig())
registerDriveBotTokenStub(reg)
reg.Register(&httpmock.Stub{
Method: "GET",
URL: "/open-apis/drive/v1/files/task_check",

7
shortcuts/drive/drive_upload.go
View File

@@ -10,7 +10,6 @@ import (
"fmt"
"io"
"net/http"
"os"
"path/filepath"
larkcore "github.com/larksuite/oapi-sdk-go/v3/core"
@@ -21,8 +20,6 @@ import (
"github.com/larksuite/cli/shortcuts/common"
)
const maxDriveUploadFileSize = 20 * 1024 * 1024 // 20MB
var DriveUpload = common.Shortcut{
Service: "drive",
Command: "+upload",
@@ -78,7 +75,7 @@ var DriveUpload = common.Shortcut{
fmt.Fprintf(runtime.IO().ErrOut, "Uploading: %s (%s)\n", fileName, common.FormatSize(fileSize))
var fileToken string
if fileSize > maxDriveUploadFileSize {
if fileSize > common.MaxDriveMediaUploadSinglePartSize {
fmt.Fprintf(runtime.IO().ErrOut, "File exceeds 20MB, using multipart upload\n")
fileToken, err = uploadFileMultipart(ctx, runtime, filePath, fileName, folderToken, fileSize)
} else {
@@ -183,7 +180,7 @@ func uploadFileMultipart(_ context.Context, runtime *common.RuntimeContext, file
partSize = remaining
}
partFile, err := os.Open(filePath)
partFile, err := vfs.Open(filePath)
if err != nil {
return "", output.ErrValidation("cannot open file: %v", err)
}

2
skills/lark-doc/references/lark-doc-media-insert.md
View File

@@ -30,7 +30,7 @@ lark-cli docs +media-insert --doc doxcnXXX --file ./arch.png --align center --ca
| 参数 | 必填 | 说明 |
|------|------|------|
| `--doc <id>` | 是 | 文档 ID 或 docx URL（仅支持 `/docx/<document_id>` 形式自动提取；**不支持 `/wiki/...` URL 自动提取**） |
| `--file <path>` | 是 | 本地文件路径（最大 20MB） |
| `--file <path>` | 是 | 本地文件路径（文件大于 20MB 时自动切换分片上传） |
| `--type <type>` | 否 | `image`（默认）或 `file` |
| `--align <align>` | 否 | 仅图片：`left` / `center`（默认）/ `right` |
| `--caption <text>` | 否 | 仅图片：图片描述 |