Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion shortcuts/common/drive_meta.go
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,7 @@ func FetchDriveMeta(runtime *RuntimeContext, token, docType string, withURL bool
body["with_url"] = true
}

data, err := runtime.CallAPI(
data, err := runtime.CallAPITyped(
"POST",
"/open-apis/drive/v1/metas/batch_query",
nil,
Expand Down
8 changes: 8 additions & 0 deletions shortcuts/common/drive_meta_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@ import (
"sync/atomic"
"testing"

"github.com/larksuite/cli/errs"
"github.com/larksuite/cli/internal/cmdutil"
"github.com/larksuite/cli/internal/core"
"github.com/larksuite/cli/internal/httpmock"
Expand Down Expand Up @@ -103,6 +104,13 @@ func TestFetchDriveMetaTitle(t *testing.T) {
if err == nil {
t.Fatal("FetchDriveMetaTitle() expected error, got nil")
}
p, ok := errs.ProblemOf(err)
if !ok {
t.Fatalf("expected typed error, got %T", err)
}
if p.Code != 99991668 {
t.Fatalf("code = %d, want 99991668", p.Code)
}
})
}

Expand Down
181 changes: 142 additions & 39 deletions shortcuts/drive/drive_inspect.go
Original file line number Diff line number Diff line change
Expand Up @@ -8,11 +8,19 @@
"fmt"
"io"
"strings"
"time"

"github.com/larksuite/cli/errs"
"github.com/larksuite/cli/shortcuts/common"
)

const (
driveInspectRateLimitRetries = 2
driveInspectRetryInitialBackoff = 200 * time.Millisecond
)

var driveInspectAfter = time.After

var DriveInspect = common.Shortcut{
Service: "drive",
Command: "+inspect",
Expand All @@ -35,32 +43,15 @@
},
},
Validate: func(ctx context.Context, runtime *common.RuntimeContext) error {
raw := strings.TrimSpace(runtime.Str("url"))
if raw == "" {
return errs.NewValidationError(errs.SubtypeInvalidArgument, "--url cannot be empty").WithParam("--url")
}

_, ok := common.ParseResourceURL(raw)
if !ok {
// Not a recognized URL pattern.
if strings.Contains(raw, "://") {
return errs.NewValidationError(errs.SubtypeInvalidArgument, "unsupported --url %q: use a recognized Lark document URL or a bare token with --type", raw).WithParam("--url")
}
// Bare token: --type is required.
if strings.TrimSpace(runtime.Str("type")) == "" {
return errs.NewValidationError(errs.SubtypeInvalidArgument, "--type is required when --url is a bare token (allowed: doc, docx, sheet, bitable, wiki, file, folder, mindnote, slides)").WithParam("--type")
}
if _, err := driveInspectResolveRef(runtime); err != nil {
return err
}
return nil
},
DryRun: func(ctx context.Context, runtime *common.RuntimeContext) *common.DryRunAPI {
raw := strings.TrimSpace(runtime.Str("url"))
ref, ok := common.ParseResourceURL(raw)
if !ok {
ref = common.ResourceRef{
Type: strings.TrimSpace(runtime.Str("type")),
Token: raw,
}
ref, err := driveInspectResolveRef(runtime)
if err != nil {
return common.NewDryRunAPI()

Check warning on line 54 in shortcuts/drive/drive_inspect.go

View check run for this annotation

Codecov / codecov/patch

shortcuts/drive/drive_inspect.go#L54

Added line #L54 was not covered by tests
}

dry := common.NewDryRunAPI()
Expand Down Expand Up @@ -91,15 +82,9 @@
},
Execute: func(ctx context.Context, runtime *common.RuntimeContext) error {
raw := strings.TrimSpace(runtime.Str("url"))

// Step 1: Parse URL to extract {type, token}.
ref, ok := common.ParseResourceURL(raw)
if !ok {
// Bare token: use --type.
ref = common.ResourceRef{
Type: strings.TrimSpace(runtime.Str("type")),
Token: raw,
}
ref, err := driveInspectResolveRef(runtime)
if err != nil {
return err

Check warning on line 87 in shortcuts/drive/drive_inspect.go

View check run for this annotation

Codecov / codecov/patch

shortcuts/drive/drive_inspect.go#L87

Added line #L87 was not covered by tests
}

inputURL := raw
Expand All @@ -111,14 +96,19 @@
// Step 2: If type is "wiki", unwrap via get_node API.
if docType == "wiki" {
fmt.Fprintf(runtime.IO().ErrOut, "Inspecting wiki node: %s\n", common.MaskToken(docToken))
data, err := runtime.CallAPITyped(
"GET",
"/open-apis/wiki/v2/spaces/get_node",
map[string]interface{}{"token": docToken},
nil,
data, err := driveInspectCallWithRetry(
ctx,
func() (map[string]interface{}, error) {
return runtime.CallAPITyped(
"GET",
"/open-apis/wiki/v2/spaces/get_node",
map[string]interface{}{"token": docToken},
nil,
)
},
)
if err != nil {
return err
return driveInspectAnnotateError("resolve_wiki", err)

Check warning on line 111 in shortcuts/drive/drive_inspect.go

View check run for this annotation

Codecov / codecov/patch

shortcuts/drive/drive_inspect.go#L111

Added line #L111 was not covered by tests
}

node := common.GetMap(data, "node")
Expand All @@ -145,9 +135,9 @@
}

// Step 3: Call batch_query to verify and get title.
title, err := common.FetchDriveMetaTitle(runtime, docToken, docType)
title, err := driveInspectFetchMetaTitle(ctx, runtime, docToken, docType)
if err != nil {
return err
return driveInspectAnnotateError("query_meta", err)
}

// Step 4: Build the resolved URL.
Expand Down Expand Up @@ -181,3 +171,116 @@
return nil
},
}

func driveInspectResolveRef(runtime *common.RuntimeContext) (common.ResourceRef, error) {
raw := strings.TrimSpace(runtime.Str("url"))
if raw == "" {
return common.ResourceRef{}, errs.NewValidationError(errs.SubtypeInvalidArgument, "--url cannot be empty").WithParam("--url")
}

inputType := strings.ToLower(strings.TrimSpace(runtime.Str("type")))
ref, ok := common.ParseResourceURL(raw)
if ok {
if inputType != "" && inputType != ref.Type {
return common.ResourceRef{}, errs.NewValidationError(
errs.SubtypeInvalidArgument,
"--type %q conflicts with URL path type %q; remove --type or use a matching value",
inputType,
ref.Type,
).WithParam("--type")
}
return ref, nil
}

if strings.Contains(raw, "://") {
return common.ResourceRef{}, errs.NewValidationError(errs.SubtypeInvalidArgument, "unsupported --url %q: use a recognized Lark document URL or a bare token with --type", raw).WithParam("--url")
}
if strings.ContainsAny(raw, "/?#") {
return common.ResourceRef{}, errs.NewValidationError(errs.SubtypeInvalidArgument, "invalid bare token %q: remove path/query fragments and pass only the raw token with --type", raw).WithParam("--url")
}
if inputType == "" {
return common.ResourceRef{}, errs.NewValidationError(errs.SubtypeInvalidArgument, "--type is required when --url is a bare token (allowed: doc, docx, sheet, bitable, wiki, file, folder, mindnote, slides)").WithParam("--type")
}
return common.ResourceRef{Type: inputType, Token: raw}, nil
}

func driveInspectFetchMetaTitle(ctx context.Context, runtime *common.RuntimeContext, token, docType string) (string, error) {
var title string
_, err := driveInspectCallWithRetry(ctx, func() (map[string]interface{}, error) {
got, callErr := common.FetchDriveMeta(runtime, token, docType, false)
if callErr != nil {
return nil, callErr
}
title = got.Title
return map[string]interface{}{"title": got.Title}, nil
})
if err != nil {
return "", err
}
return title, nil
}

func driveInspectCallWithRetry(ctx context.Context, call func() (map[string]interface{}, error)) (map[string]interface{}, error) {
var lastErr error
for attempt := 0; attempt <= driveInspectRateLimitRetries; attempt++ {
data, err := call()
if err == nil {
return data, nil
}
lastErr = err
if !driveInspectShouldRetry(err) || attempt == driveInspectRateLimitRetries {
return nil, err
}
backoff := driveInspectRetryInitialBackoff * time.Duration(1<<attempt)
if waitErr := driveInspectWait(ctx, backoff); waitErr != nil {
return nil, waitErr

Check warning on line 236 in shortcuts/drive/drive_inspect.go

View check run for this annotation

Codecov / codecov/patch

shortcuts/drive/drive_inspect.go#L236

Added line #L236 was not covered by tests
}
}
return nil, lastErr

Check warning on line 239 in shortcuts/drive/drive_inspect.go

View check run for this annotation

Codecov / codecov/patch

shortcuts/drive/drive_inspect.go#L239

Added line #L239 was not covered by tests
}

func driveInspectShouldRetry(err error) bool {
problem, ok := errs.ProblemOf(err)
if !ok || problem == nil {
return false

Check warning on line 245 in shortcuts/drive/drive_inspect.go

View check run for this annotation

Codecov / codecov/patch

shortcuts/drive/drive_inspect.go#L245

Added line #L245 was not covered by tests
}
return problem.Subtype == errs.SubtypeRateLimit || problem.Code == 99991400 || problem.Retryable
}

func driveInspectWait(ctx context.Context, d time.Duration) error {
if d <= 0 {
return nil

Check warning on line 252 in shortcuts/drive/drive_inspect.go

View check run for this annotation

Codecov / codecov/patch

shortcuts/drive/drive_inspect.go#L252

Added line #L252 was not covered by tests
}
select {
case <-ctx.Done():
return errs.WrapInternal(ctx.Err())

Check warning on line 256 in shortcuts/drive/drive_inspect.go

View check run for this annotation

Codecov / codecov/patch

shortcuts/drive/drive_inspect.go#L255-L256

Added lines #L255 - L256 were not covered by tests
case <-driveInspectAfter(d):
return nil
}
}

func driveInspectAnnotateError(stage string, err error) error {
problem, ok := errs.ProblemOf(err)
if !ok || problem == nil {
return err

Check warning on line 265 in shortcuts/drive/drive_inspect.go

View check run for this annotation

Codecov / codecov/patch

shortcuts/drive/drive_inspect.go#L265

Added line #L265 was not covered by tests
}
label := map[string]string{
"resolve_wiki": "resolve wiki node",
"query_meta": "query document metadata",
}[stage]
if label == "" {
label = stage

Check warning on line 272 in shortcuts/drive/drive_inspect.go

View check run for this annotation

Codecov / codecov/patch

shortcuts/drive/drive_inspect.go#L272

Added line #L272 was not covered by tests
}
problem.Message = fmt.Sprintf("%s failed: %s", label, problem.Message)
if strings.TrimSpace(problem.Hint) == "" {
switch stage {
case "resolve_wiki":
problem.Hint = "check that the wiki URL/token is valid and that the current identity can read the wiki node"

Check warning on line 278 in shortcuts/drive/drive_inspect.go

View check run for this annotation

Codecov / codecov/patch

shortcuts/drive/drive_inspect.go#L277-L278

Added lines #L277 - L278 were not covered by tests
case "query_meta":
problem.Hint = "check that the resolved document still exists and that the current identity can read its metadata"
}
} else if !strings.Contains(problem.Hint, label) {
problem.Hint = label + ": " + problem.Hint

Check warning on line 283 in shortcuts/drive/drive_inspect.go

View check run for this annotation

Codecov / codecov/patch

shortcuts/drive/drive_inspect.go#L282-L283

Added lines #L282 - L283 were not covered by tests
}
return err
}
101 changes: 101 additions & 0 deletions shortcuts/drive/drive_inspect_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -6,10 +6,13 @@ package drive
import (
"context"
"encoding/json"
"strings"
"testing"
"time"

"github.com/spf13/cobra"

"github.com/larksuite/cli/errs"
"github.com/larksuite/cli/internal/cmdutil"
"github.com/larksuite/cli/internal/core"
"github.com/larksuite/cli/internal/httpmock"
Expand Down Expand Up @@ -83,6 +86,34 @@ func TestDriveInspectValidate_BareTokenWithType(t *testing.T) {
}
}

func TestDriveInspectValidate_URLTypeConflict(t *testing.T) {
cmd := &cobra.Command{Use: "drive +inspect"}
cmd.Flags().String("url", "", "")
cmd.Flags().String("type", "", "")
_ = cmd.Flags().Set("url", "https://xxx.feishu.cn/docx/doxcnBareToken")
_ = cmd.Flags().Set("type", "sheet")

runtime := common.TestNewRuntimeContext(cmd, &core.CliConfig{})
err := DriveInspect.Validate(context.Background(), runtime)
if err == nil {
t.Fatal("expected error for conflicting --type, got nil")
}
}

func TestDriveInspectValidate_BareTokenWithPathFragment(t *testing.T) {
cmd := &cobra.Command{Use: "drive +inspect"}
cmd.Flags().String("url", "", "")
cmd.Flags().String("type", "", "")
_ = cmd.Flags().Set("url", "doxcnBareToken/extra")
_ = cmd.Flags().Set("type", "docx")

runtime := common.TestNewRuntimeContext(cmd, &core.CliConfig{})
err := DriveInspect.Validate(context.Background(), runtime)
if err == nil {
t.Fatal("expected error for bare token with path fragment, got nil")
}
}

func TestDriveInspectValidate_ValidDocxURL(t *testing.T) {
cmd := &cobra.Command{Use: "drive +inspect"}
cmd.Flags().String("url", "", "")
Expand Down Expand Up @@ -540,6 +571,76 @@ func TestDriveInspectExecute_BatchQueryError(t *testing.T) {
if err == nil {
t.Fatal("expected error for batch_query failure, got nil")
}
p, ok := errs.ProblemOf(err)
if !ok {
t.Fatalf("expected typed error, got %T", err)
}
if !strings.Contains(p.Message, "query document metadata failed") {
t.Fatalf("message = %q, want query document metadata prefix", p.Message)
Comment thread
fangshuyu-768 marked this conversation as resolved.
}
}

func TestDriveInspectExecute_RetriesRateLimitOnWikiResolve(t *testing.T) {
cfg := driveTestConfig()
f, stdout, _, reg := cmdutil.TestFactory(t, cfg)

reg.Register(&httpmock.Stub{
Method: "GET",
URL: "/open-apis/wiki/v2/spaces/get_node",
Body: map[string]interface{}{
"code": 99991400,
"msg": "request trigger frequency limit",
},
})
reg.Register(&httpmock.Stub{
Method: "GET",
URL: "/open-apis/wiki/v2/spaces/get_node",
Body: map[string]interface{}{
"code": 0,
"data": map[string]interface{}{
"node": map[string]interface{}{
"obj_type": "docx",
"obj_token": "doxcnUnwrapped",
"space_id": "space123",
"node_token": "wikcnNodeToken",
},
},
},
})
reg.Register(&httpmock.Stub{
Method: "POST",
URL: "/open-apis/drive/v1/metas/batch_query",
Body: map[string]interface{}{
"code": 0,
"data": map[string]interface{}{
"metas": []map[string]interface{}{
{"doc_token": "doxcnUnwrapped", "doc_type": "docx", "title": "Wiki Doc"},
},
},
},
})

origAfter := driveInspectAfter
driveInspectAfter = func(time.Duration) <-chan time.Time {
ch := make(chan time.Time, 1)
ch <- time.Now()
return ch
}
defer func() { driveInspectAfter = origAfter }()

err := mountAndRunDrive(t, DriveInspect, []string{
"+inspect",
"--url", "https://xxx.feishu.cn/wiki/wikcnABC",
"--as", "bot",
}, f, stdout)
if err != nil {
t.Fatalf("unexpected error after retry: %v", err)
}

data := decodeDriveEnvelope(t, stdout)
if data["token"] != "doxcnUnwrapped" {
t.Fatalf("token = %v, want doxcnUnwrapped", data["token"])
}
}

func TestDriveInspectExecute_PrettyFormat(t *testing.T) {
Expand Down
Loading