WeKnora deployment source with the reviewed wiki identity, summary and recovery fixes
Anydoc / anydoc parser engine (push) Canceled after 0s
App / Format, vet, test, and build (push) Canceled after 0s
CLI / test (macos-latest / go 1.26) (push) Canceled after 0s
CLI / test (ubuntu-latest / go 1.26) (push) Canceled after 0s
CLI / test (windows-latest / go 1.26) (push) Canceled after 0s
Build and Push Docker Image / build-ui (push) Canceled after 0s
Build and Push Docker Image / build-docreader (push) Canceled after 0s
Build and Push Docker Image / build-sandbox (push) Canceled after 0s
Build and Push Docker Image / build-app (amd64, linux/amd64, ubuntu-latest) (push) Canceled after 0s
Build and Push Docker Image / build-app (arm64, linux/arm64, ubuntu-24.04-arm) (push) Canceled after 0s
DocReader / Python tests and Go client tests (push) Canceled after 0s
Build & Publish (dsh-weknora) / Typecheck, test, pack (push) Canceled after 0s
Build & Publish (dsh-weknora) / End-to-end inside dsh (pinned) (push) Canceled after 0s
Build & Publish (dsh-weknora) / End-to-end inside dsh (latest, advisory) (push) Canceled after 0s
Frontend / Test, type-check, and build (push) Canceled after 0s
Go Lint Cache / warm lint cache (.) (push) Canceled after 0s
Go Lint Cache / warm lint cache (cli) (push) Canceled after 0s
Go Lint Cache / warm lint cache (client) (push) Canceled after 0s
Build & Publish (mcp-server) / Test (Python 3.10) (push) Canceled after 0s
Build & Publish (mcp-server) / Test (Python 3.11) (push) Canceled after 0s
Build & Publish (mcp-server) / Test (Python 3.12) (push) Canceled after 0s
Build & Publish (mcp-server) / Test (Python 3.13) (push) Canceled after 0s
Build and Push Docker Image / merge (push) Canceled after 0s
Build & Publish (dsh-weknora) / Publish to npm (push) Canceled after 0s
Build & Publish (mcp-server) / Build sdist + wheel (push) Canceled after 0s
Build & Publish (mcp-server) / Publish to PyPI (push) Canceled after 0s
Anydoc / anydoc parser engine (push) Canceled after 0s
App / Format, vet, test, and build (push) Canceled after 0s
CLI / test (macos-latest / go 1.26) (push) Canceled after 0s
CLI / test (ubuntu-latest / go 1.26) (push) Canceled after 0s
CLI / test (windows-latest / go 1.26) (push) Canceled after 0s
Build and Push Docker Image / build-ui (push) Canceled after 0s
Build and Push Docker Image / build-docreader (push) Canceled after 0s
Build and Push Docker Image / build-sandbox (push) Canceled after 0s
Build and Push Docker Image / build-app (amd64, linux/amd64, ubuntu-latest) (push) Canceled after 0s
Build and Push Docker Image / build-app (arm64, linux/arm64, ubuntu-24.04-arm) (push) Canceled after 0s
DocReader / Python tests and Go client tests (push) Canceled after 0s
Build & Publish (dsh-weknora) / Typecheck, test, pack (push) Canceled after 0s
Build & Publish (dsh-weknora) / End-to-end inside dsh (pinned) (push) Canceled after 0s
Build & Publish (dsh-weknora) / End-to-end inside dsh (latest, advisory) (push) Canceled after 0s
Frontend / Test, type-check, and build (push) Canceled after 0s
Go Lint Cache / warm lint cache (.) (push) Canceled after 0s
Go Lint Cache / warm lint cache (cli) (push) Canceled after 0s
Go Lint Cache / warm lint cache (client) (push) Canceled after 0s
Build & Publish (mcp-server) / Test (Python 3.10) (push) Canceled after 0s
Build & Publish (mcp-server) / Test (Python 3.11) (push) Canceled after 0s
Build & Publish (mcp-server) / Test (Python 3.12) (push) Canceled after 0s
Build & Publish (mcp-server) / Test (Python 3.13) (push) Canceled after 0s
Build and Push Docker Image / merge (push) Canceled after 0s
Build & Publish (dsh-weknora) / Publish to npm (push) Canceled after 0s
Build & Publish (mcp-server) / Build sdist + wheel (push) Canceled after 0s
Build & Publish (mcp-server) / Publish to PyPI (push) Canceled after 0s
This commit is contained in:
@@ -0,0 +1,465 @@
|
||||
package core
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/logger"
|
||||
)
|
||||
|
||||
// Feishu docx block_type integer enum (subset this connector handles).
|
||||
// https://open.feishu.cn/document/server-docs/docs/docs/docx-v1/docx-structure
|
||||
const (
|
||||
BlockTypePage = 1
|
||||
BlockTypeText = 2
|
||||
BlockTypeHeading1 = 3
|
||||
blockTypeHeading9 = 11
|
||||
BlockTypeBullet = 12
|
||||
BlockTypeOrdered = 13
|
||||
BlockTypeCode = 14
|
||||
BlockTypeQuote = 15
|
||||
BlockTypeTodo = 17
|
||||
BlockTypeBitable = 18
|
||||
BlockTypeCallout = 19
|
||||
BlockTypeDivider = 22
|
||||
BlockTypeFile = 23
|
||||
BlockTypeImage = 27
|
||||
BlockTypeSheet = 30
|
||||
BlockTypeTable = 31
|
||||
BlockTypeTableCell = 32
|
||||
)
|
||||
|
||||
// maxDocumentBlocks caps how many blocks a single document contributes, guarding
|
||||
// against pathological/adversarial documents. Far above any real Feishu doc.
|
||||
const maxDocumentBlocks = 50000
|
||||
|
||||
// TextElement is one inline run inside a text-bearing block.
|
||||
// TextRun is the text_run payload of a TextElement.
|
||||
type TextRun struct {
|
||||
Content string `json:"content"`
|
||||
}
|
||||
|
||||
type TextElement struct {
|
||||
TextRun *TextRun `json:"text_run"`
|
||||
}
|
||||
|
||||
// BlockText is the shared shape of text-bearing blocks (text, headingN, bullet…).
|
||||
type BlockText struct {
|
||||
Elements []TextElement `json:"elements"`
|
||||
}
|
||||
|
||||
// BlockTokenRef is the shared shape of sheet/bitable/image block payloads
|
||||
// (the JSON tag differs per field; only the token matters).
|
||||
type BlockTokenRef struct {
|
||||
Token string `json:"token"`
|
||||
}
|
||||
|
||||
// BlockFileRef is the file block payload: the attachment token plus its name.
|
||||
type BlockFileRef struct {
|
||||
Token string `json:"token"`
|
||||
Name string `json:"name"`
|
||||
}
|
||||
|
||||
// BlockTableProperty carries the table grid shape.
|
||||
type BlockTableProperty struct {
|
||||
ColumnSize int `json:"column_size"`
|
||||
}
|
||||
|
||||
// BlockTable is the table block payload: cell block IDs plus grid property.
|
||||
type BlockTable struct {
|
||||
Cells []string `json:"cells"`
|
||||
Property *BlockTableProperty `json:"property"`
|
||||
}
|
||||
|
||||
// DocxBlock is one node in the flat block array returned by the blocks API.
|
||||
type DocxBlock struct {
|
||||
BlockID string `json:"block_id"`
|
||||
ParentID string `json:"parent_id"`
|
||||
BlockType int `json:"block_type"`
|
||||
Children []string `json:"children"`
|
||||
|
||||
Text *BlockText `json:"text"`
|
||||
Heading1 *BlockText `json:"heading1"`
|
||||
Heading2 *BlockText `json:"heading2"`
|
||||
Heading3 *BlockText `json:"heading3"`
|
||||
Heading4 *BlockText `json:"heading4"`
|
||||
Heading5 *BlockText `json:"heading5"`
|
||||
Heading6 *BlockText `json:"heading6"`
|
||||
Heading7 *BlockText `json:"heading7"`
|
||||
Heading8 *BlockText `json:"heading8"`
|
||||
Heading9 *BlockText `json:"heading9"`
|
||||
Bullet *BlockText `json:"bullet"`
|
||||
Ordered *BlockText `json:"ordered"`
|
||||
Code *BlockText `json:"code"`
|
||||
Quote *BlockText `json:"quote"`
|
||||
Todo *BlockText `json:"todo"`
|
||||
Callout *BlockText `json:"callout"`
|
||||
|
||||
Sheet *BlockTokenRef `json:"sheet"`
|
||||
Bitable *BlockTokenRef `json:"bitable"`
|
||||
File *BlockFileRef `json:"file"`
|
||||
Image *BlockTokenRef `json:"image"`
|
||||
|
||||
Table *BlockTable `json:"table"`
|
||||
}
|
||||
|
||||
// DocxBlocksData is the data payload of DocxBlocksResponse.
|
||||
type DocxBlocksData struct {
|
||||
Items []DocxBlock `json:"items"`
|
||||
HasMore bool `json:"has_more"`
|
||||
PageToken string `json:"page_token"`
|
||||
}
|
||||
|
||||
// DocxBlocksResponse is the response for GET .../documents/:id/blocks.
|
||||
type DocxBlocksResponse struct {
|
||||
ApiResponse
|
||||
Data DocxBlocksData `json:"data"`
|
||||
}
|
||||
|
||||
// listDocumentBlocks returns every block of a docx document as a flat,
|
||||
// pre-order array. Paginates at 500 blocks/page. documentID is the obj_token.
|
||||
func (c *Client) listDocumentBlocks(ctx context.Context, documentID string) ([]DocxBlock, error) {
|
||||
var all []DocxBlock
|
||||
pageToken := ""
|
||||
for {
|
||||
path := fmt.Sprintf("/open-apis/docx/v1/documents/%s/blocks?page_size=500&document_revision_id=-1",
|
||||
url.PathEscape(documentID))
|
||||
if pageToken != "" {
|
||||
path += "&page_token=" + url.QueryEscape(pageToken)
|
||||
}
|
||||
var resp DocxBlocksResponse
|
||||
if err := c.DoRequest(ctx, http.MethodGet, path, nil, &resp); err != nil {
|
||||
return nil, fmt.Errorf("list document blocks: %w", err)
|
||||
}
|
||||
if resp.Code != 0 {
|
||||
return nil, fmt.Errorf("list document blocks error: code=%d msg=%s", resp.Code, resp.Msg)
|
||||
}
|
||||
all = append(all, resp.Data.Items...)
|
||||
if len(resp.Data.Items) == 0 {
|
||||
// Defensive: a malformed page (empty items but has_more=true and a
|
||||
// non-empty page_token) would otherwise loop until the task deadline,
|
||||
// burning API quota. There is nothing more to collect, so stop.
|
||||
break
|
||||
}
|
||||
if len(all) >= maxDocumentBlocks {
|
||||
logger.Warnf(ctx, "[Feishu] document %s exceeded %d blocks; truncating", documentID, maxDocumentBlocks)
|
||||
break
|
||||
}
|
||||
if !resp.Data.HasMore || resp.Data.PageToken == "" {
|
||||
break
|
||||
}
|
||||
pageToken = resp.Data.PageToken
|
||||
}
|
||||
return all, nil
|
||||
}
|
||||
|
||||
// maxTableRows caps how many rows of an embedded sheet/bitable are rendered,
|
||||
// protecting chunking from pathologically large tables. Beyond this the table
|
||||
// is truncated and the caller annotates the omission.
|
||||
const maxTableRows = 500
|
||||
|
||||
// sheetValueRange is the valueRange payload of sheetValuesData.
|
||||
type sheetValueRange struct {
|
||||
Values [][]any `json:"values"`
|
||||
}
|
||||
|
||||
// sheetValuesData is the data payload of sheetValuesResponse.
|
||||
type sheetValuesData struct {
|
||||
ValueRange sheetValueRange `json:"valueRange"`
|
||||
}
|
||||
|
||||
// sheetValuesResponse is the response for sheets-v2 values read.
|
||||
type sheetValuesResponse struct {
|
||||
ApiResponse
|
||||
Data sheetValuesData `json:"data"`
|
||||
}
|
||||
|
||||
// readSheetRange reads the cell values of an embedded spreadsheet block.
|
||||
// embedToken is the block's sheet.token, formatted "spreadsheetToken_sheetId".
|
||||
// Cells are stringified (display value) for RAG text retrieval. Rows are capped
|
||||
// at maxTableRows; truncated is true when the source had more rows than that.
|
||||
func (c *Client) readSheetRange(ctx context.Context, embedToken string) ([][]string, bool, error) {
|
||||
idx := strings.LastIndex(embedToken, "_")
|
||||
if idx < 0 {
|
||||
return nil, false, fmt.Errorf("invalid sheet embed token: %q", embedToken)
|
||||
}
|
||||
spreadsheetToken, sheetID := embedToken[:idx], embedToken[idx+1:]
|
||||
path := fmt.Sprintf("/open-apis/sheets/v2/spreadsheets/%s/values/%s?valueRenderOption=ToString",
|
||||
url.PathEscape(spreadsheetToken), url.PathEscape(sheetID))
|
||||
var resp sheetValuesResponse
|
||||
if err := c.DoRequest(ctx, http.MethodGet, path, nil, &resp); err != nil {
|
||||
return nil, false, fmt.Errorf("read sheet range: %w", err)
|
||||
}
|
||||
if resp.Code != 0 {
|
||||
return nil, false, fmt.Errorf("read sheet range error: code=%d msg=%s", resp.Code, resp.Msg)
|
||||
}
|
||||
raw, truncated := capRows(resp.Data.ValueRange.Values)
|
||||
return stringifyMatrix(raw), truncated, nil
|
||||
}
|
||||
|
||||
// capRows limits rows to maxTableRows, reporting whether truncation happened. It
|
||||
// is generic so the sheet ([][]any) and bitable ([][]string) read paths share one
|
||||
// truncation rule instead of each inlining their own.
|
||||
func capRows[T any](rows []T) ([]T, bool) {
|
||||
if len(rows) > maxTableRows {
|
||||
return rows[:maxTableRows], true
|
||||
}
|
||||
return rows, false
|
||||
}
|
||||
|
||||
// stringifyMatrix renders arbitrary cell values to strings; nil → "".
|
||||
func stringifyMatrix(in [][]any) [][]string {
|
||||
out := make([][]string, len(in))
|
||||
for i, row := range in {
|
||||
cells := make([]string, len(row))
|
||||
for j, v := range row {
|
||||
cells[j] = cellToString(v)
|
||||
}
|
||||
out[i] = cells
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// cellToString renders a single JSON cell value to a string; nil → "".
|
||||
func cellToString(v any) string {
|
||||
switch t := v.(type) {
|
||||
case nil:
|
||||
return ""
|
||||
case string:
|
||||
return t
|
||||
case float64:
|
||||
return strconv.FormatFloat(t, 'f', -1, 64)
|
||||
case bool:
|
||||
return strconv.FormatBool(t)
|
||||
default:
|
||||
return fmt.Sprintf("%v", t)
|
||||
}
|
||||
}
|
||||
|
||||
// bitableFieldTypeDateTime is the Feishu bitable field type for a date/datetime
|
||||
// column. Its cell value is a bare Unix-millisecond number, so without type-aware
|
||||
// formatting it would render as a 13-digit integer instead of a readable date.
|
||||
// https://open.feishu.cn/document/server-docs/docs/bitable-v1/app-table-field/guide
|
||||
const bitableFieldTypeDateTime = 5
|
||||
|
||||
// bitableColumn describes one bitable field: its name, integer type, and (for a
|
||||
// date column) its date_formatter, which distinguishes a date-only column
|
||||
// ("yyyy/MM/dd") from a datetime one ("yyyy-MM-dd HH:mm").
|
||||
type bitableColumn struct {
|
||||
name string
|
||||
fieldType int
|
||||
dateFormatter string
|
||||
}
|
||||
|
||||
// bitableFieldCell renders one bitable cell. Date columns carry a Unix-millisecond
|
||||
// UTC instant; the calendar date a user sees is that instant in the table's
|
||||
// timezone, so it is formatted in loc (rendering in UTC would shift the date, e.g.
|
||||
// a GMT+8 "2024-04-01" would show as "2024-03-31 16:00:00"). A date-only formatter
|
||||
// yields "2006-01-02"; a datetime formatter adds the "15:04" time. Every other
|
||||
// type falls through to bitableCellToString.
|
||||
func bitableFieldCell(v any, col bitableColumn, loc *time.Location) string {
|
||||
if col.fieldType == bitableFieldTypeDateTime {
|
||||
// Empty date cells arrive as nil (→ fall through to blank); ms==0 (epoch) is
|
||||
// not a real Feishu date value either — render blank, not "1970-01-01" or "0".
|
||||
if ms, ok := v.(float64); ok {
|
||||
if ms == 0 {
|
||||
return ""
|
||||
}
|
||||
t := time.UnixMilli(int64(ms)).In(loc)
|
||||
if dateFormatterHasTime(col.dateFormatter) {
|
||||
return t.Format("2006-01-02 15:04")
|
||||
}
|
||||
return t.Format("2006-01-02")
|
||||
}
|
||||
}
|
||||
return bitableCellToString(v)
|
||||
}
|
||||
|
||||
// dateFormatterHasTime reports whether a Feishu date_formatter includes a time
|
||||
// component. Feishu formatters are Java-style: 'H'/'h' hour, lowercase 'm' minute,
|
||||
// uppercase 'M' month — so a time part is present iff the string carries an hour or
|
||||
// a lowercase-'m' minute token (e.g. "yyyy-MM-dd HH:mm"). An empty formatter
|
||||
// (Feishu's default "yyyy/MM/dd") is date-only.
|
||||
func dateFormatterHasTime(f string) bool {
|
||||
return strings.ContainsAny(f, "Hh") || strings.Contains(f, "m")
|
||||
}
|
||||
|
||||
// bitableCellToString renders a bitable field value — which may be text segments,
|
||||
// a person/link object, or an attachment/multi-select array — to a readable
|
||||
// string for RAG. Scalars delegate to cellToString.
|
||||
func bitableCellToString(v any) string {
|
||||
switch t := v.(type) {
|
||||
case []any:
|
||||
parts := make([]string, 0, len(t))
|
||||
for _, e := range t {
|
||||
if s := bitableCellToString(e); s != "" {
|
||||
parts = append(parts, s)
|
||||
}
|
||||
}
|
||||
return strings.Join(parts, ", ")
|
||||
case map[string]any:
|
||||
for _, k := range []string{"text", "name", "en_name", "link"} {
|
||||
if s, ok := t[k].(string); ok && s != "" {
|
||||
return s
|
||||
}
|
||||
}
|
||||
return ""
|
||||
default:
|
||||
return cellToString(v)
|
||||
}
|
||||
}
|
||||
|
||||
// bitableFieldProperty carries a field's type-specific settings.
|
||||
type bitableFieldProperty struct {
|
||||
// DateFormatter distinguishes a date-only column from a datetime one.
|
||||
DateFormatter string `json:"date_formatter"`
|
||||
}
|
||||
|
||||
// bitableField is one entry of bitableFieldsData.Items.
|
||||
type bitableField struct {
|
||||
FieldName string `json:"field_name"`
|
||||
Type int `json:"type"`
|
||||
Property *bitableFieldProperty `json:"property"`
|
||||
}
|
||||
|
||||
// bitableFieldsData is the data payload of bitableFieldsResponse.
|
||||
type bitableFieldsData struct {
|
||||
HasMore bool `json:"has_more"`
|
||||
PageToken string `json:"page_token"`
|
||||
Items []bitableField `json:"items"`
|
||||
}
|
||||
|
||||
type bitableFieldsResponse struct {
|
||||
ApiResponse
|
||||
Data bitableFieldsData `json:"data"`
|
||||
}
|
||||
|
||||
// maxBitableFieldPageSize is the documented per-page cap for the bitable
|
||||
// list-fields endpoint (100). Unlike list-records (max 500), passing a larger
|
||||
// page_size here is rejected, so fields must be fetched 100 at a time and paged.
|
||||
const maxBitableFieldPageSize = 100
|
||||
|
||||
// bitableRecord is one entry of bitableRecordsData.Items.
|
||||
type bitableRecord struct {
|
||||
Fields map[string]any `json:"fields"`
|
||||
}
|
||||
|
||||
// bitableRecordsData is the data payload of bitableRecordsResponse.
|
||||
type bitableRecordsData struct {
|
||||
HasMore bool `json:"has_more"`
|
||||
PageToken string `json:"page_token"`
|
||||
Items []bitableRecord `json:"items"`
|
||||
}
|
||||
|
||||
type bitableRecordsResponse struct {
|
||||
ApiResponse
|
||||
Data bitableRecordsData `json:"data"`
|
||||
}
|
||||
|
||||
// readBitableRecords reads an embedded bitable block as a table: a header row of
|
||||
// field names followed by one row per record. embedToken is the block's
|
||||
// bitable.token, formatted "appToken_tableId". Record rows are capped at
|
||||
// maxTableRows; truncated is true when the source had more records than that.
|
||||
func (c *Client) readBitableRecords(ctx context.Context, embedToken string) ([][]string, bool, error) {
|
||||
idx := strings.LastIndex(embedToken, "_")
|
||||
if idx < 0 {
|
||||
return nil, false, fmt.Errorf("invalid bitable embed token: %q", embedToken)
|
||||
}
|
||||
appToken, tableID := embedToken[:idx], embedToken[idx+1:]
|
||||
|
||||
var cols []bitableColumn
|
||||
baseFPath := fmt.Sprintf("/open-apis/bitable/v1/apps/%s/tables/%s/fields?page_size=%d",
|
||||
url.PathEscape(appToken), url.PathEscape(tableID), maxBitableFieldPageSize)
|
||||
fieldPageToken := ""
|
||||
for {
|
||||
fpath := baseFPath
|
||||
if fieldPageToken != "" {
|
||||
fpath += "&page_token=" + url.QueryEscape(fieldPageToken)
|
||||
}
|
||||
var fieldsResp bitableFieldsResponse
|
||||
if err := c.DoRequest(ctx, http.MethodGet, fpath, nil, &fieldsResp); err != nil {
|
||||
return nil, false, fmt.Errorf("read bitable fields: %w", err)
|
||||
}
|
||||
if fieldsResp.Code != 0 {
|
||||
return nil, false, fmt.Errorf("read bitable fields error: code=%d msg=%s", fieldsResp.Code, fieldsResp.Msg)
|
||||
}
|
||||
for _, f := range fieldsResp.Data.Items {
|
||||
formatter := ""
|
||||
if f.Property != nil {
|
||||
formatter = f.Property.DateFormatter
|
||||
}
|
||||
cols = append(cols, bitableColumn{name: f.FieldName, fieldType: f.Type, dateFormatter: formatter})
|
||||
}
|
||||
if len(fieldsResp.Data.Items) == 0 {
|
||||
// Defensive: an empty page with has_more=true would loop forever (this
|
||||
// loop has no size cap of its own). Nothing more to read, so stop.
|
||||
break
|
||||
}
|
||||
if !fieldsResp.Data.HasMore || fieldsResp.Data.PageToken == "" {
|
||||
break
|
||||
}
|
||||
fieldPageToken = fieldsResp.Data.PageToken
|
||||
}
|
||||
|
||||
header := make([]string, len(cols))
|
||||
for i, col := range cols {
|
||||
header[i] = col.name
|
||||
}
|
||||
loc := c.tz()
|
||||
|
||||
var dataRows [][]string
|
||||
truncated := false
|
||||
// Use the Search-records endpoint (POST .../records/search). The legacy
|
||||
// GET .../records is officially deprecated ("已不推荐使用,可使用[查询记录]替代").
|
||||
// An empty body queries the table's default view, paginated via page_token
|
||||
// (we page to the end below), so a default view with a filter would omit the
|
||||
// filtered-out records — acceptable for RAG, but not literally "all records".
|
||||
baseRPath := fmt.Sprintf("/open-apis/bitable/v1/apps/%s/tables/%s/records/search?page_size=500",
|
||||
url.PathEscape(appToken), url.PathEscape(tableID))
|
||||
pageToken := ""
|
||||
for {
|
||||
rpath := baseRPath
|
||||
if pageToken != "" {
|
||||
rpath += "&page_token=" + url.QueryEscape(pageToken)
|
||||
}
|
||||
var rec bitableRecordsResponse
|
||||
if err := c.DoRequest(ctx, http.MethodPost, rpath, map[string]any{}, &rec); err != nil {
|
||||
return nil, false, fmt.Errorf("search bitable records: %w", err)
|
||||
}
|
||||
if rec.Code != 0 {
|
||||
return nil, false, fmt.Errorf("read bitable records error: code=%d msg=%s", rec.Code, rec.Msg)
|
||||
}
|
||||
for _, item := range rec.Data.Items {
|
||||
row := make([]string, len(cols))
|
||||
for i, col := range cols {
|
||||
row[i] = bitableFieldCell(item.Fields[col.name], col, loc)
|
||||
}
|
||||
dataRows = append(dataRows, row)
|
||||
}
|
||||
if len(rec.Data.Items) == 0 {
|
||||
// Defensive: an empty page with has_more=true never advances dataRows,
|
||||
// so the maxTableRows cap below would never trip — stop instead of
|
||||
// looping until the task deadline.
|
||||
break
|
||||
}
|
||||
if len(dataRows) >= maxTableRows {
|
||||
if len(dataRows) > maxTableRows || rec.Data.HasMore {
|
||||
truncated = true
|
||||
}
|
||||
break
|
||||
}
|
||||
if !rec.Data.HasMore || rec.Data.PageToken == "" {
|
||||
break
|
||||
}
|
||||
pageToken = rec.Data.PageToken
|
||||
}
|
||||
dataRows, capped := capRows(dataRows)
|
||||
truncated = truncated || capped
|
||||
rows := append([][]string{header}, dataRows...)
|
||||
return rows, truncated, nil
|
||||
}
|
||||
@@ -0,0 +1,528 @@
|
||||
package core
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strconv"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestListDocumentBlocks_Paginates(t *testing.T) {
|
||||
var gotTokens []string
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.URL.Path == "/open-apis/auth/v3/tenant_access_token/internal" {
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{"code": 0, "tenant_access_token": "t", "expire": 7200})
|
||||
return
|
||||
}
|
||||
gotTokens = append(gotTokens, r.URL.Query().Get("page_token"))
|
||||
if r.URL.Query().Get("page_size") != "500" {
|
||||
t.Errorf("page_size = %q, want 500", r.URL.Query().Get("page_size"))
|
||||
}
|
||||
if r.URL.Query().Get("page_token") == "" {
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{"code": 0, "data": map[string]any{
|
||||
"items": []map[string]any{{"block_id": "b1", "block_type": 2}, {"block_id": "b2", "block_type": 2}},
|
||||
"has_more": true,
|
||||
"page_token": "p2",
|
||||
}})
|
||||
return
|
||||
}
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{"code": 0, "data": map[string]any{
|
||||
"items": []map[string]any{{"block_id": "b3", "block_type": 2}},
|
||||
"has_more": false,
|
||||
}})
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
c := &Client{baseURL: srv.URL, appID: "a", appSecret: "s", httpClient: srv.Client()}
|
||||
blocks, err := c.listDocumentBlocks(context.Background(), "doc123")
|
||||
if err != nil {
|
||||
t.Fatalf("listDocumentBlocks: %v", err)
|
||||
}
|
||||
if len(blocks) != 3 {
|
||||
t.Fatalf("got %d blocks, want 3", len(blocks))
|
||||
}
|
||||
if blocks[0].BlockID != "b1" || blocks[2].BlockID != "b3" {
|
||||
t.Errorf("order not preserved: %+v", blocks)
|
||||
}
|
||||
if len(gotTokens) != 2 || gotTokens[1] != "p2" {
|
||||
t.Errorf("page tokens = %v, want ['', 'p2']", gotTokens)
|
||||
}
|
||||
}
|
||||
|
||||
func TestReadSheetRange_SplitsTokenAndReadsValues(t *testing.T) {
|
||||
var gotPath string
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.URL.Path == "/open-apis/auth/v3/tenant_access_token/internal" {
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{"code": 0, "tenant_access_token": "t", "expire": 7200})
|
||||
return
|
||||
}
|
||||
gotPath = r.URL.Path
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{"code": 0, "data": map[string]any{
|
||||
"valueRange": map[string]any{"values": [][]any{{"名称", "数量"}, {"苹果", 3}, {"梨", nil}}},
|
||||
}})
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
c := &Client{baseURL: srv.URL, appID: "a", appSecret: "s", httpClient: srv.Client()}
|
||||
rows, truncated, err := c.readSheetRange(context.Background(), "sht_abc_0")
|
||||
if err != nil {
|
||||
t.Fatalf("readSheetRange: %v", err)
|
||||
}
|
||||
if gotPath != "/open-apis/sheets/v2/spreadsheets/sht_abc/values/0" {
|
||||
t.Errorf("path = %q", gotPath)
|
||||
}
|
||||
if truncated {
|
||||
t.Errorf("did not expect truncation for 3 rows")
|
||||
}
|
||||
if len(rows) != 3 || rows[0][0] != "名称" || rows[1][1] != "3" || rows[2][1] != "" {
|
||||
t.Errorf("rows = %+v", rows)
|
||||
}
|
||||
}
|
||||
|
||||
func TestReadSheetRange_TruncatesLargeTable(t *testing.T) {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.URL.Path == "/open-apis/auth/v3/tenant_access_token/internal" {
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{"code": 0, "tenant_access_token": "t", "expire": 7200})
|
||||
return
|
||||
}
|
||||
values := make([][]any, 600)
|
||||
for i := range values {
|
||||
values[i] = []any{i}
|
||||
}
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{"code": 0, "data": map[string]any{
|
||||
"valueRange": map[string]any{"values": values},
|
||||
}})
|
||||
}))
|
||||
defer srv.Close()
|
||||
c := &Client{baseURL: srv.URL, appID: "a", appSecret: "s", httpClient: srv.Client()}
|
||||
rows, truncated, err := c.readSheetRange(context.Background(), "sht_x_0")
|
||||
if err != nil {
|
||||
t.Fatalf("err: %v", err)
|
||||
}
|
||||
if !truncated {
|
||||
t.Errorf("expected truncated=true for 600 rows")
|
||||
}
|
||||
if len(rows) != 500 {
|
||||
t.Errorf("rows capped to %d, want 500", len(rows))
|
||||
}
|
||||
}
|
||||
|
||||
func TestBitableCellToString(t *testing.T) {
|
||||
// text-segment array: [{type:text, text:"你好"}] → "你好"
|
||||
textSeg := []any{map[string]any{"type": "text", "text": "你好"}}
|
||||
if got := bitableCellToString(textSeg); got != "你好" {
|
||||
t.Errorf("text segment: got %q, want %q", got, "你好")
|
||||
}
|
||||
|
||||
// person array: [{id:u1, name:"张三"}] → "张三"
|
||||
person := []any{map[string]any{"id": "u1", "name": "张三"}}
|
||||
if got := bitableCellToString(person); got != "张三" {
|
||||
t.Errorf("person: got %q, want %q", got, "张三")
|
||||
}
|
||||
|
||||
// plain number (float64) → "3"
|
||||
if got := bitableCellToString(float64(3)); got != "3" {
|
||||
t.Errorf("float64: got %q, want %q", got, "3")
|
||||
}
|
||||
|
||||
// plain string → passthrough
|
||||
if got := bitableCellToString("abc"); got != "abc" {
|
||||
t.Errorf("string: got %q, want %q", got, "abc")
|
||||
}
|
||||
}
|
||||
|
||||
func TestBitableFieldCell_DateColumn(t *testing.T) {
|
||||
gmt8 := resolveLocation("") // default GMT+8
|
||||
dateOnly := bitableColumn{fieldType: bitableFieldTypeDateTime, dateFormatter: "yyyy/MM/dd"}
|
||||
dateTime := bitableColumn{fieldType: bitableFieldTypeDateTime, dateFormatter: "yyyy-MM-dd HH:mm"}
|
||||
number := bitableColumn{fieldType: 2}
|
||||
|
||||
// Feishu stores a date as a UTC instant at the table-timezone midnight. The
|
||||
// official sample 1711900800000 is 2024-04-01 00:00:00 GMT+8 (== 2024-03-31
|
||||
// 16:00:00 UTC); a date-only column must render the GMT+8 calendar date, NOT
|
||||
// the UTC-shifted "2024-03-31 16:00:00" the old code produced.
|
||||
if got := bitableFieldCell(float64(1711900800000), dateOnly, gmt8); got != "2024-04-01" {
|
||||
t.Errorf("date-only cell: got %q, want %q", got, "2024-04-01")
|
||||
}
|
||||
|
||||
// A datetime column keeps the time, rendered in the same timezone.
|
||||
// 1719802800000 == 2024-07-01 03:00 UTC == 2024-07-01 11:00 GMT+8.
|
||||
if got := bitableFieldCell(float64(1719802800000), dateTime, gmt8); got != "2024-07-01 11:00" {
|
||||
t.Errorf("datetime cell: got %q, want %q", got, "2024-07-01 11:00")
|
||||
}
|
||||
|
||||
// A non-date numeric field is left as its plain value, not a date.
|
||||
if got := bitableFieldCell(float64(42), number, gmt8); got != "42" {
|
||||
t.Errorf("number cell: got %q, want %q", got, "42")
|
||||
}
|
||||
|
||||
// An empty date cell (nil, or epoch 0) renders blank, not "1970-01-01".
|
||||
if got := bitableFieldCell(nil, dateOnly, gmt8); got != "" {
|
||||
t.Errorf("empty date cell: got %q, want empty", got)
|
||||
}
|
||||
if got := bitableFieldCell(float64(0), dateOnly, gmt8); got != "" {
|
||||
t.Errorf("epoch-0 date cell: got %q, want empty", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDateFormatterHasTime(t *testing.T) {
|
||||
withTime := []string{"yyyy-MM-dd HH:mm", "HH:mm", "yyyy/MM/dd hh:mm", "MM-dd HH:mm"}
|
||||
dateOnly := []string{"", "yyyy/MM/dd", "MM-dd", "MM/dd/yyyy", "dd/MM/yyyy"}
|
||||
for _, f := range withTime {
|
||||
if !dateFormatterHasTime(f) {
|
||||
t.Errorf("dateFormatterHasTime(%q) = false, want true", f)
|
||||
}
|
||||
}
|
||||
for _, f := range dateOnly {
|
||||
if dateFormatterHasTime(f) {
|
||||
t.Errorf("dateFormatterHasTime(%q) = true, want false", f)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveLocation(t *testing.T) {
|
||||
// Empty and unknown names both fall back to a fixed GMT+8, with no dependency
|
||||
// on system tzdata (minimal container images may omit it).
|
||||
for _, name := range []string{"", "Not/AZone"} {
|
||||
ref := time.Date(2024, 4, 1, 0, 0, 0, 0, resolveLocation(name))
|
||||
if _, off := ref.Zone(); off != defaultTimezoneOffsetSeconds {
|
||||
t.Errorf("resolveLocation(%q) offset = %d, want %d", name, off, defaultTimezoneOffsetSeconds)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestReadBitableRecords_DateColumnUsesFormatterAndTimezone proves the date fix
|
||||
// end-to-end: property.date_formatter is parsed from the fields response, a type-5
|
||||
// column is detected, and the ms instant is rendered in the client timezone
|
||||
// (GMT+8) — date-only vs datetime distinguished by the formatter.
|
||||
func TestReadBitableRecords_DateColumnUsesFormatterAndTimezone(t *testing.T) {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
switch {
|
||||
case r.URL.Path == "/open-apis/auth/v3/tenant_access_token/internal":
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{"code": 0, "tenant_access_token": "t", "expire": 7200})
|
||||
case strings.HasSuffix(r.URL.Path, "/fields"):
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{"code": 0, "data": map[string]any{
|
||||
"items": []map[string]any{
|
||||
{"field_name": "截止日期", "type": 5, "property": map[string]any{"date_formatter": "yyyy/MM/dd"}},
|
||||
{"field_name": "提醒时间", "type": 5, "property": map[string]any{"date_formatter": "yyyy-MM-dd HH:mm"}},
|
||||
},
|
||||
}})
|
||||
case strings.HasSuffix(r.URL.Path, "/records/search"):
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{"code": 0, "data": map[string]any{
|
||||
"has_more": false, "items": []map[string]any{
|
||||
{"fields": map[string]any{"截止日期": float64(1711900800000), "提醒时间": float64(1719802800000)}},
|
||||
},
|
||||
}})
|
||||
}
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
c := &Client{baseURL: srv.URL, appID: "a", appSecret: "s", httpClient: srv.Client(), location: resolveLocation("")}
|
||||
rows, _, err := c.readBitableRecords(context.Background(), "bascabc_tblxyz")
|
||||
if err != nil {
|
||||
t.Fatalf("readBitableRecords: %v", err)
|
||||
}
|
||||
if len(rows) != 2 {
|
||||
t.Fatalf("rows = %+v, want header + 1 record", rows)
|
||||
}
|
||||
if rows[1][0] != "2024-04-01" {
|
||||
t.Errorf("date-only column = %q, want 2024-04-01 (GMT+8, not UTC-shifted)", rows[1][0])
|
||||
}
|
||||
if rows[1][1] != "2024-07-01 11:00" {
|
||||
t.Errorf("datetime column = %q, want 2024-07-01 11:00 (GMT+8)", rows[1][1])
|
||||
}
|
||||
}
|
||||
|
||||
func TestReadBitableRecords_SplitsTokenAndBuildsTable(t *testing.T) {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
switch {
|
||||
case r.URL.Path == "/open-apis/auth/v3/tenant_access_token/internal":
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{"code": 0, "tenant_access_token": "t", "expire": 7200})
|
||||
case strings.HasSuffix(r.URL.Path, "/fields"):
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{"code": 0, "data": map[string]any{
|
||||
"items": []map[string]any{{"field_name": "任务"}, {"field_name": "状态"}},
|
||||
}})
|
||||
case strings.HasSuffix(r.URL.Path, "/records/search"):
|
||||
// Must use the current Search-records endpoint (POST), not the
|
||||
// deprecated GET .../records legacy interface.
|
||||
if r.Method != http.MethodPost {
|
||||
t.Errorf("records read method = %s, want POST (search endpoint)", r.Method)
|
||||
}
|
||||
if !strings.Contains(r.URL.Path, "/apps/bascabc/tables/tblxyz/") {
|
||||
t.Errorf("bad path %q", r.URL.Path)
|
||||
}
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{"code": 0, "data": map[string]any{
|
||||
"has_more": false,
|
||||
"items": []map[string]any{
|
||||
{"fields": map[string]any{"任务": "写文档", "状态": "进行中"}},
|
||||
},
|
||||
}})
|
||||
}
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
c := &Client{baseURL: srv.URL, appID: "a", appSecret: "s", httpClient: srv.Client()}
|
||||
rows, truncated, err := c.readBitableRecords(context.Background(), "bascabc_tblxyz")
|
||||
if err != nil {
|
||||
t.Fatalf("readBitableRecords: %v", err)
|
||||
}
|
||||
if truncated {
|
||||
t.Errorf("did not expect truncation")
|
||||
}
|
||||
if len(rows) != 2 || rows[0][0] != "任务" || rows[0][1] != "状态" || rows[1][0] != "写文档" || rows[1][1] != "进行中" {
|
||||
t.Errorf("rows = %+v", rows)
|
||||
}
|
||||
}
|
||||
|
||||
// TestReadBitableRecords_PaginatesFieldsAtMax100 guards the real-API constraint
|
||||
// that the list-fields endpoint caps page_size at 100 (unlike list-records, which
|
||||
// allows 500). The connector must request page_size=100 and page through fields,
|
||||
// or a table with >100 columns would silently lose its header tail (or the whole
|
||||
// bitable would degrade if Feishu rejects page_size=500). Verified against the
|
||||
// official docs: bitable-v1/app-table-field/list "最大值:100".
|
||||
func TestReadBitableRecords_PaginatesFieldsAtMax100(t *testing.T) {
|
||||
const totalFields = 150 // spans two pages: 100 + 50
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
switch {
|
||||
case r.URL.Path == "/open-apis/auth/v3/tenant_access_token/internal":
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{"code": 0, "tenant_access_token": "t", "expire": 7200})
|
||||
case strings.HasSuffix(r.URL.Path, "/fields"):
|
||||
if ps := r.URL.Query().Get("page_size"); ps != "100" {
|
||||
t.Errorf("fields page_size = %q, want 100 (endpoint max)", ps)
|
||||
}
|
||||
// Page 1 (no token): first 100 fields, has_more. Page 2 (token=f2): last 50.
|
||||
var items []map[string]any
|
||||
hasMore := false
|
||||
nextTok := ""
|
||||
if r.URL.Query().Get("page_token") == "" {
|
||||
for i := 0; i < 100; i++ {
|
||||
items = append(items, map[string]any{"field_name": fieldName(i)})
|
||||
}
|
||||
hasMore, nextTok = true, "f2"
|
||||
} else {
|
||||
for i := 100; i < totalFields; i++ {
|
||||
items = append(items, map[string]any{"field_name": fieldName(i)})
|
||||
}
|
||||
}
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{"code": 0, "data": map[string]any{
|
||||
"items": items, "has_more": hasMore, "page_token": nextTok,
|
||||
}})
|
||||
case strings.HasSuffix(r.URL.Path, "/records/search"):
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{"code": 0, "data": map[string]any{
|
||||
"has_more": false, "items": []map[string]any{},
|
||||
}})
|
||||
}
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
c := &Client{baseURL: srv.URL, appID: "a", appSecret: "s", httpClient: srv.Client()}
|
||||
rows, _, err := c.readBitableRecords(context.Background(), "bascabc_tblxyz")
|
||||
if err != nil {
|
||||
t.Fatalf("readBitableRecords: %v", err)
|
||||
}
|
||||
if len(rows) < 1 {
|
||||
t.Fatal("expected at least a header row")
|
||||
}
|
||||
header := rows[0]
|
||||
if len(header) != totalFields {
|
||||
t.Fatalf("header has %d fields, want %d (fields pagination lost columns)", len(header), totalFields)
|
||||
}
|
||||
if header[0] != fieldName(0) || header[totalFields-1] != fieldName(totalFields-1) {
|
||||
t.Errorf("header endpoints wrong: first=%q last=%q", header[0], header[totalFields-1])
|
||||
}
|
||||
}
|
||||
|
||||
func fieldName(i int) string { return "f" + strconv.Itoa(i) }
|
||||
|
||||
// bitableRows builds n bitable record items each carrying a single "col" field.
|
||||
func bitableRows(n int) []map[string]any {
|
||||
items := make([]map[string]any, n)
|
||||
for i := range items {
|
||||
items[i] = map[string]any{"fields": map[string]any{"col": "v" + strconv.Itoa(i)}}
|
||||
}
|
||||
return items
|
||||
}
|
||||
|
||||
// TestReadBitableRecords_PaginatesAndTruncatesRecords exercises the record-side
|
||||
// pagination advance and the >maxTableRows truncation — the boundary logic that
|
||||
// no prior test reached (every other bitable test returns a single sub-cap page).
|
||||
func TestReadBitableRecords_PaginatesAndTruncatesRecords(t *testing.T) {
|
||||
var secondPageToken string
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
switch {
|
||||
case r.URL.Path == "/open-apis/auth/v3/tenant_access_token/internal":
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{"code": 0, "tenant_access_token": "t", "expire": 7200})
|
||||
case strings.HasSuffix(r.URL.Path, "/fields"):
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{"code": 0, "data": map[string]any{
|
||||
"items": []map[string]any{{"field_name": "col"}},
|
||||
}})
|
||||
case strings.HasSuffix(r.URL.Path, "/records/search"):
|
||||
if tok := r.URL.Query().Get("page_token"); tok == "" {
|
||||
// Page 1: 300 records, more to come.
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{"code": 0, "data": map[string]any{
|
||||
"has_more": true, "page_token": "r2", "items": bitableRows(300),
|
||||
}})
|
||||
} else {
|
||||
secondPageToken = tok
|
||||
// Page 2: 250 more → 550 total, capped to maxTableRows.
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{"code": 0, "data": map[string]any{
|
||||
"has_more": false, "items": bitableRows(250),
|
||||
}})
|
||||
}
|
||||
}
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
c := &Client{baseURL: srv.URL, appID: "a", appSecret: "s", httpClient: srv.Client()}
|
||||
rows, truncated, err := c.readBitableRecords(context.Background(), "bascabc_tblxyz")
|
||||
if err != nil {
|
||||
t.Fatalf("readBitableRecords: %v", err)
|
||||
}
|
||||
if secondPageToken != "r2" {
|
||||
t.Errorf("second records page must be fetched with page_token=r2, got %q", secondPageToken)
|
||||
}
|
||||
if len(rows) != maxTableRows+1 { // header + maxTableRows capped data rows
|
||||
t.Errorf("rows = %d, want %d (header + %d capped)", len(rows), maxTableRows+1, maxTableRows)
|
||||
}
|
||||
if !truncated {
|
||||
t.Error("a >maxTableRows bitable must report truncated=true")
|
||||
}
|
||||
}
|
||||
|
||||
// TestReadBitableRecords_Exactly500NotTruncated guards the off-by-one at the cap:
|
||||
// exactly maxTableRows records in a single page must NOT be flagged truncated.
|
||||
func TestReadBitableRecords_Exactly500NotTruncated(t *testing.T) {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
switch {
|
||||
case r.URL.Path == "/open-apis/auth/v3/tenant_access_token/internal":
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{"code": 0, "tenant_access_token": "t", "expire": 7200})
|
||||
case strings.HasSuffix(r.URL.Path, "/fields"):
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{"code": 0, "data": map[string]any{
|
||||
"items": []map[string]any{{"field_name": "col"}},
|
||||
}})
|
||||
case strings.HasSuffix(r.URL.Path, "/records/search"):
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{"code": 0, "data": map[string]any{
|
||||
"has_more": false, "items": bitableRows(maxTableRows),
|
||||
}})
|
||||
}
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
c := &Client{baseURL: srv.URL, appID: "a", appSecret: "s", httpClient: srv.Client()}
|
||||
rows, truncated, err := c.readBitableRecords(context.Background(), "bascabc_tblxyz")
|
||||
if err != nil {
|
||||
t.Fatalf("readBitableRecords: %v", err)
|
||||
}
|
||||
if truncated {
|
||||
t.Error("exactly maxTableRows records must not be flagged truncated")
|
||||
}
|
||||
if len(rows) != maxTableRows+1 {
|
||||
t.Errorf("rows = %d, want %d (header + %d)", len(rows), maxTableRows+1, maxTableRows)
|
||||
}
|
||||
}
|
||||
|
||||
// TestReadBitableRecords_EmptyRecordPageTerminates guards the pagination
|
||||
// zero-progress break: a malformed page (empty items but has_more=true and a
|
||||
// non-empty page_token) must terminate instead of looping until the task
|
||||
// deadline. A short context deadline is a safety net so a regression fails fast
|
||||
// rather than hanging the suite.
|
||||
func TestReadBitableRecords_EmptyRecordPageTerminates(t *testing.T) {
|
||||
recordCalls := 0
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
switch {
|
||||
case r.URL.Path == "/open-apis/auth/v3/tenant_access_token/internal":
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{"code": 0, "tenant_access_token": "t", "expire": 7200})
|
||||
case strings.HasSuffix(r.URL.Path, "/fields"):
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{"code": 0, "data": map[string]any{
|
||||
"items": []map[string]any{{"field_name": "col"}},
|
||||
}})
|
||||
case strings.HasSuffix(r.URL.Path, "/records/search"):
|
||||
recordCalls++
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{"code": 0, "data": map[string]any{
|
||||
"has_more": true, "page_token": "loop", "items": []map[string]any{},
|
||||
}})
|
||||
}
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
c := &Client{baseURL: srv.URL, appID: "a", appSecret: "s", httpClient: srv.Client()}
|
||||
rows, _, err := c.readBitableRecords(ctx, "bascabc_tblxyz")
|
||||
if err != nil {
|
||||
t.Fatalf("empty-page pagination must terminate cleanly, got err: %v", err)
|
||||
}
|
||||
if recordCalls != 1 {
|
||||
t.Errorf("records endpoint called %d times, want 1 (loop must break on the empty page)", recordCalls)
|
||||
}
|
||||
if len(rows) != 1 { // header only
|
||||
t.Errorf("rows = %d, want 1 (header only, no records)", len(rows))
|
||||
}
|
||||
}
|
||||
|
||||
// TestReadBitableRecords_EmptyFieldsPageTerminates guards the fields loop, which
|
||||
// has no size cap of its own — an empty page with has_more=true must still break.
|
||||
func TestReadBitableRecords_EmptyFieldsPageTerminates(t *testing.T) {
|
||||
fieldCalls := 0
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
switch {
|
||||
case r.URL.Path == "/open-apis/auth/v3/tenant_access_token/internal":
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{"code": 0, "tenant_access_token": "t", "expire": 7200})
|
||||
case strings.HasSuffix(r.URL.Path, "/fields"):
|
||||
fieldCalls++
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{"code": 0, "data": map[string]any{
|
||||
"has_more": true, "page_token": "loop", "items": []map[string]any{},
|
||||
}})
|
||||
case strings.HasSuffix(r.URL.Path, "/records/search"):
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{"code": 0, "data": map[string]any{
|
||||
"has_more": false, "items": []map[string]any{},
|
||||
}})
|
||||
}
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
c := &Client{baseURL: srv.URL, appID: "a", appSecret: "s", httpClient: srv.Client()}
|
||||
if _, _, err := c.readBitableRecords(ctx, "bascabc_tblxyz"); err != nil {
|
||||
t.Fatalf("empty fields page must terminate cleanly, got err: %v", err)
|
||||
}
|
||||
if fieldCalls != 1 {
|
||||
t.Errorf("fields endpoint called %d times, want 1 (loop must break on the empty page)", fieldCalls)
|
||||
}
|
||||
}
|
||||
|
||||
// TestListDocumentBlocks_EmptyPageTerminates guards the blocks pagination loop.
|
||||
func TestListDocumentBlocks_EmptyPageTerminates(t *testing.T) {
|
||||
blockCalls := 0
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
switch {
|
||||
case r.URL.Path == "/open-apis/auth/v3/tenant_access_token/internal":
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{"code": 0, "tenant_access_token": "t", "expire": 7200})
|
||||
case strings.HasSuffix(r.URL.Path, "/blocks"):
|
||||
blockCalls++
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{"code": 0, "data": map[string]any{
|
||||
"has_more": true, "page_token": "loop", "items": []map[string]any{},
|
||||
}})
|
||||
}
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
c := &Client{baseURL: srv.URL, appID: "a", appSecret: "s", httpClient: srv.Client()}
|
||||
blocks, err := c.listDocumentBlocks(ctx, "doc-1")
|
||||
if err != nil {
|
||||
t.Fatalf("empty blocks page must terminate cleanly, got err: %v", err)
|
||||
}
|
||||
if blockCalls != 1 {
|
||||
t.Errorf("blocks endpoint called %d times, want 1 (loop must break on the empty page)", blockCalls)
|
||||
}
|
||||
if len(blocks) != 0 {
|
||||
t.Errorf("blocks = %d, want 0", len(blocks))
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,916 @@
|
||||
package core
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/datasource"
|
||||
"github.com/Tencent/WeKnora/internal/logger"
|
||||
)
|
||||
|
||||
// Client wraps the Feishu Open Platform API for document/wiki operations.
|
||||
type Client struct {
|
||||
baseURL string
|
||||
appID string
|
||||
appSecret string
|
||||
|
||||
// location renders bitable date cells in the table's timezone (default GMT+8).
|
||||
location *time.Location
|
||||
|
||||
httpClient *http.Client
|
||||
|
||||
// Token cache (thread-safe)
|
||||
tokenMu sync.Mutex
|
||||
tokenCache string
|
||||
tokenExpAt time.Time
|
||||
}
|
||||
|
||||
type WikiNodeListFailure struct {
|
||||
Node WikiNode
|
||||
Err error
|
||||
}
|
||||
|
||||
type PartialWikiNodeListError struct {
|
||||
Failures []WikiNodeListFailure
|
||||
}
|
||||
|
||||
func (e *PartialWikiNodeListError) Error() string {
|
||||
if e == nil || len(e.Failures) == 0 {
|
||||
return "partial wiki node listing failed"
|
||||
}
|
||||
parts := make([]string, 0, len(e.Failures))
|
||||
for _, failure := range e.Failures {
|
||||
parts = append(parts, failure.Err.Error())
|
||||
}
|
||||
return strings.Join(parts, "; ")
|
||||
}
|
||||
|
||||
// tz returns the client's date-rendering location, defaulting to GMT+8 when a
|
||||
// Client was constructed without one (e.g. in tests that build Client directly).
|
||||
func (c *Client) tz() *time.Location {
|
||||
if c.location != nil {
|
||||
return c.location
|
||||
}
|
||||
return time.FixedZone("GMT+8", defaultTimezoneOffsetSeconds)
|
||||
}
|
||||
|
||||
// NewClient creates a new Feishu API client.
|
||||
func NewClient(config *Config) *Client {
|
||||
return &Client{
|
||||
baseURL: config.GetBaseURL(),
|
||||
appID: config.AppID,
|
||||
appSecret: config.AppSecret,
|
||||
location: resolveLocation(config.Timezone),
|
||||
httpClient: datasource.NewConnectorHTTPClient(30 * time.Second),
|
||||
}
|
||||
}
|
||||
|
||||
// GetTenantAccessToken retrieves (or returns cached) tenant access token.
|
||||
// Feishu tokens expire in 2 hours; we cache with a 5-minute safety margin.
|
||||
func (c *Client) GetTenantAccessToken(ctx context.Context) (string, error) {
|
||||
c.tokenMu.Lock()
|
||||
defer c.tokenMu.Unlock()
|
||||
|
||||
if c.tokenCache != "" && time.Now().Before(c.tokenExpAt) {
|
||||
return c.tokenCache, nil
|
||||
}
|
||||
|
||||
payload, _ := json.Marshal(map[string]string{
|
||||
"app_id": c.appID,
|
||||
"app_secret": c.appSecret,
|
||||
})
|
||||
|
||||
url := c.baseURL + "/open-apis/auth/v3/tenant_access_token/internal"
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(payload))
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("create token request: %w", err)
|
||||
}
|
||||
req.Header.Set("Content-Type", "application/json; charset=utf-8")
|
||||
|
||||
resp, err := c.httpClient.Do(req)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("request token: %w", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
var result TokenResponse
|
||||
if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
|
||||
return "", fmt.Errorf("decode token response: %w", err)
|
||||
}
|
||||
if result.Code != 0 {
|
||||
return "", fmt.Errorf("feishu auth error: code=%d msg=%s", result.Code, result.Msg)
|
||||
}
|
||||
|
||||
c.tokenCache = result.TenantAccessToken
|
||||
ttl := time.Duration(result.Expire) * time.Second
|
||||
if ttl > 5*time.Minute {
|
||||
ttl -= 5 * time.Minute
|
||||
}
|
||||
c.tokenExpAt = time.Now().Add(ttl)
|
||||
|
||||
prefixLen := 8
|
||||
if len(result.TenantAccessToken) < prefixLen {
|
||||
prefixLen = len(result.TenantAccessToken)
|
||||
}
|
||||
suffixLen := 4
|
||||
if len(result.TenantAccessToken) < suffixLen {
|
||||
suffixLen = len(result.TenantAccessToken)
|
||||
}
|
||||
logger.Infof(ctx, "[Feishu] got tenant_access_token: %s...%s expire=%ds",
|
||||
result.TenantAccessToken[:prefixLen], result.TenantAccessToken[len(result.TenantAccessToken)-suffixLen:], result.Expire)
|
||||
|
||||
return c.tokenCache, nil
|
||||
}
|
||||
|
||||
// Retry policy shared by DoRequest (JSON API calls) and downloadRawBytes (file
|
||||
// downloads): 429 honours Retry-After, 5xx retries once, transport errors back off.
|
||||
const (
|
||||
feishuMaxRetries = 3
|
||||
feishuMax5xxRetries = 1
|
||||
feishuRetry5xxDelay = 2 * time.Second
|
||||
)
|
||||
|
||||
// maxFeishuDownloadBytes bounds a single file download to protect the sync
|
||||
// worker from adversarial or pathological oversized responses.
|
||||
const maxFeishuDownloadBytes = 512 * 1024 * 1024 // 512 MB
|
||||
|
||||
var feishuRetryBackoff = []time.Duration{2 * time.Second, 4 * time.Second, 8 * time.Second}
|
||||
|
||||
// DoRequest executes an authenticated API request and decodes the JSON response,
|
||||
// retrying transient failures (transport errors, HTTP 429, 5xx). Feishu's drive
|
||||
// export/wiki APIs are aggressively rate limited, and a thousand-document sync
|
||||
// issues tens of thousands of calls; without backoff a single 429 burst used to
|
||||
// fail whole swathes of documents silently. 429 responses honour Retry-After;
|
||||
// 5xx is retried once; other non-2xx statuses fail fast (no point retrying 4xx).
|
||||
func (c *Client) DoRequest(ctx context.Context, method, path string, body interface{}, result interface{}) error {
|
||||
const (
|
||||
maxRetries = feishuMaxRetries
|
||||
max5xxRetries = feishuMax5xxRetries
|
||||
retry5xxDelay = feishuRetry5xxDelay
|
||||
)
|
||||
backoff := feishuRetryBackoff
|
||||
|
||||
token, err := c.GetTenantAccessToken(ctx)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
var bodyBytes []byte
|
||||
if body != nil {
|
||||
bodyBytes, err = json.Marshal(body)
|
||||
if err != nil {
|
||||
return fmt.Errorf("marshal request body: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
url := c.baseURL + path
|
||||
var lastErr error
|
||||
|
||||
for attempt := 0; attempt <= maxRetries; attempt++ {
|
||||
var bodyReader io.Reader
|
||||
if bodyBytes != nil {
|
||||
bodyReader = bytes.NewReader(bodyBytes)
|
||||
}
|
||||
req, err := http.NewRequestWithContext(ctx, method, url, bodyReader)
|
||||
if err != nil {
|
||||
return fmt.Errorf("create request: %w", err)
|
||||
}
|
||||
req.Header.Set("Content-Type", "application/json; charset=utf-8")
|
||||
req.Header.Set("Authorization", "Bearer "+token)
|
||||
|
||||
if attempt == 0 {
|
||||
logger.Infof(ctx, "[Feishu] %s %s", method, path)
|
||||
} else {
|
||||
logger.Infof(ctx, "[Feishu] %s %s (retry %d/%d)", method, path, attempt, maxRetries)
|
||||
}
|
||||
|
||||
resp, err := c.httpClient.Do(req)
|
||||
if err != nil {
|
||||
lastErr = fmt.Errorf("execute request: %w", err)
|
||||
if attempt < maxRetries {
|
||||
if sErr := sleepCtx(ctx, backoff[attempt]); sErr != nil {
|
||||
return sErr
|
||||
}
|
||||
continue
|
||||
}
|
||||
return lastErr
|
||||
}
|
||||
|
||||
respBody, readErr := io.ReadAll(resp.Body)
|
||||
resp.Body.Close()
|
||||
if readErr != nil {
|
||||
lastErr = fmt.Errorf("read response body: %w", readErr)
|
||||
if attempt < maxRetries {
|
||||
if sErr := sleepCtx(ctx, backoff[attempt]); sErr != nil {
|
||||
return sErr
|
||||
}
|
||||
continue
|
||||
}
|
||||
return lastErr
|
||||
}
|
||||
|
||||
logger.Infof(ctx, "[Feishu] %s %s → status=%d bodyLen=%d body=%s",
|
||||
method, path, resp.StatusCode, len(respBody), truncate(string(respBody), 1000))
|
||||
|
||||
if resp.StatusCode == http.StatusTooManyRequests {
|
||||
wait := parseRetryAfter(resp.Header.Get("Retry-After"), backoff[min(attempt, len(backoff)-1)])
|
||||
lastErr = fmt.Errorf("feishu rate limited: status=429 body=%s", truncate(string(respBody), 500))
|
||||
if attempt < maxRetries {
|
||||
if sErr := sleepCtx(ctx, wait); sErr != nil {
|
||||
return sErr
|
||||
}
|
||||
continue
|
||||
}
|
||||
return lastErr
|
||||
}
|
||||
|
||||
if resp.StatusCode >= 500 && resp.StatusCode < 600 {
|
||||
lastErr = fmt.Errorf("feishu server error: status=%d body=%s", resp.StatusCode, truncate(string(respBody), 500))
|
||||
if attempt < max5xxRetries {
|
||||
if sErr := sleepCtx(ctx, retry5xxDelay); sErr != nil {
|
||||
return sErr
|
||||
}
|
||||
continue
|
||||
}
|
||||
return lastErr
|
||||
}
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
return fmt.Errorf("feishu api error: status=%d body=%s", resp.StatusCode, string(respBody))
|
||||
}
|
||||
|
||||
if result != nil {
|
||||
if err := json.Unmarshal(respBody, result); err != nil {
|
||||
return fmt.Errorf("decode response: %w", err)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
return lastErr
|
||||
}
|
||||
|
||||
// parseRetryAfter interprets a Retry-After header value (seconds) into a wait
|
||||
// duration, coercing 0/negative to a short delay and falling back when absent
|
||||
// or unparseable.
|
||||
func parseRetryAfter(header string, fallback time.Duration) time.Duration {
|
||||
if header == "" {
|
||||
return fallback
|
||||
}
|
||||
secs, err := strconv.ParseFloat(strings.TrimSpace(header), 64)
|
||||
if err != nil {
|
||||
return fallback
|
||||
}
|
||||
if secs <= 0 {
|
||||
return 100 * time.Millisecond
|
||||
}
|
||||
return time.Duration(secs * float64(time.Second))
|
||||
}
|
||||
|
||||
// sleepCtx waits for d or until ctx is cancelled, returning ctx.Err() if the
|
||||
// context ends first so retries abort promptly on task cancellation/timeout.
|
||||
func sleepCtx(ctx context.Context, d time.Duration) error {
|
||||
if d <= 0 {
|
||||
return ctx.Err()
|
||||
}
|
||||
timer := time.NewTimer(d)
|
||||
defer timer.Stop()
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return ctx.Err()
|
||||
case <-timer.C:
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
// truncate truncates a string to maxLen and appends "..." if truncated.
|
||||
func truncate(s string, maxLen int) string {
|
||||
if len(s) <= maxLen {
|
||||
return s
|
||||
}
|
||||
return s[:maxLen] + "..."
|
||||
}
|
||||
|
||||
// ListWikiSpaces returns all wiki spaces accessible to the app.
|
||||
func (c *Client) ListWikiSpaces(ctx context.Context) ([]WikiSpace, error) {
|
||||
var allSpaces []WikiSpace
|
||||
pageToken := ""
|
||||
|
||||
for {
|
||||
path := "/open-apis/wiki/v2/spaces?page_size=50"
|
||||
if pageToken != "" {
|
||||
path += "&page_token=" + pageToken
|
||||
}
|
||||
|
||||
var resp WikiSpaceListResponse
|
||||
if err := c.DoRequest(ctx, http.MethodGet, path, nil, &resp); err != nil {
|
||||
return nil, fmt.Errorf("list wiki spaces: %w", err)
|
||||
}
|
||||
if resp.Code != 0 {
|
||||
logger.Errorf(ctx, "[Feishu] ListWikiSpaces error: code=%d msg=%s", resp.Code, resp.Msg)
|
||||
return nil, fmt.Errorf("list wiki spaces error: code=%d msg=%s", resp.Code, resp.Msg)
|
||||
}
|
||||
|
||||
logger.Infof(ctx, "[Feishu] ListWikiSpaces: got %d spaces, has_more=%v", len(resp.Data.Items), resp.Data.HasMore)
|
||||
for i, s := range resp.Data.Items {
|
||||
logger.Infof(ctx, "[Feishu] space[%d]: id=%s name=%q visibility=%s", i, s.SpaceID, s.Name, s.Visibility)
|
||||
}
|
||||
|
||||
allSpaces = append(allSpaces, resp.Data.Items...)
|
||||
|
||||
if !resp.Data.HasMore || resp.Data.PageToken == "" {
|
||||
break
|
||||
}
|
||||
pageToken = resp.Data.PageToken
|
||||
}
|
||||
|
||||
logger.Infof(ctx, "[Feishu] ListWikiSpaces: total %d spaces", len(allSpaces))
|
||||
return allSpaces, nil
|
||||
}
|
||||
|
||||
// ListWikiNodes returns all nodes (documents) under a wiki space.
|
||||
// If parentNodeToken is empty, returns top-level nodes.
|
||||
func (c *Client) ListWikiNodes(ctx context.Context, spaceID string, parentNodeToken string) ([]WikiNode, error) {
|
||||
var allNodes []WikiNode
|
||||
pageToken := ""
|
||||
|
||||
for {
|
||||
path := fmt.Sprintf("/open-apis/wiki/v2/spaces/%s/nodes?page_size=50", spaceID)
|
||||
if parentNodeToken != "" {
|
||||
path += "&parent_node_token=" + parentNodeToken
|
||||
}
|
||||
if pageToken != "" {
|
||||
path += "&page_token=" + pageToken
|
||||
}
|
||||
|
||||
var resp WikiNodeListResponse
|
||||
if err := c.DoRequest(ctx, http.MethodGet, path, nil, &resp); err != nil {
|
||||
return nil, fmt.Errorf("list wiki nodes: %w", err)
|
||||
}
|
||||
if resp.Code != 0 {
|
||||
return nil, fmt.Errorf("list wiki nodes error: code=%d msg=%s", resp.Code, resp.Msg)
|
||||
}
|
||||
|
||||
for _, node := range resp.Data.Items {
|
||||
if parentNodeToken != "" && node.ParentNodeID == "" {
|
||||
node.ParentNodeID = parentNodeToken
|
||||
}
|
||||
if node.SpaceID == "" {
|
||||
node.SpaceID = spaceID
|
||||
}
|
||||
allNodes = append(allNodes, node)
|
||||
}
|
||||
|
||||
if !resp.Data.HasMore || resp.Data.PageToken == "" {
|
||||
break
|
||||
}
|
||||
pageToken = resp.Data.PageToken
|
||||
}
|
||||
|
||||
return allNodes, nil
|
||||
}
|
||||
|
||||
// GetWikiNode returns metadata for a single wiki node.
|
||||
func (c *Client) GetWikiNode(ctx context.Context, spaceID string, nodeToken string) (WikiNode, error) {
|
||||
path := fmt.Sprintf("/open-apis/wiki/v2/spaces/get_node?token=%s", url.QueryEscape(nodeToken))
|
||||
|
||||
var resp WikiNodeInfoResponse
|
||||
if err := c.DoRequest(ctx, http.MethodGet, path, nil, &resp); err != nil {
|
||||
return WikiNode{}, fmt.Errorf("get wiki node: %w", err)
|
||||
}
|
||||
if resp.Code != 0 {
|
||||
return WikiNode{}, fmt.Errorf("get wiki node error: code=%d msg=%s", resp.Code, resp.Msg)
|
||||
}
|
||||
|
||||
node := resp.Data.Node
|
||||
if node.SpaceID == "" {
|
||||
node.SpaceID = spaceID
|
||||
}
|
||||
return node, nil
|
||||
}
|
||||
|
||||
// listAllWikiNodesRecursive recursively lists all nodes under a wiki space.
|
||||
// It walks the tree depth-first to discover all nested documents.
|
||||
func (c *Client) listAllWikiNodesRecursive(ctx context.Context, spaceID string) ([]WikiNode, error) {
|
||||
// Start with top-level nodes
|
||||
topNodes, err := c.ListWikiNodes(ctx, spaceID, "")
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
var allNodes []WikiNode
|
||||
var failures []WikiNodeListFailure
|
||||
var walk func(nodes []WikiNode)
|
||||
|
||||
walk = func(nodes []WikiNode) {
|
||||
for _, node := range nodes {
|
||||
allNodes = append(allNodes, node)
|
||||
|
||||
// Recurse into child nodes if this node has children
|
||||
if node.HasChild {
|
||||
children, err := c.ListWikiNodes(ctx, spaceID, node.NodeToken)
|
||||
if err != nil {
|
||||
wrappedErr := fmt.Errorf("list children of %s: %w", node.NodeToken, err)
|
||||
failures = append(failures, WikiNodeListFailure{
|
||||
Node: node,
|
||||
Err: wrappedErr,
|
||||
})
|
||||
logger.Warnf(ctx, "[Feishu] partial wiki node listing failure: space=%s node=%s err=%v",
|
||||
spaceID, node.NodeToken, err)
|
||||
continue
|
||||
}
|
||||
walk(children)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
walk(topNodes)
|
||||
if len(failures) > 0 {
|
||||
return allNodes, &PartialWikiNodeListError{Failures: failures}
|
||||
}
|
||||
|
||||
return allNodes, nil
|
||||
}
|
||||
|
||||
// ListWikiNodesRecursiveFrom returns a wiki node and all descendants below it.
|
||||
func (c *Client) ListWikiNodesRecursiveFrom(ctx context.Context, spaceID string, nodeToken string) ([]WikiNode, error) {
|
||||
if nodeToken == "" {
|
||||
return c.listAllWikiNodesRecursive(ctx, spaceID)
|
||||
}
|
||||
|
||||
root, err := c.GetWikiNode(ctx, spaceID, nodeToken)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
nodes, err := c.listWikiNodeDescendants(ctx, spaceID, root)
|
||||
if err != nil {
|
||||
return append([]WikiNode{root}, nodes...), err
|
||||
}
|
||||
return append([]WikiNode{root}, nodes...), nil
|
||||
}
|
||||
|
||||
func (c *Client) listWikiNodeDescendants(ctx context.Context, spaceID string, root WikiNode) ([]WikiNode, error) {
|
||||
if !root.HasChild {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
children, err := c.ListWikiNodes(ctx, spaceID, root.NodeToken)
|
||||
if err != nil {
|
||||
wrappedErr := fmt.Errorf("list children of %s: %w", root.NodeToken, err)
|
||||
logger.Warnf(ctx, "[Feishu] partial wiki node listing failure: space=%s node=%s err=%v",
|
||||
spaceID, root.NodeToken, err)
|
||||
return nil, &PartialWikiNodeListError{
|
||||
Failures: []WikiNodeListFailure{{
|
||||
Node: root,
|
||||
Err: wrappedErr,
|
||||
}},
|
||||
}
|
||||
}
|
||||
|
||||
var allNodes []WikiNode
|
||||
var failures []WikiNodeListFailure
|
||||
var walk func(nodes []WikiNode)
|
||||
|
||||
walk = func(nodes []WikiNode) {
|
||||
for _, node := range nodes {
|
||||
allNodes = append(allNodes, node)
|
||||
if !node.HasChild {
|
||||
continue
|
||||
}
|
||||
|
||||
grandChildren, err := c.ListWikiNodes(ctx, spaceID, node.NodeToken)
|
||||
if err != nil {
|
||||
wrappedErr := fmt.Errorf("list children of %s: %w", node.NodeToken, err)
|
||||
failures = append(failures, WikiNodeListFailure{
|
||||
Node: node,
|
||||
Err: wrappedErr,
|
||||
})
|
||||
logger.Warnf(ctx, "[Feishu] partial wiki node listing failure: space=%s node=%s err=%v",
|
||||
spaceID, node.NodeToken, err)
|
||||
continue
|
||||
}
|
||||
walk(grandChildren)
|
||||
}
|
||||
}
|
||||
|
||||
walk(children)
|
||||
if len(failures) > 0 {
|
||||
return allNodes, &PartialWikiNodeListError{Failures: failures}
|
||||
}
|
||||
return allNodes, nil
|
||||
}
|
||||
|
||||
// getDocumentRawContent retrieves the raw text content of a Feishu docx document.
|
||||
// This returns plain text (not rich text / block structure).
|
||||
// Deprecated: prefer ExportAndDownload which preserves formatting.
|
||||
func (c *Client) getDocumentRawContent(ctx context.Context, documentID string) (string, error) {
|
||||
path := fmt.Sprintf("/open-apis/docx/v1/documents/%s/raw_content", documentID)
|
||||
|
||||
var resp docRawContentResponse
|
||||
if err := c.DoRequest(ctx, http.MethodGet, path, nil, &resp); err != nil {
|
||||
return "", fmt.Errorf("get document raw content: %w", err)
|
||||
}
|
||||
if resp.Code != 0 {
|
||||
return "", fmt.Errorf("get document raw content error: code=%d msg=%s", resp.Code, resp.Msg)
|
||||
}
|
||||
|
||||
return resp.Data.Content, nil
|
||||
}
|
||||
|
||||
// Ping verifies the credentials by attempting to get a tenant access token.
|
||||
func (c *Client) Ping(ctx context.Context) error {
|
||||
_, err := c.GetTenantAccessToken(ctx)
|
||||
return err
|
||||
}
|
||||
|
||||
// ──────────────────────────────────────────────────────────────────────
|
||||
// Export task API: export docx/sheet/bitable to downloadable files
|
||||
//
|
||||
// Flow:
|
||||
// 1. POST /drive/v1/export_tasks → create export task, get ticket
|
||||
// 2. GET /drive/v1/export_tasks/:ticket → poll until status=0 (success)
|
||||
// 3. GET /drive/v1/export_tasks/file/:ticket/download → download file bytes
|
||||
// ──────────────────────────────────────────────────────────────────────
|
||||
|
||||
// createExportTask creates an async export task for a Feishu document.
|
||||
// - token: the obj_token of the document (e.g. docx token, sheet token)
|
||||
// - objType: the Feishu obj_type ("docx", "doc", "sheet", "bitable")
|
||||
// - fileExtension: desired output format ("docx", "xlsx", "pdf")
|
||||
func (c *Client) createExportTask(ctx context.Context, token, objType, fileExtension string) (string, error) {
|
||||
body := map[string]string{
|
||||
"file_extension": fileExtension,
|
||||
"token": token,
|
||||
"type": objType,
|
||||
}
|
||||
|
||||
var resp ExportTaskCreateResponse
|
||||
if err := c.DoRequest(ctx, http.MethodPost, "/open-apis/drive/v1/export_tasks", body, &resp); err != nil {
|
||||
return "", fmt.Errorf("create export task: %w", err)
|
||||
}
|
||||
if resp.Code != 0 {
|
||||
return "", fmt.Errorf("create export task error: code=%d msg=%s", resp.Code, resp.Msg)
|
||||
}
|
||||
|
||||
return resp.Data.Ticket, nil
|
||||
}
|
||||
|
||||
// getExportTaskStatus polls the status of an export task.
|
||||
// Returns (fileToken, fileName, error). fileToken is non-empty only when the job succeeds.
|
||||
// The token parameter is the obj_token of the document being exported (required by the API).
|
||||
func (c *Client) getExportTaskStatus(ctx context.Context, ticket string, token string) (string, string, error) {
|
||||
path := fmt.Sprintf("/open-apis/drive/v1/export_tasks/%s?token=%s", ticket, token)
|
||||
|
||||
var resp ExportTaskStatusResponse
|
||||
if err := c.DoRequest(ctx, http.MethodGet, path, nil, &resp); err != nil {
|
||||
return "", "", fmt.Errorf("get export task status: %w", err)
|
||||
}
|
||||
if resp.Code != 0 {
|
||||
return "", "", fmt.Errorf("get export task status error: code=%d msg=%s", resp.Code, resp.Msg)
|
||||
}
|
||||
|
||||
r := resp.Data.Result
|
||||
switch r.JobStatus {
|
||||
case 0: // success
|
||||
return r.FileToken, r.FileName, nil
|
||||
case 1, 2: // initializing, processing
|
||||
return "", "", nil // not ready yet
|
||||
default:
|
||||
return "", "", fmt.Errorf("export task failed: status=%d msg=%s", r.JobStatus, r.JobErrorMsg)
|
||||
}
|
||||
}
|
||||
|
||||
// downloadExportFile downloads the exported file by its file_token.
|
||||
// The file_token is returned by getExportTaskStatus when the export job completes.
|
||||
// The file must be downloaded within 10 minutes of export completion.
|
||||
func (c *Client) downloadExportFile(ctx context.Context, fileToken string) ([]byte, error) {
|
||||
path := fmt.Sprintf("/open-apis/drive/v1/export_tasks/file/%s/download", fileToken)
|
||||
return c.downloadRawBytes(ctx, path)
|
||||
}
|
||||
|
||||
// ExportAndDownload is a high-level helper that creates an export task, polls until
|
||||
// completion, and downloads the resulting file. Returns (fileBytes, fileName, error).
|
||||
//
|
||||
// Timeout: 60 seconds. Poll interval: 2 seconds.
|
||||
func (c *Client) ExportAndDownload(ctx context.Context, objToken, objType string) ([]byte, string, error) {
|
||||
// Determine export format
|
||||
fileExt, ok := ObjTypeToExportFileExtension[objType]
|
||||
if !ok {
|
||||
return nil, "", fmt.Errorf("unsupported obj_type for export: %s", objType)
|
||||
}
|
||||
|
||||
exportType, ok := ObjTypeToExportType[objType]
|
||||
if !ok {
|
||||
return nil, "", fmt.Errorf("unsupported obj_type for export: %s", objType)
|
||||
}
|
||||
|
||||
// Step 1: create export task
|
||||
ticket, err := c.createExportTask(ctx, objToken, exportType, fileExt)
|
||||
if err != nil {
|
||||
return nil, "", err
|
||||
}
|
||||
|
||||
// Step 2: poll until ready (max 60s, every 2s)
|
||||
deadline := time.Now().Add(60 * time.Second)
|
||||
var fileToken, fileName string
|
||||
|
||||
for time.Now().Before(deadline) {
|
||||
fileToken, fileName, err = c.getExportTaskStatus(ctx, ticket, objToken)
|
||||
if err != nil {
|
||||
return nil, "", err
|
||||
}
|
||||
if fileToken != "" {
|
||||
break // export ready
|
||||
}
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return nil, "", ctx.Err()
|
||||
case <-time.After(2 * time.Second):
|
||||
}
|
||||
}
|
||||
|
||||
if fileToken == "" {
|
||||
return nil, "", fmt.Errorf("export task timed out after 60s (ticket=%s)", ticket)
|
||||
}
|
||||
|
||||
// Step 3: download file using file_token (NOT ticket)
|
||||
data, err := c.downloadExportFile(ctx, fileToken)
|
||||
if err != nil {
|
||||
return nil, "", err
|
||||
}
|
||||
|
||||
// Build a sensible file name
|
||||
if fileName == "" {
|
||||
fileName = "export" + ExportFileExtToSuffix[fileExt]
|
||||
}
|
||||
|
||||
return data, fileName, nil
|
||||
}
|
||||
|
||||
// ──────────────────────────────────────────────────────────────────────
|
||||
// Drive file download: for "file" type wiki nodes (uploaded PDF/Word/etc.)
|
||||
// ──────────────────────────────────────────────────────────────────────
|
||||
|
||||
// DownloadDriveFile downloads a file from Feishu Drive by its file token.
|
||||
// Used for wiki nodes with obj_type="file" (user-uploaded PDF, Word, images, etc.).
|
||||
func (c *Client) DownloadDriveFile(ctx context.Context, fileToken string) ([]byte, error) {
|
||||
path := fmt.Sprintf("/open-apis/drive/v1/files/%s/download", fileToken)
|
||||
return c.downloadRawBytes(ctx, path)
|
||||
}
|
||||
|
||||
// downloadMediaFile downloads embedded media (attachments/images referenced by
|
||||
// document File/Image blocks) by its media token. Embedded block media live in a
|
||||
// different token space than standalone Drive files and must use the /medias/
|
||||
// endpoint rather than /files/.
|
||||
func (c *Client) downloadMediaFile(ctx context.Context, fileToken string) ([]byte, error) {
|
||||
path := fmt.Sprintf("/open-apis/drive/v1/medias/%s/download", url.PathEscape(fileToken))
|
||||
return c.downloadRawBytes(ctx, path)
|
||||
}
|
||||
|
||||
// downloadRawBytes performs an authenticated GET and returns the raw response body.
|
||||
func (c *Client) downloadRawBytes(ctx context.Context, path string) ([]byte, error) {
|
||||
token, err := c.GetTenantAccessToken(ctx)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
url := c.baseURL + path
|
||||
var lastErr error
|
||||
|
||||
for attempt := 0; attempt <= feishuMaxRetries; attempt++ {
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("create download request: %w", err)
|
||||
}
|
||||
req.Header.Set("Authorization", "Bearer "+token)
|
||||
|
||||
if attempt == 0 {
|
||||
logger.Infof(ctx, "[Feishu] download GET %s", path)
|
||||
} else {
|
||||
logger.Infof(ctx, "[Feishu] download GET %s (retry %d/%d)", path, attempt, feishuMaxRetries)
|
||||
}
|
||||
|
||||
resp, err := c.httpClient.Do(req)
|
||||
if err != nil {
|
||||
lastErr = fmt.Errorf("download request: %w", err)
|
||||
if attempt < feishuMaxRetries {
|
||||
if sErr := sleepCtx(ctx, feishuRetryBackoff[attempt]); sErr != nil {
|
||||
return nil, sErr
|
||||
}
|
||||
continue
|
||||
}
|
||||
return nil, lastErr
|
||||
}
|
||||
|
||||
if resp.StatusCode == http.StatusTooManyRequests {
|
||||
body, _ := io.ReadAll(resp.Body)
|
||||
resp.Body.Close()
|
||||
wait := parseRetryAfter(resp.Header.Get("Retry-After"), feishuRetryBackoff[min(attempt, len(feishuRetryBackoff)-1)])
|
||||
lastErr = fmt.Errorf("download rate limited: status=429 body=%s", truncate(string(body), 500))
|
||||
if attempt < feishuMaxRetries {
|
||||
if sErr := sleepCtx(ctx, wait); sErr != nil {
|
||||
return nil, sErr
|
||||
}
|
||||
continue
|
||||
}
|
||||
return nil, lastErr
|
||||
}
|
||||
|
||||
if resp.StatusCode >= 500 && resp.StatusCode < 600 {
|
||||
body, _ := io.ReadAll(resp.Body)
|
||||
resp.Body.Close()
|
||||
lastErr = fmt.Errorf("download server error: status=%d body=%s", resp.StatusCode, truncate(string(body), 500))
|
||||
if attempt < feishuMax5xxRetries {
|
||||
if sErr := sleepCtx(ctx, feishuRetry5xxDelay); sErr != nil {
|
||||
return nil, sErr
|
||||
}
|
||||
continue
|
||||
}
|
||||
return nil, lastErr
|
||||
}
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
body, _ := io.ReadAll(resp.Body)
|
||||
resp.Body.Close()
|
||||
logger.Errorf(ctx, "[Feishu] download GET %s → status=%d body=%s", path, resp.StatusCode, truncate(string(body), 500))
|
||||
return nil, fmt.Errorf("download failed: status=%d body=%s", resp.StatusCode, string(body))
|
||||
}
|
||||
|
||||
data, readErr := io.ReadAll(io.LimitReader(resp.Body, maxFeishuDownloadBytes+1))
|
||||
if readErr == nil && int64(len(data)) > maxFeishuDownloadBytes {
|
||||
resp.Body.Close()
|
||||
return nil, fmt.Errorf("download exceeds max size (%d bytes): %s", maxFeishuDownloadBytes, path)
|
||||
}
|
||||
resp.Body.Close()
|
||||
if readErr != nil {
|
||||
lastErr = fmt.Errorf("read download body: %w", readErr)
|
||||
if attempt < feishuMaxRetries {
|
||||
if sErr := sleepCtx(ctx, feishuRetryBackoff[attempt]); sErr != nil {
|
||||
return nil, sErr
|
||||
}
|
||||
continue
|
||||
}
|
||||
return nil, lastErr
|
||||
}
|
||||
|
||||
logger.Infof(ctx, "[Feishu] download GET %s → OK, %d bytes", path, len(data))
|
||||
return data, nil
|
||||
}
|
||||
|
||||
return nil, lastErr
|
||||
}
|
||||
|
||||
// ──────────────────────────────────────────────────────────────────────
|
||||
// Drive (云盘) file listing: for feishu_drive / lark_drive connectors.
|
||||
// Mirrors the wiki ListWikiNodes / ListWikiNodesRecursiveFrom shape so the
|
||||
// Drive connector's FetchStream mirrors the wiki connector's. See ADR-0001/0002.
|
||||
// ──────────────────────────────────────────────────────────────────────
|
||||
|
||||
// listDriveFiles lists files in a Drive folder (non-recursive), one page at a
|
||||
// time. Pass pageToken="" for the first page; the returned nextPageToken is ""
|
||||
// when there are no more pages.
|
||||
//
|
||||
// folderToken == "" is rejected: the root folder is not paginated and does
|
||||
// not return shortcuts (Feishu API limitation), which would silently drop
|
||||
// content and risk an unbounded single response. See ADR-0004.
|
||||
func (c *Client) listDriveFiles(ctx context.Context, folderToken, pageToken string) ([]DriveFile, string, error) {
|
||||
if folderToken == "" {
|
||||
return nil, "", fmt.Errorf("root folder not supported; specify a concrete folder_token (root folder is not paginated and does not return shortcuts)")
|
||||
}
|
||||
|
||||
path := "/open-apis/drive/v1/files?folder_token=" + url.QueryEscape(folderToken)
|
||||
path += "&page_size=200" // max
|
||||
path += "&order_by=EditedTime&direction=DESC"
|
||||
if pageToken != "" {
|
||||
path += "&page_token=" + url.QueryEscape(pageToken)
|
||||
}
|
||||
|
||||
var resp DriveFileListResponse
|
||||
if err := c.DoRequest(ctx, http.MethodGet, path, nil, &resp); err != nil {
|
||||
return nil, "", fmt.Errorf("list drive files: %w", err)
|
||||
}
|
||||
if resp.Code != 0 {
|
||||
return nil, "", fmt.Errorf("list drive files error: code=%d msg=%s", resp.Code, resp.Msg)
|
||||
}
|
||||
|
||||
logger.Infof(ctx, "[FeishuDrive] listDriveFiles: folder=%s got %d files, has_more=%v",
|
||||
folderToken, len(resp.Data.Files), resp.Data.HasMore)
|
||||
return resp.Data.Files, resp.Data.NextPageToken, nil
|
||||
}
|
||||
|
||||
// GetDriveFolderMeta returns the metadata (name, owner, etc.) of a single Drive
|
||||
// folder. Used to resolve a root folder's human-readable name - the list API
|
||||
// only returns the folder's children, not the folder itself.
|
||||
//
|
||||
// GET /open-apis/drive/explorer/v2/folder/:folderToken/meta
|
||||
func (c *Client) GetDriveFolderMeta(ctx context.Context, folderToken string) (driveFolderMetaResponse, error) {
|
||||
var resp driveFolderMetaResponse
|
||||
if folderToken == "" {
|
||||
return resp, fmt.Errorf("root folder not supported; specify a concrete folder_token")
|
||||
}
|
||||
path := "/open-apis/drive/explorer/v2/folder/" + url.QueryEscape(folderToken) + "/meta"
|
||||
if err := c.DoRequest(ctx, http.MethodGet, path, nil, &resp); err != nil {
|
||||
return resp, fmt.Errorf("get drive folder meta: %w", err)
|
||||
}
|
||||
if resp.Code != 0 {
|
||||
return resp, fmt.Errorf("get drive folder meta error: code=%d msg=%s", resp.Code, resp.Msg)
|
||||
}
|
||||
return resp, nil
|
||||
}
|
||||
|
||||
// ListDriveFilesAllPages lists every direct child of a folder across all pages.
|
||||
func (c *Client) ListDriveFilesAllPages(ctx context.Context, folderToken string) ([]DriveFile, error) {
|
||||
var all []DriveFile
|
||||
pageToken := ""
|
||||
for {
|
||||
files, next, err := c.listDriveFiles(ctx, folderToken, pageToken)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
all = append(all, files...)
|
||||
if next == "" {
|
||||
break
|
||||
}
|
||||
pageToken = next
|
||||
}
|
||||
return all, nil
|
||||
}
|
||||
|
||||
// ListDriveFilesRecursiveFrom walks a Drive folder subtree depth-first,
|
||||
// returning all non-folder files. Mirrors ListWikiNodesRecursiveFrom.
|
||||
//
|
||||
// - folder -> recurse (visited is a pure-defensive cycle guard; Drive folders
|
||||
// have no shortcut concept so cycles are not expected - see glossary).
|
||||
// - shortcut -> expand to its target (target_type is never "folder", verified)
|
||||
// and include the target as a regular file. No extra API call: shortcut_info
|
||||
// is returned by the list API.
|
||||
// - other -> collect.
|
||||
//
|
||||
// Partial failures (a sub-folder listing returns an error) are collected into a
|
||||
// *PartialDriveFileListError and the walk continues, mirroring the wiki
|
||||
// connector's PartialWikiNodeListError semantics.
|
||||
func (c *Client) ListDriveFilesRecursiveFrom(ctx context.Context, folderToken string) ([]DriveFile, error) {
|
||||
visited := make(map[string]bool)
|
||||
var all []DriveFile
|
||||
var failures []DriveFileListFailure
|
||||
|
||||
var walk func(folderToken string)
|
||||
walk = func(folderToken string) {
|
||||
if visited[folderToken] {
|
||||
return
|
||||
}
|
||||
visited[folderToken] = true
|
||||
|
||||
files, err := c.ListDriveFilesAllPages(ctx, folderToken)
|
||||
if err != nil {
|
||||
wrappedErr := fmt.Errorf("list children of %s: %w", folderToken, err)
|
||||
failures = append(failures, DriveFileListFailure{
|
||||
FolderToken: folderToken,
|
||||
Err: wrappedErr,
|
||||
})
|
||||
logger.Warnf(ctx, "[FeishuDrive] partial drive file listing failure: folder=%s err=%v",
|
||||
folderToken, err)
|
||||
return
|
||||
}
|
||||
|
||||
for _, f := range files {
|
||||
switch f.Type {
|
||||
case "folder":
|
||||
walk(f.Token)
|
||||
case "shortcut":
|
||||
// Expand to target. target_type is never "folder" (verified), so
|
||||
// no recursion here - the target is a regular file.
|
||||
if f.ShortcutInfo != nil && f.ShortcutInfo.TargetToken != "" {
|
||||
expanded := DriveFile{
|
||||
Token: f.ShortcutInfo.TargetToken,
|
||||
Name: f.Name,
|
||||
Type: f.ShortcutInfo.TargetType,
|
||||
ParentToken: f.ParentToken,
|
||||
URL: f.URL,
|
||||
CreatedTime: f.CreatedTime,
|
||||
ModifiedTime: f.ModifiedTime,
|
||||
OwnerID: f.OwnerID,
|
||||
}
|
||||
all = append(all, expanded)
|
||||
}
|
||||
default:
|
||||
all = append(all, f)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
walk(folderToken)
|
||||
if len(failures) > 0 {
|
||||
return all, &PartialDriveFileListError{Failures: failures}
|
||||
}
|
||||
return all, nil
|
||||
}
|
||||
@@ -0,0 +1,173 @@
|
||||
package core
|
||||
|
||||
import (
|
||||
"context"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
// retryTestServer builds a server that always answers the auth-token call and
|
||||
// routes the given target path to h, so tests can drive DoRequest's retry loop.
|
||||
func retryTestServer(target string, h http.HandlerFunc) (*httptest.Server, *Config) {
|
||||
mux := http.NewServeMux()
|
||||
mux.HandleFunc("/open-apis/auth/v3/tenant_access_token/internal", func(w http.ResponseWriter, r *http.Request) {
|
||||
writeJSON(w, TokenResponse{
|
||||
ApiResponse: ApiResponse{Code: 0},
|
||||
TenantAccessToken: "fake-token",
|
||||
Expire: 7200,
|
||||
})
|
||||
})
|
||||
mux.HandleFunc(target, h)
|
||||
ts := httptest.NewServer(mux)
|
||||
return ts, &Config{AppID: "a", AppSecret: "b", BaseURL: ts.URL}
|
||||
}
|
||||
|
||||
func TestDoRequest_RetriesOn429ThenSucceeds(t *testing.T) {
|
||||
var attempts int
|
||||
ts, cfg := retryTestServer("/target", func(w http.ResponseWriter, r *http.Request) {
|
||||
attempts++
|
||||
if attempts == 1 {
|
||||
// "0" is coerced to a short delay inside the client so the test stays fast.
|
||||
w.Header().Set("Retry-After", "0")
|
||||
w.WriteHeader(http.StatusTooManyRequests)
|
||||
_, _ = io.WriteString(w, `{"code":99991400,"msg":"rate limited"}`)
|
||||
return
|
||||
}
|
||||
writeJSON(w, ApiResponse{Code: 0})
|
||||
})
|
||||
defer ts.Close()
|
||||
|
||||
c := NewClient(cfg)
|
||||
var resp ApiResponse
|
||||
if err := c.DoRequest(context.Background(), http.MethodGet, "/target", nil, &resp); err != nil {
|
||||
t.Fatalf("expected success after retry, got %v", err)
|
||||
}
|
||||
if attempts < 2 {
|
||||
t.Errorf("attempts = %d, want >= 2 (should retry after 429)", attempts)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDoRequest_429ExhaustsRetries(t *testing.T) {
|
||||
var attempts int
|
||||
ts, cfg := retryTestServer("/target", func(w http.ResponseWriter, r *http.Request) {
|
||||
attempts++
|
||||
w.Header().Set("Retry-After", "0")
|
||||
w.WriteHeader(http.StatusTooManyRequests)
|
||||
_, _ = io.WriteString(w, `{"code":99991400,"msg":"rate limited"}`)
|
||||
})
|
||||
defer ts.Close()
|
||||
|
||||
c := NewClient(cfg)
|
||||
err := c.DoRequest(context.Background(), http.MethodGet, "/target", nil, nil)
|
||||
if err == nil {
|
||||
t.Fatal("expected error when 429s exceed the retry budget")
|
||||
}
|
||||
if attempts != 4 { // initial + 3 retries
|
||||
t.Errorf("attempts = %d, want 4 (1 + 3 retries)", attempts)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDoRequest_5xxRetriesOnce(t *testing.T) {
|
||||
var attempts int
|
||||
ts, cfg := retryTestServer("/target", func(w http.ResponseWriter, r *http.Request) {
|
||||
attempts++
|
||||
w.WriteHeader(http.StatusInternalServerError)
|
||||
_, _ = io.WriteString(w, `{"code":1,"msg":"internal error"}`)
|
||||
})
|
||||
defer ts.Close()
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
|
||||
c := NewClient(cfg)
|
||||
if err := c.DoRequest(ctx, http.MethodGet, "/target", nil, nil); err == nil {
|
||||
t.Fatal("expected error after 5xx exhaustion")
|
||||
}
|
||||
if attempts != 2 { // initial + 1 retry
|
||||
t.Errorf("attempts = %d, want 2 (5xx retries exactly once)", attempts)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDoRequest_4xxNotRetried(t *testing.T) {
|
||||
var attempts int
|
||||
ts, cfg := retryTestServer("/target", func(w http.ResponseWriter, r *http.Request) {
|
||||
attempts++
|
||||
w.WriteHeader(http.StatusBadRequest)
|
||||
_, _ = io.WriteString(w, `{"code":1,"msg":"bad request"}`)
|
||||
})
|
||||
defer ts.Close()
|
||||
|
||||
c := NewClient(cfg)
|
||||
if err := c.DoRequest(context.Background(), http.MethodGet, "/target", nil, nil); err == nil {
|
||||
t.Fatal("expected error on 400")
|
||||
}
|
||||
if attempts != 1 {
|
||||
t.Errorf("attempts = %d, want 1 (non-429/5xx 4xx must not retry)", attempts)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDownloadRawBytes_RetriesOn429ThenSucceeds(t *testing.T) {
|
||||
var attempts int
|
||||
ts, cfg := retryTestServer("/dl", func(w http.ResponseWriter, r *http.Request) {
|
||||
attempts++
|
||||
if attempts == 1 {
|
||||
w.Header().Set("Retry-After", "0")
|
||||
w.WriteHeader(http.StatusTooManyRequests)
|
||||
return
|
||||
}
|
||||
w.Header().Set("Content-Type", "application/octet-stream")
|
||||
_, _ = w.Write([]byte("payload-bytes"))
|
||||
})
|
||||
defer ts.Close()
|
||||
|
||||
c := NewClient(cfg)
|
||||
data, err := c.downloadRawBytes(context.Background(), "/dl")
|
||||
if err != nil {
|
||||
t.Fatalf("expected success after retry, got %v", err)
|
||||
}
|
||||
if string(data) != "payload-bytes" {
|
||||
t.Errorf("data = %q, want %q", string(data), "payload-bytes")
|
||||
}
|
||||
if attempts < 2 {
|
||||
t.Errorf("attempts = %d, want >= 2 (should retry download after 429)", attempts)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDownloadRawBytes_4xxNotRetried(t *testing.T) {
|
||||
var attempts int
|
||||
ts, cfg := retryTestServer("/dl", func(w http.ResponseWriter, r *http.Request) {
|
||||
attempts++
|
||||
w.WriteHeader(http.StatusForbidden)
|
||||
})
|
||||
defer ts.Close()
|
||||
|
||||
c := NewClient(cfg)
|
||||
if _, err := c.downloadRawBytes(context.Background(), "/dl"); err == nil {
|
||||
t.Fatal("expected error on 403")
|
||||
}
|
||||
if attempts != 1 {
|
||||
t.Errorf("attempts = %d, want 1 (403 must not retry)", attempts)
|
||||
}
|
||||
}
|
||||
|
||||
func TestParseRetryAfter(t *testing.T) {
|
||||
fallback := 5 * time.Second
|
||||
tests := []struct {
|
||||
header string
|
||||
want time.Duration
|
||||
}{
|
||||
{"", fallback},
|
||||
{"0", 100 * time.Millisecond},
|
||||
{"-1", 100 * time.Millisecond}, // negative coerced to a short delay
|
||||
{"3", 3 * time.Second},
|
||||
{"abc", fallback}, // unparseable
|
||||
}
|
||||
for _, tt := range tests {
|
||||
if got := parseRetryAfter(tt.header, fallback); got != tt.want {
|
||||
t.Errorf("parseRetryAfter(%q) = %v, want %v", tt.header, got, tt.want)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,108 @@
|
||||
package core
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// feishuFailure classifies a raw connector/API error into a stable i18n code
|
||||
// (+ optional numeric feishu code) and an English fallback message. The frontend
|
||||
// localises by code; the raw status/body/log_id stays in the server logs and
|
||||
// must never appear in the code, codeValue or fallback.
|
||||
func TestFeishuFailure(t *testing.T) {
|
||||
cases := []struct {
|
||||
name string
|
||||
err error
|
||||
wantCode string
|
||||
wantCodeValue string
|
||||
fallbackHas []string
|
||||
noLeak []string // must appear in neither code nor fallback
|
||||
}{
|
||||
{
|
||||
name: "rate limited",
|
||||
err: errors.New("feishu rate limited: status=429 body={\"code\":99991400,\"msg\":\"too many request\"}"),
|
||||
wantCode: "feishu_rate_limited",
|
||||
fallbackHas: []string{"retry"},
|
||||
noLeak: []string{"body=", "{", "99991400"},
|
||||
},
|
||||
{
|
||||
name: "server 5xx",
|
||||
err: errors.New("feishu server error: status=500 body={\"code\":1663}"),
|
||||
wantCode: "feishu_server_unavailable",
|
||||
noLeak: []string{"body=", "{"},
|
||||
},
|
||||
{
|
||||
name: "export timeout",
|
||||
err: errors.New("export 季度报告 (docx): export task timed out after 60s (ticket=abc)"),
|
||||
wantCode: "feishu_timeout",
|
||||
noLeak: []string{"ticket=", "abc"},
|
||||
},
|
||||
{
|
||||
name: "api error carries the feishu code as a param",
|
||||
err: errors.New("feishu api error: status=500 body={\"code\":1663,\"msg\":\"internal error\",\"Error\":{\"log_id\":\"20260\"}}"),
|
||||
wantCode: "feishu_api_error",
|
||||
wantCodeValue: "1663",
|
||||
noLeak: []string{"log_id", "body=", "{"},
|
||||
},
|
||||
{
|
||||
name: "api error without a code is generic",
|
||||
err: errors.New("feishu api error: status=502 body=bad gateway"),
|
||||
wantCode: "feishu_api_error_generic",
|
||||
},
|
||||
{
|
||||
name: "auth failure is actionable, not a retry",
|
||||
err: errors.New("feishu auth error: code=99991663 msg=Invalid access token"),
|
||||
wantCode: "feishu_auth_or_permission",
|
||||
noLeak: []string{"Invalid access token", "99991663"},
|
||||
},
|
||||
{
|
||||
name: "unknown error falls back to sync_failed",
|
||||
err: errors.New("some totally unexpected failure"),
|
||||
wantCode: "sync_failed",
|
||||
noLeak: []string{"totally unexpected"},
|
||||
},
|
||||
{
|
||||
// "decode"/"encode"/"unicode" all contain the substring "code" but
|
||||
// are not Feishu API errors — they must not be classified as one.
|
||||
name: "decode error is not a feishu api error",
|
||||
err: errors.New("failed to decode response body"),
|
||||
wantCode: "sync_failed",
|
||||
},
|
||||
{
|
||||
// The failed item is retried on the next sync (the cursor is not
|
||||
// advanced), which for a manual-only source is not automatic. The
|
||||
// message must not promise an unconditional automatic retry.
|
||||
name: "transient message promises retry on next sync, not automatic",
|
||||
err: errors.New("feishu rate limited: status=429"),
|
||||
wantCode: "feishu_rate_limited",
|
||||
fallbackHas: []string{"next sync"},
|
||||
noLeak: []string{"automatically"},
|
||||
},
|
||||
}
|
||||
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
code, codeValue, fallback := feishuFailure(tc.err)
|
||||
if code != tc.wantCode {
|
||||
t.Errorf("code = %q, want %q", code, tc.wantCode)
|
||||
}
|
||||
if codeValue != tc.wantCodeValue {
|
||||
t.Errorf("codeValue = %q, want %q", codeValue, tc.wantCodeValue)
|
||||
}
|
||||
if strings.TrimSpace(fallback) == "" {
|
||||
t.Errorf("fallback message must not be empty")
|
||||
}
|
||||
for _, s := range tc.fallbackHas {
|
||||
if !strings.Contains(strings.ToLower(fallback), s) {
|
||||
t.Errorf("fallback %q missing %q", fallback, s)
|
||||
}
|
||||
}
|
||||
for _, s := range tc.noLeak {
|
||||
if strings.Contains(code, s) || strings.Contains(fallback, s) {
|
||||
t.Errorf("raw detail %q leaked into code/fallback (code=%q fallback=%q)", s, code, fallback)
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,263 @@
|
||||
package core
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/datasource"
|
||||
"github.com/Tencent/WeKnora/internal/logger"
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
)
|
||||
|
||||
// engine.go holds the single generic streaming sync engine shared by the wiki
|
||||
// Connector and the Drive DriveConnector. The per-connector differences (node
|
||||
// type, listing API, edit-time field, cursor wire format, fetch dispatch,
|
||||
// log tag) are isolated behind the NodeOps adapter interface. FetchAll /
|
||||
// FetchIncremental are thin wrappers over the same engine that collect Emits
|
||||
// instead of streaming them.
|
||||
//
|
||||
// Behaviour note (deliberate, see ADR-0005 / design §2.4):
|
||||
// - Resume/incremental fast-path: a node recorded at its current edit time
|
||||
// is skipped, keeping the cursor entry.
|
||||
// - A fetch failure does NOT advance the cursor: the prior edit time is
|
||||
// retained so the node is retried next run instead of being permanently
|
||||
// skipped on a transient export failure (Tencent/WeKnora#2136). This now
|
||||
// also holds for the FetchIncremental path (previously it advanced the
|
||||
// cursor before fetching, a latent #2136 bug).
|
||||
// - Logs use the "stream progress/summary" wording uniformly; the
|
||||
// FetchIncremental path additionally gains per-100 progress + tally
|
||||
// summary logs it did not Emit before (log-only change).
|
||||
|
||||
// NodeOps adapts one connector's node type to the shared sync engine. Every
|
||||
// method is a pure accessor or a thin wrapper - no engine logic lives here.
|
||||
type NodeOps[N any] interface {
|
||||
// List returns every syncable node under resourceID. A non-nil partial
|
||||
// (with err == nil) signals a partial listing: nodes is still usable and
|
||||
// the sync continues, but the caller surfaces the failed sub-trees via
|
||||
// ListFailureItems. A non-nil err is fatal and aborts the sync.
|
||||
List(ctx context.Context, client *Client, resourceID string) (nodes []N, partial error, err error)
|
||||
|
||||
Token(n N) string
|
||||
Title(n N) string
|
||||
ObjType(n N) string
|
||||
// EditTime is the change-detection timestamp string stored in the cursor.
|
||||
EditTime(n N) string
|
||||
|
||||
// fetch retrieves one node's content; (nil, nil) means an unsupported
|
||||
// type that yields no item.
|
||||
Fetch(ctx context.Context, client *Client, n N, resourceID string, multimodal bool) ([]*types.FetchedItem, error)
|
||||
|
||||
// ListFailureItems converts a partial-listing error into error FetchedItems.
|
||||
ListFailureItems(resourceID string, partial error) []types.FetchedItem
|
||||
// ResourceNoun is the noun in user-visible error text: "nodes" / "files".
|
||||
ResourceNoun() string
|
||||
// EmptyResourceIDsError is the connector-specific message when no resource
|
||||
// IDs are configured (wiki/drive text differs and is preserved verbatim).
|
||||
EmptyResourceIDsError() string
|
||||
// LogTag is the log prefix: "[Feishu]" / "[FeishuDrive]".
|
||||
LogTag() string
|
||||
|
||||
// DecodeCursorTimes extracts the per-resource edit-time map from a
|
||||
// persisted ConnectorCursor (nil-safe: returns nil when absent).
|
||||
DecodeCursorTimes(m map[string]interface{}) map[string]map[string]string
|
||||
// EncodeCursor wraps the engine's internal times map into the connector's
|
||||
// wire-format SyncCursor. JSON-marshals for snapshot isolation, mirroring
|
||||
// the original FeishuCursor.toSyncCursor / FeishuDriveCursor.toSyncCursor.
|
||||
EncodeCursor(times map[string]map[string]string, lastSync time.Time) *types.SyncCursor
|
||||
}
|
||||
|
||||
// CollectHandler is the StreamHandler used by FetchAll / FetchIncremental to
|
||||
// gather every Emitted item into a slice instead of streaming. Checkpoint is a
|
||||
// no-op: those paths return a single cursor at the end.
|
||||
type CollectHandler struct {
|
||||
items []types.FetchedItem
|
||||
}
|
||||
|
||||
func (h *CollectHandler) Emit(_ context.Context, item types.FetchedItem) error {
|
||||
h.items = append(h.items, item)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (h *CollectHandler) Checkpoint(_ context.Context, _ *types.SyncCursor) error { return nil }
|
||||
|
||||
// runSync is the single implementation behind FetchStream / FetchAll /
|
||||
// FetchIncremental. With cursor == nil it fetches everything (full sync);
|
||||
// with a cursor it skips nodes whose recorded edit time is unchanged
|
||||
// (incremental + resume). resourceIDs comes from the caller: FetchStream /
|
||||
// FetchIncremental pass config.ResourceIDs (after a non-empty check),
|
||||
// FetchAll passes its own argument without one.
|
||||
func runSync[N any](
|
||||
ctx context.Context, client *Client, config *types.DataSourceConfig,
|
||||
resourceIDs []string, cursor *types.SyncCursor, h datasource.StreamHandler,
|
||||
ops NodeOps[N],
|
||||
) (*types.SyncCursor, error) {
|
||||
var prevTimes map[string]map[string]string
|
||||
if cursor != nil && cursor.ConnectorCursor != nil {
|
||||
prevTimes = ops.DecodeCursorTimes(cursor.ConnectorCursor)
|
||||
}
|
||||
|
||||
newTimes := make(map[string]map[string]string)
|
||||
lastSync := time.Now()
|
||||
|
||||
processed := 0
|
||||
lastCheckpoint := time.Now()
|
||||
for _, resourceID := range resourceIDs {
|
||||
nodes, partial, err := ops.List(ctx, client, resourceID)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("list %s for resource %s: %w", ops.ResourceNoun(), resourceID, err)
|
||||
}
|
||||
if partial != nil {
|
||||
for _, item := range ops.ListFailureItems(resourceID, partial) {
|
||||
if eerr := h.Emit(ctx, item); eerr != nil {
|
||||
return nil, eerr
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
newTimes[resourceID] = make(map[string]string)
|
||||
// On a partial listing, carry prior edit times forward so a later full
|
||||
// listing can still detect changes and deletions.
|
||||
if partial != nil && prevTimes != nil {
|
||||
if prev, ok := prevTimes[resourceID]; ok {
|
||||
for tok, et := range prev {
|
||||
newTimes[resourceID][tok] = et
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
currentNodes := make(map[string]bool)
|
||||
tally := newFetchTally(len(nodes))
|
||||
for i, node := range nodes {
|
||||
tok := ops.Token(node)
|
||||
currentNodes[tok] = true
|
||||
editTimeStr := ops.EditTime(node)
|
||||
|
||||
var prevEdit string
|
||||
var hadPrev bool
|
||||
if prevTimes != nil {
|
||||
if prev, ok := prevTimes[resourceID]; ok {
|
||||
prevEdit, hadPrev = prev[tok]
|
||||
}
|
||||
}
|
||||
|
||||
// Resume/incremental fast-path: a node recorded at its current edit
|
||||
// time is unchanged (or already synced this run) - keep the record
|
||||
// and Skip re-fetching.
|
||||
if hadPrev && prevEdit == editTimeStr {
|
||||
newTimes[resourceID][tok] = editTimeStr
|
||||
continue
|
||||
}
|
||||
|
||||
items, ferr := ops.Fetch(ctx, client, node, resourceID, config.MultimodalEnabled)
|
||||
if ferr != nil {
|
||||
tally.fail()
|
||||
// Do NOT advance the cursor: the content was never fetched.
|
||||
// Retain the prior edit time (if any) so prev != current next
|
||||
// run and the node is retried, instead of being permanently
|
||||
// skipped on a transient export failure (Tencent/WeKnora#2136).
|
||||
if hadPrev {
|
||||
newTimes[resourceID][tok] = prevEdit
|
||||
}
|
||||
if eerr := h.Emit(ctx, types.FetchedItem{
|
||||
ExternalID: tok,
|
||||
Title: ops.Title(node),
|
||||
SourceResourceID: resourceID,
|
||||
Metadata: FeishuErrorItemMeta(ferr, nil),
|
||||
}); eerr != nil {
|
||||
return nil, eerr
|
||||
}
|
||||
} else {
|
||||
// Fetched, or an unsupported type (nothing to fetch): record
|
||||
// the current edit time so the node is not re-processed next run.
|
||||
newTimes[resourceID][tok] = editTimeStr
|
||||
if len(items) > 0 {
|
||||
tally.fetch()
|
||||
for _, it := range items {
|
||||
if eerr := h.Emit(ctx, *it); eerr != nil {
|
||||
return nil, eerr
|
||||
}
|
||||
}
|
||||
} else {
|
||||
// Unsupported type (mindnote/slides/…): no item.
|
||||
tally.Skip(ops.ObjType(node))
|
||||
}
|
||||
}
|
||||
|
||||
processed++
|
||||
if processed%FeishuStreamCheckpointInterval == 0 || time.Since(lastCheckpoint) >= FeishuStreamCheckpointMaxInterval {
|
||||
if cerr := h.Checkpoint(ctx, ops.EncodeCursor(newTimes, lastSync)); cerr != nil {
|
||||
logger.Warnf(ctx, "%s stream Checkpoint failed: %v", ops.LogTag(), cerr)
|
||||
}
|
||||
lastCheckpoint = time.Now()
|
||||
}
|
||||
if n := i + 1; n%100 == 0 {
|
||||
logger.Infof(ctx, "%s stream progress resource=%s %d/%d (%s)",
|
||||
ops.LogTag(), resourceID, n, len(nodes), tally.summary())
|
||||
}
|
||||
}
|
||||
|
||||
// Detect deleted nodes (only when the full tree was listed successfully).
|
||||
// A partial listing did not enumerate the whole subtree, so deletion
|
||||
// detection would false-positive.
|
||||
if partial == nil && prevTimes != nil {
|
||||
if prev, ok := prevTimes[resourceID]; ok {
|
||||
for tok := range prev {
|
||||
if !currentNodes[tok] {
|
||||
if eerr := h.Emit(ctx, types.FetchedItem{
|
||||
ExternalID: tok,
|
||||
IsDeleted: true,
|
||||
SourceResourceID: resourceID,
|
||||
}); eerr != nil {
|
||||
return nil, eerr
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
logger.Infof(ctx, "%s stream summary resource=%s %s", ops.LogTag(), resourceID, tally.summary())
|
||||
}
|
||||
|
||||
return ops.EncodeCursor(newTimes, lastSync), nil
|
||||
}
|
||||
|
||||
// FetchStreamEngine runs the streaming sync. FetchStream / FetchIncremental
|
||||
// shells differ only in whether they pass a cursor and how they collect
|
||||
// results; both route here.
|
||||
func FetchStreamEngine[N any](
|
||||
ctx context.Context, client *Client, config *types.DataSourceConfig,
|
||||
cursor *types.SyncCursor, h datasource.StreamHandler, ops NodeOps[N],
|
||||
) (*types.SyncCursor, error) {
|
||||
return runSync(ctx, client, config, config.ResourceIDs, cursor, h, ops)
|
||||
}
|
||||
|
||||
// FetchAllEngine runs a full sync, collecting every item into a slice. It
|
||||
// passes the caller-supplied resourceIDs without a non-empty check (mirroring
|
||||
// the original FetchAll which accepted an empty list and returned no items).
|
||||
// On error any partially collected items are discarded, matching the original
|
||||
// behaviour (a fatal list error dropped everything collected so far).
|
||||
func FetchAllEngine[N any](
|
||||
ctx context.Context, client *Client, config *types.DataSourceConfig,
|
||||
resourceIDs []string, ops NodeOps[N],
|
||||
) ([]types.FetchedItem, error) {
|
||||
ch := &CollectHandler{}
|
||||
if _, err := runSync(ctx, client, config, resourceIDs, nil, ch, ops); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return ch.items, nil
|
||||
}
|
||||
|
||||
// FetchIncrementalEngine runs an incremental sync against cursor, collecting
|
||||
// items. resourceIDs come from config.ResourceIDs (non-empty check is the
|
||||
// caller's responsibility).
|
||||
func FetchIncrementalEngine[N any](
|
||||
ctx context.Context, client *Client, config *types.DataSourceConfig,
|
||||
cursor *types.SyncCursor, ops NodeOps[N],
|
||||
) ([]types.FetchedItem, *types.SyncCursor, error) {
|
||||
ch := &CollectHandler{}
|
||||
next, err := runSync(ctx, client, config, config.ResourceIDs, cursor, ch, ops)
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
return ch.items, next, nil
|
||||
}
|
||||
@@ -0,0 +1,54 @@
|
||||
package core
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"os"
|
||||
"testing"
|
||||
|
||||
secutils "github.com/Tencent/WeKnora/internal/utils"
|
||||
)
|
||||
|
||||
func TestMain(m *testing.M) {
|
||||
os.Setenv("SSRF_WHITELIST", "127.0.0.1,localhost,open.feishu.cn,open.larksuite.com")
|
||||
secutils.ResetSSRFWhitelistForTest()
|
||||
os.Exit(m.Run())
|
||||
}
|
||||
|
||||
func writeJSON(w http.ResponseWriter, v interface{}) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
json.NewEncoder(w).Encode(v)
|
||||
}
|
||||
|
||||
// blk constructors mirroring wiki/connector_golden_test.go, for core-internal
|
||||
// block-rendering tests (DocxBlock is same-package here, no core. prefix).
|
||||
func sheetBlk(id, token string) DocxBlock {
|
||||
return DocxBlock{BlockID: id, BlockType: BlockTypeSheet, Sheet: &BlockTokenRef{Token: token}}
|
||||
}
|
||||
|
||||
func bitableBlk(id, token string) DocxBlock {
|
||||
return DocxBlock{BlockID: id, BlockType: BlockTypeBitable, Bitable: &BlockTokenRef{Token: token}}
|
||||
}
|
||||
|
||||
func imageBlk(id, token string) DocxBlock {
|
||||
return DocxBlock{BlockID: id, BlockType: BlockTypeImage, Image: &BlockTokenRef{Token: token}}
|
||||
}
|
||||
|
||||
func fileBlk(id, token, name string) DocxBlock {
|
||||
return DocxBlock{BlockID: id, BlockType: BlockTypeFile, File: &BlockFileRef{Token: token, Name: name}}
|
||||
}
|
||||
|
||||
func cellBlk(id string) DocxBlock {
|
||||
return DocxBlock{BlockID: id, BlockType: BlockTypeTableCell, Children: []string{id + "_txt"}}
|
||||
}
|
||||
|
||||
func cellTextBlk(id, text string) DocxBlock {
|
||||
return DocxBlock{BlockID: id + "_txt", BlockType: BlockTypeText, Text: txt(text)}
|
||||
}
|
||||
|
||||
func tableBlk(id string, cols int, cellIDs ...string) DocxBlock {
|
||||
b := DocxBlock{BlockID: id, BlockType: BlockTypeTable}
|
||||
b.Table = &BlockTable{Cells: cellIDs}
|
||||
b.Table.Property = &BlockTableProperty{ColumnSize: cols}
|
||||
return b
|
||||
}
|
||||
@@ -0,0 +1,320 @@
|
||||
package core
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// sheetReader is the subset of *Client that blocksToMarkdown needs, so the
|
||||
// converter can be unit-tested with a fake or nil client.
|
||||
type sheetReader interface {
|
||||
readSheetRange(ctx context.Context, embedToken string) ([][]string, bool, error)
|
||||
readBitableRecords(ctx context.Context, embedToken string) ([][]string, bool, error)
|
||||
}
|
||||
|
||||
// pendingAttachment is an embedded file block awaiting a download decision by
|
||||
// the connector (whitelist/size filtering happens there, not here).
|
||||
type pendingAttachment struct {
|
||||
FileToken string
|
||||
Name string
|
||||
}
|
||||
|
||||
// blocksToMarkdown renders a flat docx block array to Markdown, inlining
|
||||
// embedded spreadsheet/bitable tables (Task 5) and collecting downloadable
|
||||
// attachments. client may be nil when the block set has no downdrill blocks.
|
||||
func blocksToMarkdown(ctx context.Context, client sheetReader, blocks []DocxBlock) ([]byte, []pendingAttachment, error) {
|
||||
byID := make(map[string]DocxBlock, len(blocks))
|
||||
for _, b := range blocks {
|
||||
byID[b.BlockID] = b
|
||||
}
|
||||
// Blocks nested inside a native table (its cells and their content blocks)
|
||||
// are rendered by renderNativeTable. Mark them so the flat loop below does
|
||||
// not also Emit them as stray top-level paragraphs.
|
||||
consumed := tableDescendants(blocks, byID)
|
||||
|
||||
var sb strings.Builder
|
||||
var atts []pendingAttachment
|
||||
for _, b := range blocks {
|
||||
if consumed[b.BlockID] {
|
||||
continue
|
||||
}
|
||||
switch b.BlockType {
|
||||
case BlockTypePage:
|
||||
continue // root container, no text of its own
|
||||
case BlockTypeText:
|
||||
writePara(&sb, plainText(textBearingField(b)))
|
||||
case BlockTypeBullet:
|
||||
writePara(&sb, "- "+plainText(textBearingField(b)))
|
||||
case BlockTypeOrdered:
|
||||
writePara(&sb, "1. "+plainText(textBearingField(b)))
|
||||
case BlockTypeCode:
|
||||
writePara(&sb, "```\n"+plainText(textBearingField(b))+"\n```")
|
||||
case BlockTypeQuote:
|
||||
writePara(&sb, "> "+plainText(textBearingField(b)))
|
||||
case BlockTypeDivider:
|
||||
writePara(&sb, "---")
|
||||
case BlockTypeTable:
|
||||
writePara(&sb, renderNativeTable(b, byID))
|
||||
case BlockTypeTableCell:
|
||||
continue // rendered by its parent table (also covered by `consumed`)
|
||||
case BlockTypeSheet:
|
||||
if b.Sheet != nil {
|
||||
writePara(&sb, inlineTable(ctx, client, b.Sheet.Token, "sheet"))
|
||||
}
|
||||
case BlockTypeBitable:
|
||||
if b.Bitable != nil {
|
||||
writePara(&sb, inlineTable(ctx, client, b.Bitable.Token, "bitable"))
|
||||
}
|
||||
case BlockTypeImage:
|
||||
// Emit a token-free placeholder: images carry no retrievable text, and
|
||||
// leaking the internal media token would pollute embeddings. A neutral
|
||||
// marker preserves surrounding context (e.g. "如下图所示").
|
||||
writePara(&sb, "![图片]()")
|
||||
case BlockTypeTodo:
|
||||
if t := plainText(textBearingField(b)); t != "" {
|
||||
writePara(&sb, "- [ ] "+t)
|
||||
}
|
||||
case BlockTypeCallout:
|
||||
// Callout is a container; its body is usually in child blocks (rendered
|
||||
// separately). Emit its own text only if it carries direct inline text,
|
||||
// so the container case is a safe no-op.
|
||||
if t := plainText(textBearingField(b)); t != "" {
|
||||
writePara(&sb, "> "+t)
|
||||
}
|
||||
case BlockTypeFile:
|
||||
if b.File != nil {
|
||||
name := b.File.Name
|
||||
if name == "" {
|
||||
name = b.File.Token
|
||||
}
|
||||
writePara(&sb, "📎 附件:"+name)
|
||||
atts = append(atts, pendingAttachment{FileToken: b.File.Token, Name: b.File.Name})
|
||||
}
|
||||
default:
|
||||
if b.BlockType >= BlockTypeHeading1 && b.BlockType <= blockTypeHeading9 {
|
||||
level := b.BlockType - BlockTypeHeading1 + 1
|
||||
writePara(&sb, strings.Repeat("#", level)+" "+plainText(textBearingField(b)))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
out := strings.TrimRight(sb.String(), "\n")
|
||||
if out != "" {
|
||||
out += "\n"
|
||||
}
|
||||
return []byte(out), atts, nil
|
||||
}
|
||||
|
||||
// writePara appends a block of text followed by a blank line; empty text is skipped.
|
||||
func writePara(sb *strings.Builder, s string) {
|
||||
if s == "" {
|
||||
return
|
||||
}
|
||||
sb.WriteString(s)
|
||||
sb.WriteString("\n\n")
|
||||
}
|
||||
|
||||
// plainText concatenates the text runs of a text-bearing block.
|
||||
func plainText(bt *BlockText) string {
|
||||
if bt == nil {
|
||||
return ""
|
||||
}
|
||||
var sb strings.Builder
|
||||
for _, e := range bt.Elements {
|
||||
if e.TextRun != nil {
|
||||
sb.WriteString(e.TextRun.Content)
|
||||
}
|
||||
}
|
||||
return sb.String()
|
||||
}
|
||||
|
||||
// headingText returns the heading field for the block's level, or Text as fallback.
|
||||
func headingText(b DocxBlock) *BlockText {
|
||||
fields := []*BlockText{
|
||||
b.Heading1, b.Heading2, b.Heading3, b.Heading4, b.Heading5,
|
||||
b.Heading6, b.Heading7, b.Heading8, b.Heading9,
|
||||
}
|
||||
if idx := b.BlockType - BlockTypeHeading1; idx >= 0 && idx < len(fields) && fields[idx] != nil {
|
||||
return fields[idx]
|
||||
}
|
||||
return b.Text
|
||||
}
|
||||
|
||||
// tableDescendants returns the set of block IDs that belong to a native table —
|
||||
// every cell listed in a table block plus everything reachable through those
|
||||
// cells' Children. blocksToMarkdown skips these so table content is emitted only
|
||||
// by the table renderer, never a second time as loose paragraphs.
|
||||
func tableDescendants(blocks []DocxBlock, byID map[string]DocxBlock) map[string]bool {
|
||||
consumed := make(map[string]bool)
|
||||
var mark func(id string)
|
||||
mark = func(id string) {
|
||||
b, ok := byID[id]
|
||||
if !ok || consumed[id] {
|
||||
return
|
||||
}
|
||||
// Attachment/media blocks nested in a cell must still be collected (and
|
||||
// their reference emitted) by the main loop — the table renderer only
|
||||
// extracts text — so do not consume them, only their text structure.
|
||||
if b.BlockType == BlockTypeFile || b.BlockType == BlockTypeImage {
|
||||
return
|
||||
}
|
||||
consumed[id] = true
|
||||
for _, c := range b.Children {
|
||||
mark(c)
|
||||
}
|
||||
}
|
||||
for _, b := range blocks {
|
||||
// Only consume cells of tables we will actually render. A table that
|
||||
// renderNativeTable would bail on (missing property / zero columns) must
|
||||
// NOT have its cells consumed, or their text would be dropped entirely —
|
||||
// leave them for the flat loop to Emit as loose paragraphs instead.
|
||||
if b.BlockType == BlockTypeTable && tableRenderable(b) {
|
||||
for _, cid := range b.Table.Cells {
|
||||
mark(cid)
|
||||
}
|
||||
}
|
||||
}
|
||||
return consumed
|
||||
}
|
||||
|
||||
// tableRenderable reports whether a native table block carries enough structure
|
||||
// (a column count) for renderNativeTable to produce a Markdown table. It is the
|
||||
// single predicate shared by the consume pass and the render pass so the two
|
||||
// never disagree about which tables are handled by the table renderer.
|
||||
func tableRenderable(b DocxBlock) bool {
|
||||
return b.Table != nil && b.Table.Property != nil && b.Table.Property.ColumnSize > 0
|
||||
}
|
||||
|
||||
// textBearingField returns the inline-text payload for whichever type-named
|
||||
// field a block populates (docx stores a block's text in a field named after
|
||||
// its type), so cell content of any text-like type can be extracted uniformly.
|
||||
func textBearingField(b DocxBlock) *BlockText {
|
||||
switch b.BlockType {
|
||||
case BlockTypeText:
|
||||
return b.Text
|
||||
case BlockTypeBullet:
|
||||
return b.Bullet
|
||||
case BlockTypeOrdered:
|
||||
return b.Ordered
|
||||
case BlockTypeCode:
|
||||
return b.Code
|
||||
case BlockTypeQuote:
|
||||
return b.Quote
|
||||
case BlockTypeTodo:
|
||||
return b.Todo
|
||||
case BlockTypeCallout:
|
||||
return b.Callout
|
||||
}
|
||||
if b.BlockType >= BlockTypeHeading1 && b.BlockType <= blockTypeHeading9 {
|
||||
return headingText(b)
|
||||
}
|
||||
return b.Text
|
||||
}
|
||||
|
||||
// cellText renders a native table cell to a single string. A Feishu table_cell
|
||||
// (block_type 32) is a container: its text lives in child blocks, not on the
|
||||
// cell itself, so we concatenate the text of each child block.
|
||||
func cellText(cell DocxBlock, byID map[string]DocxBlock) string {
|
||||
var parts []string
|
||||
for _, childID := range cell.Children {
|
||||
if t := plainText(textBearingField(byID[childID])); t != "" {
|
||||
parts = append(parts, t)
|
||||
}
|
||||
}
|
||||
return strings.Join(parts, " ")
|
||||
}
|
||||
|
||||
// renderNativeTable renders a native docx table block into a Markdown table.
|
||||
func renderNativeTable(b DocxBlock, byID map[string]DocxBlock) string {
|
||||
if !tableRenderable(b) {
|
||||
return ""
|
||||
}
|
||||
cols := b.Table.Property.ColumnSize
|
||||
var cells []string
|
||||
for _, cid := range b.Table.Cells {
|
||||
cells = append(cells, cellText(byID[cid], byID))
|
||||
}
|
||||
var rows [][]string
|
||||
for i := 0; i < len(cells); i += cols {
|
||||
end := i + cols
|
||||
if end > len(cells) {
|
||||
end = len(cells)
|
||||
}
|
||||
rows = append(rows, cells[i:end])
|
||||
}
|
||||
return markdownTable(rows)
|
||||
}
|
||||
|
||||
// markdownTable renders a [][]string (first row = header) as a GFM table.
|
||||
func markdownTable(rows [][]string) string {
|
||||
// A zero-column header (an embedded sheet/bitable with no columns) would Emit
|
||||
// a header line of "| |" and a separator of just "|" — malformed GFM. Render
|
||||
// nothing instead.
|
||||
if len(rows) == 0 || len(rows[0]) == 0 {
|
||||
return ""
|
||||
}
|
||||
var sb strings.Builder
|
||||
cols := len(rows[0])
|
||||
sb.WriteString("| " + strings.Join(escapePipes(rows[0]), " | ") + " |\n")
|
||||
sb.WriteString("|" + strings.Repeat(" --- |", cols) + "\n")
|
||||
for _, r := range rows[1:] {
|
||||
// Ragged data: a row wider than the header would Emit more cells than the
|
||||
// header/separator declare, producing a malformed GFM table. Clamp to the
|
||||
// header width (truncate overflow, pad shortfall).
|
||||
if len(r) > cols {
|
||||
r = r[:cols]
|
||||
}
|
||||
for len(r) < cols {
|
||||
r = append(r, "")
|
||||
}
|
||||
sb.WriteString("| " + strings.Join(escapePipes(r), " | ") + " |\n")
|
||||
}
|
||||
return strings.TrimRight(sb.String(), "\n")
|
||||
}
|
||||
|
||||
// escapePipes makes cell values safe for a single Markdown table row.
|
||||
func escapePipes(row []string) []string {
|
||||
out := make([]string, len(row))
|
||||
for i, c := range row {
|
||||
out[i] = strings.ReplaceAll(strings.ReplaceAll(c, "\n", " "), "|", "\\|")
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// inlineTable reads an embedded sheet/bitable and renders it as a Markdown
|
||||
// table. Read/permission errors degrade to an inline note rather than failing
|
||||
// the whole document (partial content beats no content). A reader-reported
|
||||
// truncation appends a note.
|
||||
func inlineTable(ctx context.Context, client sheetReader, token, kind string) string {
|
||||
if client == nil {
|
||||
return ""
|
||||
}
|
||||
var (
|
||||
rows [][]string
|
||||
truncated bool
|
||||
err error
|
||||
noun string
|
||||
)
|
||||
if kind == "sheet" {
|
||||
rows, truncated, err = client.readSheetRange(ctx, token)
|
||||
noun = "内嵌电子表格"
|
||||
} else {
|
||||
rows, truncated, err = client.readBitableRecords(ctx, token)
|
||||
noun = "内嵌多维表格"
|
||||
}
|
||||
if err != nil {
|
||||
return fmt.Sprintf("> [无法读取%s]", noun)
|
||||
}
|
||||
// markdownTable returns "" when there is nothing renderable (no rows, or a
|
||||
// header with no columns). Skip the truncation note too in that case, since it
|
||||
// would otherwise dangle without a table above it.
|
||||
table := markdownTable(rows)
|
||||
if table == "" {
|
||||
return ""
|
||||
}
|
||||
if truncated {
|
||||
table += fmt.Sprintf("\n\n> 表格已截断(仅显示前 %d 行)", maxTableRows)
|
||||
}
|
||||
return table
|
||||
}
|
||||
@@ -0,0 +1,295 @@
|
||||
package core
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// txt builds a text-bearing BlockText from plain content.
|
||||
func txt(s string) *BlockText {
|
||||
return &BlockText{Elements: []TextElement{{TextRun: &TextRun{Content: s}}}}
|
||||
}
|
||||
|
||||
type fakeReader struct {
|
||||
sheet [][]string
|
||||
sheetTruncated bool
|
||||
sheetErr error
|
||||
bitable [][]string
|
||||
bitableTruncated bool
|
||||
bitableErr error
|
||||
}
|
||||
|
||||
func (f fakeReader) readSheetRange(_ context.Context, _ string) ([][]string, bool, error) {
|
||||
return f.sheet, f.sheetTruncated, f.sheetErr
|
||||
}
|
||||
|
||||
func (f fakeReader) readBitableRecords(_ context.Context, _ string) ([][]string, bool, error) {
|
||||
return f.bitable, f.bitableTruncated, f.bitableErr
|
||||
}
|
||||
|
||||
func TestBlocksToMarkdown_EmbeddedSheetAndFile(t *testing.T) {
|
||||
blocks := []DocxBlock{
|
||||
{BlockID: "root", BlockType: BlockTypePage},
|
||||
{BlockID: "s", BlockType: BlockTypeSheet, Sheet: &BlockTokenRef{Token: "sht_a_0"}},
|
||||
{BlockID: "img", BlockType: BlockTypeImage, Image: &BlockTokenRef{Token: "img_t"}},
|
||||
{BlockID: "f", BlockType: BlockTypeFile, File: &BlockFileRef{Token: "file_t", Name: "报表.pdf"}},
|
||||
}
|
||||
fr := fakeReader{sheet: [][]string{{"名称", "数量"}, {"苹果", "3"}}}
|
||||
md, atts, err := blocksToMarkdown(context.Background(), fr, blocks)
|
||||
if err != nil {
|
||||
t.Fatalf("err: %v", err)
|
||||
}
|
||||
if !strings.Contains(string(md), "| 名称 | 数量 |") || !strings.Contains(string(md), "| 苹果 | 3 |") {
|
||||
t.Errorf("sheet not inlined:\n%s", md)
|
||||
}
|
||||
if !strings.Contains(string(md), "![图片]()") {
|
||||
t.Errorf("image placeholder missing:\n%s", md)
|
||||
}
|
||||
if strings.Contains(string(md), "feishu-media") {
|
||||
t.Errorf("internal media token leaked into markdown:\n%s", md)
|
||||
}
|
||||
if len(atts) != 1 || atts[0].FileToken != "file_t" || atts[0].Name != "报表.pdf" {
|
||||
t.Errorf("attachments = %+v", atts)
|
||||
}
|
||||
if !strings.Contains(string(md), "📎 附件:报表.pdf") {
|
||||
t.Errorf("attachment inline reference missing:\n%s", md)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBlocksToMarkdown_SheetTruncatedNote(t *testing.T) {
|
||||
blocks := []DocxBlock{
|
||||
{BlockID: "root", BlockType: BlockTypePage},
|
||||
{BlockID: "s", BlockType: BlockTypeSheet, Sheet: &BlockTokenRef{Token: "sht_a_0"}},
|
||||
}
|
||||
fr := fakeReader{sheet: [][]string{{"h"}, {"1"}}, sheetTruncated: true}
|
||||
md, _, err := blocksToMarkdown(context.Background(), fr, blocks)
|
||||
if err != nil {
|
||||
t.Fatalf("err: %v", err)
|
||||
}
|
||||
if !strings.Contains(string(md), "表格已截断") {
|
||||
t.Errorf("want truncation note, got:\n%s", md)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBlocksToMarkdown_SheetPermissionDegrades(t *testing.T) {
|
||||
blocks := []DocxBlock{
|
||||
{BlockID: "root", BlockType: BlockTypePage},
|
||||
{BlockID: "s", BlockType: BlockTypeSheet, Sheet: &BlockTokenRef{Token: "sht_a_0"}},
|
||||
}
|
||||
fr := fakeReader{sheetErr: fmt.Errorf("code=99991672 permission denied")}
|
||||
md, _, err := blocksToMarkdown(context.Background(), fr, blocks)
|
||||
if err != nil {
|
||||
t.Fatalf("should not fail on permission error, got %v", err)
|
||||
}
|
||||
if !strings.Contains(string(md), "无法读取内嵌电子表格") {
|
||||
t.Errorf("want degraded placeholder, got:\n%s", md)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBlocksToMarkdown_BitableInlinedAndDegrades(t *testing.T) {
|
||||
mk := func(fr fakeReader) string {
|
||||
blocks := []DocxBlock{
|
||||
{BlockID: "root", BlockType: BlockTypePage},
|
||||
{BlockID: "bt", BlockType: BlockTypeBitable, Bitable: &BlockTokenRef{Token: "bascabc_tblxyz"}},
|
||||
}
|
||||
md, _, err := blocksToMarkdown(context.Background(), fr, blocks)
|
||||
if err != nil {
|
||||
t.Fatalf("err: %v", err)
|
||||
}
|
||||
return string(md)
|
||||
}
|
||||
|
||||
inlined := mk(fakeReader{bitable: [][]string{{"任务", "状态"}, {"写文档", "进行中"}}})
|
||||
if !strings.Contains(inlined, "| 任务 | 状态 |") || !strings.Contains(inlined, "| 写文档 | 进行中 |") {
|
||||
t.Errorf("bitable not inlined:\n%s", inlined)
|
||||
}
|
||||
|
||||
degraded := mk(fakeReader{bitableErr: fmt.Errorf("code=99991672 permission denied")})
|
||||
if !strings.Contains(degraded, "无法读取内嵌多维表格") {
|
||||
t.Errorf("want bitable degradation note, got:\n%s", degraded)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBlocksToMarkdown_NativeTableFromCellChildren(t *testing.T) {
|
||||
// Real Feishu shape: a table_cell (block_type 32) is a container whose text
|
||||
// lives in child text blocks, not on the cell. The renderer must read cell
|
||||
// text from those children AND must not also Emit them as loose paragraphs.
|
||||
blocks := []DocxBlock{
|
||||
{BlockID: "root", BlockType: BlockTypePage},
|
||||
tableBlk("t", 2, "c1", "c2", "c3", "c4"),
|
||||
cellBlk("c1"), cellBlk("c2"), cellBlk("c3"), cellBlk("c4"),
|
||||
cellTextBlk("c1", "姓名"), cellTextBlk("c2", "分数"),
|
||||
cellTextBlk("c3", "张三"), cellTextBlk("c4", "95"),
|
||||
}
|
||||
md, _, err := blocksToMarkdown(context.Background(), nil, blocks)
|
||||
if err != nil {
|
||||
t.Fatalf("err: %v", err)
|
||||
}
|
||||
s := string(md)
|
||||
if !strings.Contains(s, "| 姓名 | 分数 |") || !strings.Contains(s, "| 张三 | 95 |") {
|
||||
t.Errorf("native table not rendered from cell children:\n%s", s)
|
||||
}
|
||||
// Cell text must appear ONLY inside the table, never as a stray paragraph.
|
||||
if n := strings.Count(s, "姓名"); n != 1 {
|
||||
t.Errorf("cell text leaked outside the table (count=%d, want 1):\n%s", n, s)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBlocksToMarkdown_AttachmentInsideTableCellStillCollected(t *testing.T) {
|
||||
// A file block nested inside a table cell must still be collected as an
|
||||
// attachment — the table "consumed" set marks a cell's text structure but
|
||||
// must NOT swallow attachment/media blocks, or embedded files silently vanish.
|
||||
blocks := []DocxBlock{
|
||||
{BlockID: "root", BlockType: BlockTypePage},
|
||||
tableBlk("t", 1, "c1"),
|
||||
{BlockID: "c1", BlockType: BlockTypeTableCell, Children: []string{"f1"}},
|
||||
{BlockID: "f1", BlockType: BlockTypeFile, File: &BlockFileRef{Token: "tok-in-cell", Name: "内嵌.pdf"}},
|
||||
}
|
||||
_, atts, err := blocksToMarkdown(context.Background(), nil, blocks)
|
||||
if err != nil {
|
||||
t.Fatalf("err: %v", err)
|
||||
}
|
||||
if len(atts) != 1 || atts[0].FileToken != "tok-in-cell" {
|
||||
t.Fatalf("attachment nested in a table cell was not collected: %+v", atts)
|
||||
}
|
||||
}
|
||||
|
||||
func TestMarkdownTable_RaggedRowClampedToHeader(t *testing.T) {
|
||||
// Header declares 2 columns; a data row with 4 cells (trailing populated
|
||||
// cells beyond the header range) must be clamped to 2, or the emitted table
|
||||
// has more body cells than header/separator columns → malformed GFM.
|
||||
rows := [][]string{
|
||||
{"名称", "数量"},
|
||||
{"苹果", "3", "多余1", "多余2"},
|
||||
}
|
||||
out := markdownTable(rows)
|
||||
lines := strings.Split(out, "\n")
|
||||
if len(lines) != 3 {
|
||||
t.Fatalf("want 3 lines (header, separator, 1 row), got %d:\n%s", len(lines), out)
|
||||
}
|
||||
// Every row must have the same pipe count as the 2-column header.
|
||||
want := strings.Count(lines[0], "|")
|
||||
for i, ln := range lines {
|
||||
if got := strings.Count(ln, "|"); got != want {
|
||||
t.Errorf("line %d %q has %d pipes, want %d (ragged row not clamped)", i, ln, got, want)
|
||||
}
|
||||
}
|
||||
if strings.Contains(out, "多余") {
|
||||
t.Errorf("overflow cells must be dropped, got:\n%s", out)
|
||||
}
|
||||
}
|
||||
|
||||
func TestMarkdownTable_ZeroColumnRendersNothing(t *testing.T) {
|
||||
// An embedded sheet/bitable with a header row but no columns (e.g. a bitable
|
||||
// whose fields were all deleted) yields rows == [][]string{{}}. len(rows)==1
|
||||
// passes a naive empty guard, but cols==0 would Emit a malformed GFM table
|
||||
// ("| |" header + a bare "|" separator). It must render nothing instead.
|
||||
for _, rows := range [][][]string{
|
||||
{{}}, // one empty header row, no data
|
||||
{{}, {}, {}}, // empty header + empty data rows
|
||||
} {
|
||||
out := markdownTable(rows)
|
||||
if out != "" {
|
||||
t.Errorf("zero-column table must render nothing, got %q for rows=%+v", out, rows)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestBlocksToMarkdown_UnrenderableTablePreservesCellText(t *testing.T) {
|
||||
// A native table block with no column property (Table!=nil, Property==nil)
|
||||
// cannot be rendered as a Markdown table. Its cells must NOT be consumed, so
|
||||
// their text still reaches the output as loose paragraphs rather than
|
||||
// vanishing entirely.
|
||||
blocks := []DocxBlock{
|
||||
{BlockID: "root", BlockType: BlockTypePage},
|
||||
// tableBlk sets Property (renderable); here we want an UNrenderable one.
|
||||
{BlockID: "t", BlockType: BlockTypeTable, Table: &BlockTable{Cells: []string{"c1"}}},
|
||||
{BlockID: "c1", BlockType: BlockTypeTableCell, Children: []string{"c1_txt"}},
|
||||
cellTextBlk("c1", "重要内容"),
|
||||
}
|
||||
md, _, err := blocksToMarkdown(context.Background(), nil, blocks)
|
||||
if err != nil {
|
||||
t.Fatalf("err: %v", err)
|
||||
}
|
||||
if !strings.Contains(string(md), "重要内容") {
|
||||
t.Errorf("cell text of an unrenderable table was dropped:\n%s", md)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBlocksToMarkdown_TextConstructs(t *testing.T) {
|
||||
blocks := []DocxBlock{
|
||||
{BlockID: "root", BlockType: BlockTypePage, Children: []string{"h", "p", "b"}},
|
||||
{BlockID: "h", BlockType: BlockTypeHeading1, Heading1: txt("标题")},
|
||||
{BlockID: "p", BlockType: BlockTypeText, Text: txt("一段正文")},
|
||||
{BlockID: "b", BlockType: BlockTypeBullet, Bullet: txt("要点")},
|
||||
}
|
||||
md, atts, err := blocksToMarkdown(context.Background(), nil, blocks)
|
||||
if err != nil {
|
||||
t.Fatalf("blocksToMarkdown: %v", err)
|
||||
}
|
||||
if len(atts) != 0 {
|
||||
t.Errorf("want 0 attachments, got %d", len(atts))
|
||||
}
|
||||
want := "# 标题\n\n一段正文\n\n- 要点\n"
|
||||
if string(md) != want {
|
||||
t.Errorf("md =\n%q\nwant\n%q", md, want)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBlocksToMarkdown_BlankDocRendersEmpty(t *testing.T) {
|
||||
// A page with no renderable children (or only empty-text/container blocks)
|
||||
// must render to empty Markdown. FetchDocxWithBlocks relies on this: an empty
|
||||
// render triggers the export fallback instead of emitting a content-less main
|
||||
// item that would wrongly ingest the login-gated wiki URL. Any File/Image
|
||||
// block writes a placeholder, so a truly empty render also implies no
|
||||
// attachment/image downdrill is lost by that fallback.
|
||||
blocks := []DocxBlock{
|
||||
{BlockID: "root", BlockType: BlockTypePage, Children: []string{"p"}},
|
||||
{BlockID: "p", BlockType: BlockTypeText, Text: txt("")},
|
||||
}
|
||||
md, atts, err := blocksToMarkdown(context.Background(), nil, blocks)
|
||||
if err != nil {
|
||||
t.Fatalf("blocksToMarkdown: %v", err)
|
||||
}
|
||||
if strings.TrimSpace(string(md)) != "" {
|
||||
t.Errorf("blank doc should render empty Markdown, got:\n%q", md)
|
||||
}
|
||||
if len(atts) != 0 {
|
||||
t.Errorf("blank doc should collect no attachments, got %d", len(atts))
|
||||
}
|
||||
}
|
||||
|
||||
func TestBlocksToMarkdown_TodoAndCallout(t *testing.T) {
|
||||
blocks := []DocxBlock{
|
||||
{BlockID: "root", BlockType: BlockTypePage},
|
||||
{BlockID: "t", BlockType: BlockTypeTodo, Todo: txt("买牛奶")},
|
||||
{BlockID: "c", BlockType: BlockTypeCallout, Callout: txt("注意事项")},
|
||||
}
|
||||
md, _, err := blocksToMarkdown(context.Background(), nil, blocks)
|
||||
if err != nil {
|
||||
t.Fatalf("err: %v", err)
|
||||
}
|
||||
if !strings.Contains(string(md), "- [ ] 买牛奶") {
|
||||
t.Errorf("todo not rendered:\n%s", md)
|
||||
}
|
||||
if !strings.Contains(string(md), "> 注意事项") {
|
||||
t.Errorf("callout direct text not rendered:\n%s", md)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBlocksToMarkdown_CalloutContainerNoOp(t *testing.T) {
|
||||
// A callout with no direct text (container form) must Emit nothing itself.
|
||||
blocks := []DocxBlock{
|
||||
{BlockID: "root", BlockType: BlockTypePage},
|
||||
{BlockID: "c", BlockType: BlockTypeCallout},
|
||||
}
|
||||
md, _, err := blocksToMarkdown(context.Background(), nil, blocks)
|
||||
if err != nil {
|
||||
t.Fatalf("err: %v", err)
|
||||
}
|
||||
if strings.TrimSpace(string(md)) != "" {
|
||||
t.Errorf("empty callout should Emit nothing, got:\n%q", md)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,81 @@
|
||||
package core
|
||||
|
||||
import "github.com/Tencent/WeKnora/internal/types"
|
||||
|
||||
// Open Platform API origins. Feishu and Lark are the same product deployed on
|
||||
// two isolated clouds; the wiki/docx/drive APIs this connector uses are
|
||||
// identical, only the host differs.
|
||||
const (
|
||||
feishuOpenBaseURL = "https://open.feishu.cn"
|
||||
larkOpenBaseURL = "https://open.larksuite.com"
|
||||
)
|
||||
|
||||
// Web origins used to build human-facing links to wiki spaces and nodes. These
|
||||
// are the end-user app hosts, not the API hosts.
|
||||
const (
|
||||
feishuWebBaseURL = "https://feishu.cn"
|
||||
larkWebBaseURL = "https://larksuite.com"
|
||||
)
|
||||
|
||||
// Region selects which cloud the connector syncs from. Apps, tenants, tokens
|
||||
// and document tokens are scoped to a single cloud and are never valid on the
|
||||
// other, so a Region is fixed per connector instance at registration.
|
||||
type Region struct {
|
||||
// ConnectorType is the types.ConnectorType* identifier this instance serves.
|
||||
ConnectorType string
|
||||
// OpenBaseURL is the Open Platform API origin, without a trailing slash.
|
||||
// It is only the default: a data source may override it via base_url.
|
||||
OpenBaseURL string
|
||||
// WebBaseURL is the app origin used to build wiki links shown to users.
|
||||
WebBaseURL string
|
||||
// Label prefixes log lines so operators can tell the clouds apart.
|
||||
Label string
|
||||
}
|
||||
|
||||
var (
|
||||
// RegionFeishu is the Chinese mainland cloud (飞书).
|
||||
RegionFeishu = Region{
|
||||
ConnectorType: types.ConnectorTypeFeishu,
|
||||
OpenBaseURL: feishuOpenBaseURL,
|
||||
WebBaseURL: feishuWebBaseURL,
|
||||
Label: "Feishu",
|
||||
}
|
||||
|
||||
// RegionLark is the international cloud (Lark).
|
||||
RegionLark = Region{
|
||||
ConnectorType: types.ConnectorTypeLark,
|
||||
OpenBaseURL: larkOpenBaseURL,
|
||||
WebBaseURL: larkWebBaseURL,
|
||||
Label: "Lark",
|
||||
}
|
||||
|
||||
// RegionFeishuDrive is the Chinese mainland cloud, Drive (云盘) mode.
|
||||
// Shares the feishu connector package with RegionFeishu; only the connector
|
||||
// type differs so the registry dispatches to the Drive connector.
|
||||
RegionFeishuDrive = Region{
|
||||
ConnectorType: types.ConnectorTypeFeishuDrive,
|
||||
OpenBaseURL: feishuOpenBaseURL,
|
||||
WebBaseURL: feishuWebBaseURL,
|
||||
Label: "FeishuDrive",
|
||||
}
|
||||
|
||||
// RegionLarkDrive is the international cloud, Drive mode.
|
||||
RegionLarkDrive = Region{
|
||||
ConnectorType: types.ConnectorTypeLarkDrive,
|
||||
OpenBaseURL: larkOpenBaseURL,
|
||||
WebBaseURL: larkWebBaseURL,
|
||||
Label: "LarkDrive",
|
||||
}
|
||||
)
|
||||
|
||||
// WikiURL builds the user-facing link to a wiki space or node on this cloud.
|
||||
// The token is either a space_id or a node_token; both live under /wiki/.
|
||||
func (r Region) WikiURL(token string) string {
|
||||
return r.WebBaseURL + "/wiki/" + token
|
||||
}
|
||||
|
||||
// DriveFolderURL builds the user-facing link to a Drive folder on this cloud.
|
||||
// A folder_token lives under /drive/folder/.
|
||||
func (r Region) DriveFolderURL(folderToken string) string {
|
||||
return r.WebBaseURL + "/drive/folder/" + folderToken
|
||||
}
|
||||
@@ -0,0 +1,76 @@
|
||||
package core
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
)
|
||||
|
||||
func TestRegions_AreDistinct(t *testing.T) {
|
||||
if RegionFeishu.ConnectorType == RegionLark.ConnectorType {
|
||||
t.Errorf("regions share ConnectorType %q", RegionFeishu.ConnectorType)
|
||||
}
|
||||
if RegionFeishu.OpenBaseURL == RegionLark.OpenBaseURL {
|
||||
t.Errorf("regions share OpenBaseURL %q", RegionFeishu.OpenBaseURL)
|
||||
}
|
||||
if RegionFeishu.WebBaseURL == RegionLark.WebBaseURL {
|
||||
t.Errorf("regions share WebBaseURL %q", RegionFeishu.WebBaseURL)
|
||||
}
|
||||
if RegionFeishu.ConnectorType != types.ConnectorTypeFeishu {
|
||||
t.Errorf("RegionFeishu.ConnectorType = %q, want %q", RegionFeishu.ConnectorType, types.ConnectorTypeFeishu)
|
||||
}
|
||||
if RegionLark.ConnectorType != types.ConnectorTypeLark {
|
||||
t.Errorf("RegionLark.ConnectorType = %q, want %q", RegionLark.ConnectorType, types.ConnectorTypeLark)
|
||||
}
|
||||
}
|
||||
|
||||
// TestConnectorType_FollowsRegion lives in package wiki (it exercises
|
||||
// wiki.NewConnector); see wiki/connector_test.go TestConnectorType.
|
||||
|
||||
// The wiki link shown in the resource picker must point at the cloud the data
|
||||
// actually lives on. Feishu's value is the pre-existing one and must not drift.
|
||||
func TestRegion_WikiURL(t *testing.T) {
|
||||
if got, want := RegionFeishu.WikiURL("spc123"), "https://feishu.cn/wiki/spc123"; got != want {
|
||||
t.Errorf("Feishu WikiURL = %q, want %q", got, want)
|
||||
}
|
||||
if got, want := RegionLark.WikiURL("spc123"), "https://larksuite.com/wiki/spc123"; got != want {
|
||||
t.Errorf("Lark WikiURL = %q, want %q", got, want)
|
||||
}
|
||||
// A Lark link must never point at the Feishu host — that was the bug.
|
||||
if strings.Contains(RegionLark.WikiURL("x"), "feishu") {
|
||||
t.Errorf("Lark WikiURL leaks the Feishu host: %q", RegionLark.WikiURL("x"))
|
||||
}
|
||||
}
|
||||
|
||||
// base_url stays an override so data sources created before the lark connector
|
||||
// existed — a "feishu" connector pointed at open.larksuite.com — keep working.
|
||||
func TestParseFeishuConfig_BaseURLDefaultsToRegion(t *testing.T) {
|
||||
cases := []struct {
|
||||
name string
|
||||
region Region
|
||||
baseURL string
|
||||
want string
|
||||
}{
|
||||
{"feishu default", RegionFeishu, "", "https://open.feishu.cn"},
|
||||
{"lark default", RegionLark, "", "https://open.larksuite.com"},
|
||||
{"explicit override wins on feishu", RegionFeishu, "https://open.larksuite.com", "https://open.larksuite.com"},
|
||||
{"explicit override wins on lark", RegionLark, "https://open.feishu.cn", "https://open.feishu.cn"},
|
||||
}
|
||||
|
||||
for _, c := range cases {
|
||||
t.Run(c.name, func(t *testing.T) {
|
||||
creds := map[string]interface{}{"app_id": "cli_x", "app_secret": "s"}
|
||||
if c.baseURL != "" {
|
||||
creds["base_url"] = c.baseURL
|
||||
}
|
||||
cfg, err := ParseFeishuConfig(&types.DataSourceConfig{Credentials: creds}, c.region)
|
||||
if err != nil {
|
||||
t.Fatalf("ParseFeishuConfig: %v", err)
|
||||
}
|
||||
if got := cfg.GetBaseURL(); got != c.want {
|
||||
t.Errorf("GetBaseURL() = %q, want %q", got, c.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,502 @@
|
||||
package core
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"maps"
|
||||
"net/http"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"regexp"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
"unicode/utf8"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/datasource"
|
||||
"github.com/Tencent/WeKnora/internal/logger"
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
)
|
||||
|
||||
const FeishuWikiNodeResourceSeparator = ":"
|
||||
|
||||
// shared.go holds the helpers used by BOTH the wiki Connector (connector.go)
|
||||
// and the Drive DriveConnector (connector.go): error classification,
|
||||
// config parsing, stream-Checkpoint tuning, the fetch tally, filename/time
|
||||
// utilities, attachment rules, and the docx blocks fetch path. Anything that
|
||||
// is specific to one connector stays in that connector's own file.
|
||||
|
||||
// FeishuStreamCheckpointInterval is how many processed nodes pass between
|
||||
// cursor checkpoints during a streaming fetch. Small enough that a timed-out
|
||||
// sync loses little work on resume, large enough that Checkpoint persistence
|
||||
// (a DB write) does not dominate. Overridable in tests. See FetchStream.
|
||||
var FeishuStreamCheckpointInterval = 50
|
||||
|
||||
// FeishuStreamCheckpointMaxInterval bounds checkpointing by wall-clock time as
|
||||
// well as node count. Without it, a sync of fewer than
|
||||
// FeishuStreamCheckpointInterval very slow (rate-limited) exports could reach
|
||||
// the 2h task timeout having never checkpointed, and resume from scratch every
|
||||
// retry — the #2136 "never fully syncs" case. Overridable in tests.
|
||||
var FeishuStreamCheckpointMaxInterval = 30 * time.Second
|
||||
|
||||
// fetchTally accumulates the outcome of fetching a wiki node subtree so the
|
||||
// connector can Emit a single actionable summary. Without it, unsupported nodes
|
||||
// (mindnote/slides/etc.) vanish with no item, no error and no log, leaving users
|
||||
// unable to explain why "13 documents synced only 3" (Tencent/WeKnora#2136).
|
||||
type fetchTally struct {
|
||||
discovered int
|
||||
fetched int
|
||||
failed int
|
||||
skippedByType map[string]int
|
||||
}
|
||||
|
||||
func newFetchTally(discovered int) *fetchTally {
|
||||
return &fetchTally{discovered: discovered, skippedByType: map[string]int{}}
|
||||
}
|
||||
|
||||
func (t *fetchTally) fetch() { t.fetched++ }
|
||||
func (t *fetchTally) fail() { t.failed++ }
|
||||
func (t *fetchTally) Skip(objType string) { t.skippedByType[objType]++ }
|
||||
|
||||
func (t *fetchTally) skipped() int {
|
||||
n := 0
|
||||
for _, c := range t.skippedByType {
|
||||
n += c
|
||||
}
|
||||
return n
|
||||
}
|
||||
|
||||
func (t *fetchTally) summary() string {
|
||||
return fmt.Sprintf("discovered=%d fetched=%d failed=%d skipped_unsupported=%d by_type=%v",
|
||||
t.discovered, t.fetched, t.failed, t.skipped(), t.skippedByType)
|
||||
}
|
||||
|
||||
var reFeishuErrorCode = regexp.MustCompile(`code["\s]*[:=]\s*(\d+)`)
|
||||
|
||||
// feishuErrorCode extracts the numeric Feishu error code from a raw error string
|
||||
// (e.g. `body={"code":1663,...}` or `code=1663`), best-effort.
|
||||
func feishuErrorCode(raw string) string {
|
||||
if m := reFeishuErrorCode.FindStringSubmatch(raw); len(m) == 2 {
|
||||
return m[1]
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
// feishuFailure classifies a raw connector/API error into a stable i18n code
|
||||
// (mapped to a localized string on the frontend), an optional numeric Feishu
|
||||
// error code for interpolation, and an English fallback message for clients
|
||||
// without the i18n key. The raw status/JSON body/log_id is never returned here —
|
||||
// it stays in the server logs. Dumping it in the UI is the anti-pattern
|
||||
// Airbyte/Fivetran/Onyx warn against. Transient errors are retried next sync
|
||||
// (the cursor is retained); auth/permission errors point at the fix instead.
|
||||
func feishuFailure(err error) (code, codeValue, fallback string) {
|
||||
if err == nil {
|
||||
return "sync_failed", "", "Sync failed; will retry on the next sync"
|
||||
}
|
||||
s := strings.ToLower(err.Error())
|
||||
|
||||
switch {
|
||||
case strings.Contains(s, "auth error"),
|
||||
strings.Contains(s, "invalid access token"),
|
||||
strings.Contains(s, "permission"),
|
||||
strings.Contains(s, "forbidden"),
|
||||
strings.Contains(s, "status=403"):
|
||||
return "feishu_auth_or_permission", "", "Authentication or permission error; check credentials and app scopes"
|
||||
case strings.Contains(s, "rate limited"), strings.Contains(s, "status=429"):
|
||||
return "feishu_rate_limited", "", "Feishu API rate limited; will retry on the next sync"
|
||||
case strings.Contains(s, "timed out"),
|
||||
strings.Contains(s, "timeout"),
|
||||
strings.Contains(s, "deadline exceeded"):
|
||||
return "feishu_timeout", "", "Export or request timed out; will retry on the next sync"
|
||||
case strings.Contains(s, "server error"):
|
||||
return "feishu_server_unavailable", "", "Feishu service temporarily unavailable; will retry on the next sync"
|
||||
case strings.Contains(s, "api error"),
|
||||
strings.Contains(s, "export task failed"),
|
||||
strings.Contains(s, "download failed"):
|
||||
if v := feishuErrorCode(err.Error()); v != "" {
|
||||
return "feishu_api_error", v, fmt.Sprintf("Feishu API error (code=%s); will retry on the next sync", v)
|
||||
}
|
||||
return "feishu_api_error_generic", "", "Feishu API error; will retry on the next sync"
|
||||
default:
|
||||
return "sync_failed", "", "Sync failed; will retry on the next sync"
|
||||
}
|
||||
}
|
||||
|
||||
// FeishuErrorItemMeta builds the metadata for a failed item: the raw error (for
|
||||
// server logs) plus the classified i18n code / param / fallback (for a
|
||||
// localisable SyncItemError in the UI), merged with any caller-supplied extras.
|
||||
func FeishuErrorItemMeta(err error, extra map[string]string) map[string]string {
|
||||
code, codeValue, fallback := feishuFailure(err)
|
||||
m := map[string]string{
|
||||
"error": err.Error(),
|
||||
"error_reason_code": code,
|
||||
"error_reason": fallback,
|
||||
}
|
||||
if codeValue != "" {
|
||||
m["error_reason_code_value"] = codeValue
|
||||
}
|
||||
maps.Copy(m, extra)
|
||||
return m
|
||||
}
|
||||
|
||||
// parseableAttachmentExts are attachment extensions worth ingesting as their
|
||||
// own knowledge entries; other files (icons, tiny decor) are skipped.
|
||||
var parseableAttachmentExts = map[string]bool{
|
||||
".pdf": true, ".doc": true, ".docx": true, ".xls": true, ".xlsx": true,
|
||||
".ppt": true, ".pptx": true, ".txt": true, ".md": true, ".csv": true,
|
||||
}
|
||||
|
||||
// MinAttachmentBytes filters out decorative micro-files.
|
||||
const MinAttachmentBytes = 2 * 1024
|
||||
|
||||
// SupportedImageExt sniffs image bytes and returns the filename extension and
|
||||
// content type WeKnora accepts for a standalone image knowledge item (png/jpg/
|
||||
// gif — the image set isValidFileType admits). ok is false for non-image or
|
||||
// unsupported formats (e.g. webp/bmp), which the caller skips rather than
|
||||
// mislabel — a wrong extension would fail parsing. The detected content type is
|
||||
// returned even when ok is false so the caller can log it without re-sniffing.
|
||||
func SupportedImageExt(data []byte) (ext, contentType string, ok bool) {
|
||||
switch ct := http.DetectContentType(data); ct {
|
||||
case "image/png":
|
||||
return ".png", ct, true
|
||||
case "image/jpeg":
|
||||
return ".jpg", ct, true
|
||||
case "image/gif":
|
||||
return ".gif", ct, true
|
||||
default:
|
||||
return "", ct, false
|
||||
}
|
||||
}
|
||||
|
||||
// ParseFeishuConfig extracts and validates Feishu/Lark-specific configuration.
|
||||
//
|
||||
// base_url stays an explicit override so existing data sources that pointed a
|
||||
// "feishu" connector at open.larksuite.com keep working; when it is unset the
|
||||
// region's own host is filled in, making the resolved Config.BaseURL concrete
|
||||
// for everything downstream.
|
||||
func ParseFeishuConfig(config *types.DataSourceConfig, region Region) (*Config, error) {
|
||||
if config == nil {
|
||||
return nil, fmt.Errorf("config is nil")
|
||||
}
|
||||
|
||||
credBytes, err := json.Marshal(config.Credentials)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("marshal credentials: %w", err)
|
||||
}
|
||||
|
||||
var feishuConfig Config
|
||||
if err := json.Unmarshal(credBytes, &feishuConfig); err != nil {
|
||||
return nil, fmt.Errorf("parse %s credentials: %w", region.ConnectorType, err)
|
||||
}
|
||||
|
||||
if feishuConfig.AppID == "" || feishuConfig.AppSecret == "" {
|
||||
return nil, fmt.Errorf("%s app_id and app_secret are required", region.ConnectorType)
|
||||
}
|
||||
|
||||
if feishuConfig.BaseURL == "" {
|
||||
feishuConfig.BaseURL = region.OpenBaseURL
|
||||
}
|
||||
|
||||
// Timezone is a display setting (bitable date rendering), not a credential, so
|
||||
// it lives in Settings. Empty falls back to GMT+8 in resolveLocation.
|
||||
if feishuConfig.Timezone == "" && config.Settings != nil {
|
||||
if tz, ok := config.Settings["timezone"].(string); ok {
|
||||
feishuConfig.Timezone = strings.TrimSpace(tz)
|
||||
}
|
||||
}
|
||||
|
||||
if err := datasource.ValidateConnectorBaseURL(feishuConfig.GetBaseURL()); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return &feishuConfig, nil
|
||||
}
|
||||
|
||||
// IsSupportedDocType checks if a Feishu document type can be synced.
|
||||
// mindnote and slides have no content read API and are skipped.
|
||||
func IsSupportedDocType(objType string) bool {
|
||||
switch objType {
|
||||
case "docx", "doc", "sheet", "bitable", "file":
|
||||
return true
|
||||
default:
|
||||
// mindnote, slides — no content retrieval API available
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
// ParseFeishuTimestamp parses a Feishu unix timestamp string (seconds) into time.Time.
|
||||
func ParseFeishuTimestamp(ts string) time.Time {
|
||||
if ts == "" {
|
||||
return time.Time{}
|
||||
}
|
||||
sec, err := strconv.ParseInt(ts, 10, 64)
|
||||
if err != nil {
|
||||
return time.Time{}
|
||||
}
|
||||
return time.Unix(sec, 0)
|
||||
}
|
||||
|
||||
// SanitizeFileName removes characters that are invalid in filenames and
|
||||
// truncates at a UTF-8 rune boundary. Raw byte truncation would split a
|
||||
// multi-byte codepoint (Chinese chars are 3 bytes) and produce invalid UTF-8
|
||||
// that downstream validation (utf8.ValidString) rejects.
|
||||
//
|
||||
// The extension is preserved across truncation: only the base name is trimmed,
|
||||
// so a long attachment name like "很长的名字….pdf" keeps its ".pdf" suffix that
|
||||
// downstream file-type classification depends on.
|
||||
func SanitizeFileName(name string) string {
|
||||
if name == "" {
|
||||
return "untitled"
|
||||
}
|
||||
replacer := strings.NewReplacer(
|
||||
"/", "_", "\\", "_", ":", "_", "*", "_",
|
||||
"?", "_", "\"", "_", "<", "_", ">", "_", "|", "_",
|
||||
)
|
||||
result := replacer.Replace(name)
|
||||
const maxBytes = 200
|
||||
if len(result) <= maxBytes {
|
||||
return result
|
||||
}
|
||||
ext := filepath.Ext(result)
|
||||
if len(ext) >= maxBytes {
|
||||
// pathological: extension alone overflows the budget → drop it
|
||||
ext = ""
|
||||
}
|
||||
base := truncateUTF8(result[:len(result)-len(ext)], maxBytes-len(ext))
|
||||
return base + ext
|
||||
}
|
||||
|
||||
// truncateUTF8 shortens s to at most maxBytes bytes without splitting a
|
||||
// multi-byte rune: after a hard byte cut it trims any trailing partial codepoint.
|
||||
func truncateUTF8(s string, maxBytes int) string {
|
||||
if len(s) <= maxBytes {
|
||||
return s
|
||||
}
|
||||
s = s[:maxBytes]
|
||||
for len(s) > 0 {
|
||||
r, size := utf8.DecodeLastRuneInString(s)
|
||||
if r != utf8.RuneError || size != 1 {
|
||||
break
|
||||
}
|
||||
s = s[:len(s)-1]
|
||||
}
|
||||
return s
|
||||
}
|
||||
|
||||
// DocxFetchInput is the unified description of one docx document from either
|
||||
// source (wiki node or Drive file) that FetchDocxWithBlocks needs.
|
||||
type DocxFetchInput struct {
|
||||
// WeKnora external_id: wiki=node.NodeToken, drive=file.Token
|
||||
DocToken string
|
||||
// Feishu docx document token
|
||||
ObjToken string
|
||||
Title string
|
||||
URL string
|
||||
ResourceID string
|
||||
EditTime time.Time
|
||||
BaseMeta map[string]string
|
||||
MultimodalEnabled bool
|
||||
}
|
||||
|
||||
// FetchDocxWithBlocks retrieves a docx document via the blocks API, converts it
|
||||
// to Markdown, and returns a main item plus any parseable attachment/image
|
||||
// sub-items. Falls back to the export API if the blocks API errors or renders
|
||||
// empty. Shared by the wiki Connector and the Drive DriveConnector.
|
||||
func FetchDocxWithBlocks(ctx context.Context, client *Client, in DocxFetchInput) ([]*types.FetchedItem, error) {
|
||||
// FEISHU_DOCX_PARSE_MODE selects the docx parsing path. The blocks path
|
||||
// renders image blocks as empty `![图片]()` placeholders and fans images out
|
||||
// into separate knowledge items, which breaks image↔document association in
|
||||
// retrieval/wiki/agent. The export path yields a .docx that docreader parses
|
||||
// inline, so images are bound to the parent document via parent_chunk_id
|
||||
// (same as a regular docx upload). Default (unset / "export") uses export so
|
||||
// images associate with the document; set "blocks" for the blocks-first
|
||||
// behaviour (faster, keeps docx attachments, but images are detached).
|
||||
|
||||
// This is a temporary solution. If a better parsing solution is available later, this environment variable will be removed and replaced with a better one.
|
||||
parsingMode := strings.TrimSpace(os.Getenv("FEISHU_DOCX_PARSE_MODE"))
|
||||
if parsingMode == "" {
|
||||
parsingMode = "export"
|
||||
}
|
||||
|
||||
if strings.EqualFold(parsingMode, "export") {
|
||||
item, err := exportDocxFallback(ctx, client, in)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return []*types.FetchedItem{item}, nil
|
||||
}
|
||||
|
||||
blocks, err := client.listDocumentBlocks(ctx, in.ObjToken)
|
||||
if err != nil {
|
||||
logger.Warnf(ctx, "[Feishu] blocks API failed for %s (%s), falling back to export: %v",
|
||||
in.Title, in.ObjToken, err)
|
||||
item, ferr := exportDocxFallback(ctx, client, in)
|
||||
if ferr != nil {
|
||||
return nil, ferr
|
||||
}
|
||||
// Do NOT set ReplacesSubtree here (see the wiki history: a transient
|
||||
// blocks failure must not sweep good attachment children from the prior
|
||||
// blocks-path sync with nothing to replace them).
|
||||
return []*types.FetchedItem{item}, nil
|
||||
}
|
||||
|
||||
md, atts, err := blocksToMarkdown(ctx, client, blocks)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("convert blocks %s: %w", in.Title, err)
|
||||
}
|
||||
|
||||
if len(strings.TrimSpace(string(md))) == 0 {
|
||||
logger.Infof(ctx, "[Feishu] doc %s (%s): blocks rendered empty Markdown, falling back to export",
|
||||
in.Title, in.ObjToken)
|
||||
item, ferr := exportDocxFallback(ctx, client, in)
|
||||
if ferr != nil {
|
||||
return nil, ferr
|
||||
}
|
||||
return []*types.FetchedItem{item}, nil
|
||||
}
|
||||
|
||||
main := &types.FetchedItem{
|
||||
ExternalID: in.DocToken,
|
||||
Title: in.Title,
|
||||
Content: md,
|
||||
ContentType: "text/markdown",
|
||||
FileName: SanitizeFileName(in.Title) + ".md",
|
||||
URL: in.URL,
|
||||
UpdatedAt: in.EditTime,
|
||||
SourceResourceID: in.ResourceID,
|
||||
Metadata: in.BaseMeta,
|
||||
ReplacesSubtree: true, // sweep stale attachment sub-items on re-sync
|
||||
}
|
||||
items := []*types.FetchedItem{main}
|
||||
|
||||
keep := make([]string, 0, len(atts))
|
||||
childMeta := func() map[string]string {
|
||||
m := maps.Clone(in.BaseMeta)
|
||||
m["parent_node_token"] = in.DocToken
|
||||
m["attachment"] = "true"
|
||||
return m
|
||||
}
|
||||
for _, a := range atts {
|
||||
childID := types.SubtreeChildID(in.DocToken, "file", a.FileToken)
|
||||
keep = append(keep, childID) // present in the doc → never sweep as stale
|
||||
ext := strings.ToLower(filepath.Ext(a.Name))
|
||||
if ext == "" {
|
||||
logger.Warnf(ctx, "[Feishu] doc %s: skipping attachment with no usable filename (token=%s name=%q)",
|
||||
in.ObjToken, a.FileToken, a.Name)
|
||||
continue
|
||||
}
|
||||
if !parseableAttachmentExts[ext] {
|
||||
continue
|
||||
}
|
||||
data, derr := client.downloadMediaFile(ctx, a.FileToken)
|
||||
if derr != nil {
|
||||
logger.Warnf(ctx, "[Feishu] doc %s: attachment %q (token=%s) download failed: %v",
|
||||
in.ObjToken, a.Name, a.FileToken, derr)
|
||||
items = append(items, &types.FetchedItem{
|
||||
ExternalID: childID,
|
||||
Title: a.Name,
|
||||
SourceResourceID: in.ResourceID,
|
||||
Metadata: FeishuErrorItemMeta(derr, childMeta()),
|
||||
})
|
||||
continue
|
||||
}
|
||||
if len(data) < MinAttachmentBytes {
|
||||
logger.Infof(ctx, "[Feishu] doc %s: skipping tiny attachment %q (token=%s, %d bytes < %d)",
|
||||
in.ObjToken, a.Name, a.FileToken, len(data), MinAttachmentBytes)
|
||||
continue
|
||||
}
|
||||
items = append(items, &types.FetchedItem{
|
||||
ExternalID: childID,
|
||||
Title: a.Name,
|
||||
Content: data,
|
||||
ContentType: "application/octet-stream",
|
||||
FileName: SanitizeFileName(a.Name),
|
||||
URL: in.URL,
|
||||
UpdatedAt: in.EditTime,
|
||||
SourceResourceID: in.ResourceID,
|
||||
Metadata: childMeta(),
|
||||
})
|
||||
}
|
||||
|
||||
imgMeta := func() map[string]string {
|
||||
m := maps.Clone(in.BaseMeta)
|
||||
m["parent_node_token"] = in.DocToken
|
||||
m["embedded_image"] = "true"
|
||||
return m
|
||||
}
|
||||
for _, b := range blocks {
|
||||
if b.BlockType != BlockTypeImage || b.Image == nil || b.Image.Token == "" {
|
||||
continue
|
||||
}
|
||||
childID := types.SubtreeChildID(in.DocToken, "image", b.Image.Token)
|
||||
keep = append(keep, childID) // present in the doc → never sweep as stale
|
||||
if !in.MultimodalEnabled {
|
||||
continue // KB can't OCR images; the inline placeholder is all we keep
|
||||
}
|
||||
data, derr := client.downloadMediaFile(ctx, b.Image.Token)
|
||||
if derr != nil {
|
||||
logger.Warnf(ctx, "[Feishu] doc %s: image (token=%s) download failed: %v",
|
||||
in.ObjToken, b.Image.Token, derr)
|
||||
items = append(items, &types.FetchedItem{
|
||||
ExternalID: childID,
|
||||
Title: fmt.Sprintf("%s(内嵌图片)", in.Title),
|
||||
SourceResourceID: in.ResourceID,
|
||||
Metadata: FeishuErrorItemMeta(derr, imgMeta()),
|
||||
})
|
||||
continue
|
||||
}
|
||||
if len(data) < MinAttachmentBytes {
|
||||
continue // decorative micro-image (icon/spacer)
|
||||
}
|
||||
ext, contentType, ok := SupportedImageExt(data)
|
||||
if !ok {
|
||||
logger.Warnf(ctx, "[Feishu] doc %s: skipping image (token=%s) of unsupported type %q",
|
||||
in.ObjToken, b.Image.Token, contentType)
|
||||
continue
|
||||
}
|
||||
items = append(items, &types.FetchedItem{
|
||||
ExternalID: childID,
|
||||
Title: fmt.Sprintf("%s(内嵌图片)", in.Title),
|
||||
Content: data,
|
||||
ContentType: contentType,
|
||||
FileName: "image-" + b.Image.Token + ext,
|
||||
URL: in.URL,
|
||||
UpdatedAt: in.EditTime,
|
||||
SourceResourceID: in.ResourceID,
|
||||
Metadata: imgMeta(),
|
||||
})
|
||||
}
|
||||
|
||||
main.SubtreeKeep = keep
|
||||
return items, nil
|
||||
}
|
||||
|
||||
// exportDocxFallback exports a docx document via the async export API and
|
||||
// returns a single FetchedItem containing the exported .docx binary. Used by
|
||||
// FetchDocxWithBlocks when the blocks API is unavailable or renders empty.
|
||||
func exportDocxFallback(ctx context.Context, client *Client, in DocxFetchInput) (*types.FetchedItem, error) {
|
||||
data, fileName, err := client.ExportAndDownload(ctx, in.ObjToken, "docx")
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("export %s (docx): %w", in.Title, err)
|
||||
}
|
||||
|
||||
ext := ExportFileExtToSuffix[ObjTypeToExportFileExtension["docx"]]
|
||||
if fileName == "" {
|
||||
fileName = SanitizeFileName(in.Title) + ext
|
||||
} else if !strings.HasSuffix(strings.ToLower(fileName), ext) {
|
||||
fileName = SanitizeFileName(fileName) + ext
|
||||
}
|
||||
|
||||
return &types.FetchedItem{
|
||||
ExternalID: in.DocToken,
|
||||
Title: in.Title,
|
||||
Content: data,
|
||||
ContentType: "application/octet-stream",
|
||||
FileName: fileName,
|
||||
URL: in.URL,
|
||||
UpdatedAt: in.EditTime,
|
||||
SourceResourceID: in.ResourceID,
|
||||
Metadata: in.BaseMeta,
|
||||
}, nil
|
||||
}
|
||||
@@ -0,0 +1,45 @@
|
||||
package core
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestFetchTally_CountsAndSummary(t *testing.T) {
|
||||
tally := newFetchTally(13)
|
||||
tally.fetch()
|
||||
tally.fetch()
|
||||
tally.fetch()
|
||||
tally.Skip("mindnote")
|
||||
tally.Skip("mindnote")
|
||||
tally.Skip("slides")
|
||||
tally.fail()
|
||||
|
||||
if got := tally.skipped(); got != 3 {
|
||||
t.Errorf("skipped() = %d, want 3", got)
|
||||
}
|
||||
|
||||
summary := tally.summary()
|
||||
for _, want := range []string{
|
||||
"discovered=13",
|
||||
"fetched=3",
|
||||
"failed=1",
|
||||
"skipped_unsupported=3",
|
||||
"mindnote:2",
|
||||
"slides:1",
|
||||
} {
|
||||
if !strings.Contains(summary, want) {
|
||||
t.Errorf("summary() = %q, missing %q", summary, want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestFetchTally_EmptyHasNoSkips(t *testing.T) {
|
||||
tally := newFetchTally(0)
|
||||
if got := tally.skipped(); got != 0 {
|
||||
t.Errorf("skipped() = %d, want 0", got)
|
||||
}
|
||||
if !strings.Contains(tally.summary(), "skipped_unsupported=0") {
|
||||
t.Errorf("summary() = %q, want skipped_unsupported=0", tally.summary())
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,359 @@
|
||||
// Package feishu implements the Feishu (飞书/Lark) data source connector for WeKnora.
|
||||
//
|
||||
// It syncs documents from Feishu Wiki spaces and cloud documents into WeKnora knowledge bases.
|
||||
//
|
||||
// Feishu API docs:
|
||||
// - Wiki spaces: https://open.feishu.cn/document/server-docs/docs/wiki-v2/space/list
|
||||
// - Wiki nodes: https://open.feishu.cn/document/server-docs/docs/wiki-v2/space-node/list
|
||||
// - Export tasks: https://open.feishu.cn/document/server-docs/docs/drive-v1/export_task/export-user-guide
|
||||
// - File download: https://open.feishu.cn/document/server-docs/docs/drive-v1/file/download
|
||||
// - Auth: https://open.feishu.cn/document/server-docs/authentication-management/access-token/tenant_access_token_internal
|
||||
package core
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
// Config holds Feishu-specific configuration for the data source connector.
|
||||
// Uses the self-built app (企业自建应用) authentication model.
|
||||
type Config struct {
|
||||
// App ID from Feishu developer console
|
||||
AppID string `json:"app_id"`
|
||||
|
||||
// App Secret from Feishu developer console
|
||||
AppSecret string `json:"app_secret"`
|
||||
|
||||
// Base URL for Feishu API (default: https://open.feishu.cn)
|
||||
// Use https://open.larksuite.com for Lark (international) deployments
|
||||
BaseURL string `json:"base_url,omitempty"`
|
||||
|
||||
// Timezone is the IANA name (e.g. "Asia/Shanghai", "America/New_York") used to
|
||||
// render bitable date cells. Feishu stores a date as a UTC instant, but the
|
||||
// calendar date a user sees is that instant in the table's timezone, so
|
||||
// rendering in UTC shifts the date. Empty defaults to GMT+8, which matches
|
||||
// Feishu (mainland) tenants; set it for Lark tenants in other zones.
|
||||
Timezone string `json:"timezone,omitempty"`
|
||||
}
|
||||
|
||||
// defaultTimezoneOffsetSeconds is GMT+8 (the Feishu mainland default), used when
|
||||
// no timezone is configured. A fixed zone avoids any dependency on system tzdata,
|
||||
// which minimal container images may omit.
|
||||
const defaultTimezoneOffsetSeconds = 8 * 3600
|
||||
|
||||
// resolveLocation returns the *time.Location used to render bitable date cells.
|
||||
// An empty name yields a fixed GMT+8 zone (no tzdata dependency); a named zone is
|
||||
// loaded from the system tz database, falling back to GMT+8 if it is unavailable.
|
||||
func resolveLocation(name string) *time.Location {
|
||||
if name != "" {
|
||||
if loc, err := time.LoadLocation(name); err == nil {
|
||||
return loc
|
||||
}
|
||||
}
|
||||
return time.FixedZone("GMT+8", defaultTimezoneOffsetSeconds)
|
||||
}
|
||||
|
||||
// defaultBaseURL is the default Feishu Open Platform API base URL.
|
||||
const defaultBaseURL = "https://open.feishu.cn"
|
||||
|
||||
// larkBaseURL is the Lark (international) API base URL.
|
||||
const larkBaseURL = "https://open.larksuite.com"
|
||||
|
||||
// GetBaseURL returns the effective base URL, defaulting to Feishu if not set.
|
||||
func (c *Config) GetBaseURL() string {
|
||||
if c.BaseURL != "" {
|
||||
return c.BaseURL
|
||||
}
|
||||
return defaultBaseURL
|
||||
}
|
||||
|
||||
// --- Export format constants ---
|
||||
// Used by the export task API: POST /drive/v1/export_tasks
|
||||
|
||||
const (
|
||||
// ExportTypeDocx exports Feishu documents to .docx format.
|
||||
ExportTypeDocx = "docx"
|
||||
// ExportTypeXlsx exports spreadsheets / bitable to .xlsx format.
|
||||
ExportTypeXlsx = "xlsx"
|
||||
// ExportTypePDF exports documents to .pdf format (fallback).
|
||||
ExportTypePDF = "pdf"
|
||||
)
|
||||
|
||||
// ObjTypeToExportFileExtension maps Feishu obj_type to the best export file_extension.
|
||||
var ObjTypeToExportFileExtension = map[string]string{
|
||||
"docx": ExportTypeDocx,
|
||||
"doc": ExportTypeDocx,
|
||||
"sheet": ExportTypeXlsx,
|
||||
"bitable": ExportTypeXlsx,
|
||||
}
|
||||
|
||||
// ObjTypeToExportType maps Feishu obj_type to the export API "type" parameter.
|
||||
// See: https://open.feishu.cn/document/server-docs/docs/drive-v1/export_task/create
|
||||
var ObjTypeToExportType = map[string]string{
|
||||
"docx": "docx",
|
||||
"doc": "doc",
|
||||
"sheet": "sheet",
|
||||
"bitable": "bitable",
|
||||
}
|
||||
|
||||
// ExportFileExtToSuffix maps export file_extension to the file suffix for FileName.
|
||||
var ExportFileExtToSuffix = map[string]string{
|
||||
ExportTypeDocx: ".docx",
|
||||
ExportTypeXlsx: ".xlsx",
|
||||
ExportTypePDF: ".pdf",
|
||||
}
|
||||
|
||||
// --- Feishu API response structures ---
|
||||
|
||||
// ApiResponse is the common Feishu API response wrapper.
|
||||
type ApiResponse struct {
|
||||
Code int `json:"code"`
|
||||
Msg string `json:"msg"`
|
||||
}
|
||||
|
||||
// TokenResponse is the response for tenant_access_token API.
|
||||
type TokenResponse struct {
|
||||
ApiResponse
|
||||
TenantAccessToken string `json:"tenant_access_token"`
|
||||
Expire int `json:"expire"` // seconds
|
||||
}
|
||||
|
||||
// WikiSpaceListData is the data payload of WikiSpaceListResponse.
|
||||
type WikiSpaceListData struct {
|
||||
Items []WikiSpace `json:"items"`
|
||||
HasMore bool `json:"has_more"`
|
||||
PageToken string `json:"page_token"`
|
||||
}
|
||||
|
||||
// WikiSpaceListResponse is the response for GET /open-apis/wiki/v2/spaces.
|
||||
type WikiSpaceListResponse struct {
|
||||
ApiResponse
|
||||
Data WikiSpaceListData `json:"data"`
|
||||
}
|
||||
|
||||
// WikiSpace represents a Feishu Wiki space.
|
||||
type WikiSpace struct {
|
||||
SpaceID string `json:"space_id"`
|
||||
Name string `json:"name"`
|
||||
Description string `json:"description"`
|
||||
Visibility string `json:"visibility"` // "public" or "private"
|
||||
}
|
||||
|
||||
// WikiNodeListData is the data payload of WikiNodeListResponse.
|
||||
type WikiNodeListData struct {
|
||||
Items []WikiNode `json:"items"`
|
||||
HasMore bool `json:"has_more"`
|
||||
PageToken string `json:"page_token"`
|
||||
}
|
||||
|
||||
// WikiNodeListResponse is the response for GET /open-apis/wiki/v2/spaces/:space_id/nodes.
|
||||
type WikiNodeListResponse struct {
|
||||
ApiResponse
|
||||
Data WikiNodeListData `json:"data"`
|
||||
}
|
||||
|
||||
// WikiNode represents a node (document or folder) in a Feishu Wiki space.
|
||||
type WikiNode struct {
|
||||
SpaceID string `json:"space_id"`
|
||||
NodeToken string `json:"node_token"`
|
||||
ObjToken string `json:"obj_token"` // document token
|
||||
ObjType string `json:"obj_type"` // "doc", "sheet", "mindnote", "bitable", "file", "docx", "slides"
|
||||
ParentNodeID string `json:"parent_node_token"`
|
||||
NodeType string `json:"node_type"` // "origin" or "shortcut"
|
||||
OriginNodeID string `json:"origin_node_id"`
|
||||
OriginSpaceID string `json:"origin_space_id"`
|
||||
HasChild bool `json:"has_child"`
|
||||
Title string `json:"title"`
|
||||
Creator string `json:"creator"`
|
||||
Owner string `json:"owner"`
|
||||
ObjCreateTime string `json:"obj_create_time"` // document creation time (unix timestamp string)
|
||||
ObjEditTime string `json:"obj_edit_time"` // document last edit time (unix timestamp string) — tracks content changes
|
||||
NodeCreateTime string `json:"node_create_time"` // node creation time (unix timestamp string)
|
||||
NodeEditTime string `json:"node_edit_time"` // node edit time (unix timestamp string) — only tracks node attribute changes
|
||||
}
|
||||
|
||||
// WikiNodeInfoData is the data payload of WikiNodeInfoResponse.
|
||||
type WikiNodeInfoData struct {
|
||||
Node WikiNode `json:"node"`
|
||||
}
|
||||
|
||||
// WikiNodeInfoResponse is the response for GET /open-apis/wiki/v2/spaces/get_node.
|
||||
type WikiNodeInfoResponse struct {
|
||||
ApiResponse
|
||||
Data WikiNodeInfoData `json:"data"`
|
||||
}
|
||||
|
||||
// --- Export task API responses ---
|
||||
|
||||
// docRawContentData is the data payload of docRawContentResponse.
|
||||
type docRawContentData struct {
|
||||
Content string `json:"content"`
|
||||
}
|
||||
|
||||
// docRawContentResponse is the response for GET /open-apis/docx/v1/documents/:document_id/raw_content.
|
||||
// Deprecated: prefer export API for full-fidelity document export.
|
||||
type docRawContentResponse struct {
|
||||
ApiResponse
|
||||
Data docRawContentData `json:"data"`
|
||||
}
|
||||
|
||||
// ExportTaskCreateData is the data payload of ExportTaskCreateResponse.
|
||||
type ExportTaskCreateData struct {
|
||||
Ticket string `json:"ticket"`
|
||||
}
|
||||
|
||||
// ExportTaskCreateResponse is the response for POST /drive/v1/export_tasks.
|
||||
type ExportTaskCreateResponse struct {
|
||||
ApiResponse
|
||||
Data ExportTaskCreateData `json:"data"`
|
||||
}
|
||||
|
||||
// ExportTaskResult is the per-task result inside ExportTaskStatusData.
|
||||
type ExportTaskResult struct {
|
||||
FileToken string `json:"file_token"`
|
||||
FileSize int64 `json:"file_size"`
|
||||
// JobStatus: 0=success, 1=initializing, 2=processing
|
||||
JobStatus int `json:"job_status"`
|
||||
JobErrorMsg string `json:"job_error_msg"`
|
||||
FileName string `json:"file_name"`
|
||||
}
|
||||
|
||||
// ExportTaskStatusData is the data payload of ExportTaskStatusResponse.
|
||||
type ExportTaskStatusData struct {
|
||||
Result ExportTaskResult `json:"result"`
|
||||
}
|
||||
|
||||
// ExportTaskStatusResponse is the response for GET /drive/v1/export_tasks/{ticket}.
|
||||
type ExportTaskStatusResponse struct {
|
||||
ApiResponse
|
||||
Data ExportTaskStatusData `json:"data"`
|
||||
}
|
||||
|
||||
// --- File download response ---
|
||||
|
||||
// driveFileMeta is one entry of driveFileMetaData.Metas.
|
||||
type driveFileMeta struct {
|
||||
DocToken string `json:"doc_token"`
|
||||
DocType string `json:"doc_type"`
|
||||
Title string `json:"title"`
|
||||
}
|
||||
|
||||
// driveFileMetaData is the data payload of driveFileMetaResponse.
|
||||
type driveFileMetaData struct {
|
||||
Metas []driveFileMeta `json:"metas"`
|
||||
}
|
||||
|
||||
// driveFileMetaResponse is the response for GET /drive/v1/metas for file type nodes.
|
||||
type driveFileMetaResponse struct {
|
||||
ApiResponse
|
||||
Data driveFileMetaData `json:"data"`
|
||||
}
|
||||
|
||||
// FeishuCursor stores incremental sync state for Feishu.
|
||||
type FeishuCursor struct {
|
||||
// LastSyncTime is the timestamp of the last successful sync.
|
||||
LastSyncTime time.Time `json:"last_sync_time"`
|
||||
|
||||
// SpaceNodeTimes maps space_id -> node_token -> last known edit time.
|
||||
// Used to detect which nodes have changed since last sync.
|
||||
SpaceNodeTimes map[string]map[string]string `json:"space_node_times,omitempty"`
|
||||
}
|
||||
|
||||
// --- Drive (云盘) file listing types (feishu_drive / lark_drive connectors) ---
|
||||
// Added by feat/datasource-feishu-drive. These are independent of the wiki
|
||||
// types above and do not affect the wiki connector.
|
||||
|
||||
// DriveFile represents a file/folder in Feishu Drive (云空间). Returned by
|
||||
// GET /open-apis/drive/v1/files?folder_token=xxx. The list API returns
|
||||
// modified_time directly (verified), so no batch_query/metas call is needed for
|
||||
// incremental detection - see ADR-0002.
|
||||
type DriveFile struct {
|
||||
Token string `json:"token"`
|
||||
Name string `json:"name"`
|
||||
Type string `json:"type"` // doc/docx/sheet/bitable/file/folder/shortcut/mindnote/slides/board
|
||||
ParentToken string `json:"parent_token"`
|
||||
URL string `json:"url"`
|
||||
CreatedTime string `json:"created_time"` // unix seconds string
|
||||
ModifiedTime string `json:"modified_time"` // unix seconds string - 等价知识库 obj_edit_time
|
||||
OwnerID string `json:"owner_id"`
|
||||
// ShortcutInfo is populated only for type=="shortcut". target_type can only
|
||||
// be doc/sheet/mindnote/bitable/file/docx (Feishu does not allow shortcuts to
|
||||
// folders, verified) - see ADR-0002 / glossary shortcut entry.
|
||||
ShortcutInfo *driveShortcutInfo `json:"shortcut_info,omitempty"`
|
||||
}
|
||||
|
||||
// driveShortcutInfo is the target metadata of a Drive shortcut.
|
||||
type driveShortcutInfo struct {
|
||||
TargetToken string `json:"target_token"`
|
||||
TargetType string `json:"target_type"`
|
||||
}
|
||||
|
||||
// DriveFileListData is the data payload of DriveFileListResponse.
|
||||
type DriveFileListData struct {
|
||||
Files []DriveFile `json:"files"`
|
||||
HasMore bool `json:"has_more"`
|
||||
NextPageToken string `json:"next_page_token"`
|
||||
}
|
||||
|
||||
// DriveFileListResponse is the response for GET /open-apis/drive/v1/files.
|
||||
type DriveFileListResponse struct {
|
||||
ApiResponse
|
||||
Data DriveFileListData `json:"data"`
|
||||
}
|
||||
|
||||
// driveFolderMetaData is the data payload of driveFolderMetaResponse.
|
||||
type driveFolderMetaData struct {
|
||||
ID string `json:"id"`
|
||||
Name string `json:"name"`
|
||||
Token string `json:"token"`
|
||||
CreateUid string `json:"createUid"`
|
||||
EditUid string `json:"editUid"`
|
||||
ParentID string `json:"parentId"`
|
||||
OwnUid string `json:"ownUid"`
|
||||
}
|
||||
|
||||
// driveFolderMetaResponse is the response for GET /open-apis/drive/explorer/v2/folder/:folderToken/meta.
|
||||
// Used to resolve a root folder's human-readable name (the list API only returns
|
||||
// the folder's children, not the folder itself).
|
||||
type driveFolderMetaResponse struct {
|
||||
ApiResponse
|
||||
Data driveFolderMetaData `json:"data"`
|
||||
}
|
||||
|
||||
// DriveFileListFailure records a single sub-folder listing that failed during a
|
||||
// recursive walk. Mirrors WikiNodeListFailure.
|
||||
type DriveFileListFailure struct {
|
||||
FolderToken string
|
||||
Err error
|
||||
}
|
||||
|
||||
// PartialDriveFileListError aggregates per-folder listing failures so the walk
|
||||
// can continue and the caller can still surface the partial result. Mirrors
|
||||
// PartialWikiNodeListError.
|
||||
type PartialDriveFileListError struct {
|
||||
Failures []DriveFileListFailure
|
||||
}
|
||||
|
||||
func (e *PartialDriveFileListError) Error() string {
|
||||
if e == nil || len(e.Failures) == 0 {
|
||||
return "partial drive file listing failed"
|
||||
}
|
||||
parts := make([]string, 0, len(e.Failures))
|
||||
for _, failure := range e.Failures {
|
||||
parts = append(parts, failure.Err.Error())
|
||||
}
|
||||
return strings.Join(parts, "; ")
|
||||
}
|
||||
|
||||
// FeishuDriveCursor stores incremental sync state for Feishu Drive (云盘).
|
||||
// Structurally symmetric with FeishuCursor: outer key = resourceID
|
||||
// ("folderToken" or "folderToken:fileToken"), inner key = file_token,
|
||||
// value = modified_time. See ADR-0001.
|
||||
type FeishuDriveCursor struct {
|
||||
// LastSyncTime is the timestamp of the last successful sync.
|
||||
LastSyncTime time.Time `json:"last_sync_time"`
|
||||
|
||||
// FileTimes maps resourceID -> file_token -> last known modified_time.
|
||||
// Used to detect which files have changed since last sync.
|
||||
FileTimes map[string]map[string]string `json:"file_times,omitempty"`
|
||||
}
|
||||
@@ -0,0 +1,615 @@
|
||||
package drive
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/datasource"
|
||||
"github.com/Tencent/WeKnora/internal/datasource/connector/feishu/core"
|
||||
"github.com/Tencent/WeKnora/internal/logger"
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
)
|
||||
|
||||
// DriveConnector implements the datasource.Connector (and StreamingConnector)
|
||||
// interface for Feishu/Lark Drive (云盘) mode. It shares core.Client/core.Config/core.Region
|
||||
// and the export/download logic with the wiki Connector; only resource
|
||||
// enumeration and Fetch dispatch differ. See 飞书云盘数据源设计.md and
|
||||
// ADR-0001..0004.
|
||||
type DriveConnector struct {
|
||||
region core.Region
|
||||
}
|
||||
|
||||
// NewDriveConnector creates a Drive connector for the given region
|
||||
// (RegionFeishuDrive or RegionLarkDrive).
|
||||
func NewDriveConnector(region core.Region) *DriveConnector {
|
||||
return &DriveConnector{region: region}
|
||||
}
|
||||
|
||||
// Drive supports resumable streaming sync; the service prefers FetchStream over
|
||||
// FetchAll/FetchIncremental when a connector implements StreamingConnector.
|
||||
var _ datasource.StreamingConnector = (*DriveConnector)(nil)
|
||||
|
||||
// Type returns the connector type identifier.
|
||||
func (c *DriveConnector) Type() string {
|
||||
return c.region.ConnectorType
|
||||
}
|
||||
|
||||
// Validate verifies that the Drive configuration is valid by testing
|
||||
// connectivity. It does not validate folder_token here - that is done in
|
||||
// ListResources when the user loads the tree root. Mirrors the wiki Connector.
|
||||
func (c *DriveConnector) Validate(ctx context.Context, config *types.DataSourceConfig) error {
|
||||
feishuConfig, err := core.ParseFeishuConfig(config, c.region)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
client := core.NewClient(feishuConfig)
|
||||
if err := client.Ping(ctx); err != nil {
|
||||
return fmt.Errorf("%s connection failed: %w", c.region.Label, err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// ListResources lists Feishu Drive resources for selection, loading the tree
|
||||
// lazily one level at a time. Mirrors the wiki Connector.ListResources shape
|
||||
// but with the Drive difference that the root is user-supplied (config.
|
||||
// ResourceIDs[0]) rather than enumerated via ListWikiSpaces.
|
||||
//
|
||||
// - parentID == "" -> return the user-supplied root
|
||||
// folder (from config.ResourceIDs[0]) as the single root resource
|
||||
// (HasChildren=true). Drive has no "space list" API, so the root is
|
||||
// user-supplied. folder_token == "" is rejected (ADR-0004).
|
||||
// - parentID == folderToken -> ListDriveFiles(folderToken)
|
||||
// returns the direct children.
|
||||
// - parentID == "folderToken:subFolderToken" -> ListDriveFiles(subFolderToken)
|
||||
// returns that sub-folder's direct children.
|
||||
//
|
||||
// Each core.DriveFile becomes a Resource: folder HasChildren=true, others false.
|
||||
// resourceID encoding: root = folderToken; child = folderToken + ":" + fileToken
|
||||
// (reuses core.FeishuWikiNodeResourceSeparator). See ADR-0001 §3.4.
|
||||
func (c *DriveConnector) ListResources(
|
||||
ctx context.Context, config *types.DataSourceConfig, parentID string,
|
||||
) ([]types.Resource, error) {
|
||||
feishuConfig, err := core.ParseFeishuConfig(config, c.region)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
client := core.NewClient(feishuConfig)
|
||||
|
||||
if parentID == "" {
|
||||
// Root load: read the user-supplied folder_token from config.ResourceIDs.
|
||||
rootFolderToken := driveRootFolderToken(config)
|
||||
if rootFolderToken == "" {
|
||||
return nil, fmt.Errorf("folder_token is required; specify a Drive folder token")
|
||||
}
|
||||
// Validate access by listing the root's direct children (also lazy-loads
|
||||
// the first level for the picker). Reuse the list call rather than a
|
||||
// separate ping.
|
||||
files, err := client.ListDriveFilesAllPages(ctx, rootFolderToken)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("list feishu drive folder %s: %w", rootFolderToken, err)
|
||||
}
|
||||
_ = files // children returned via the parentID == rootFolderToken branch below
|
||||
// Resolve the root folder's human-readable name via the folder meta API.
|
||||
// Best-effort: on failure (no permission / not found) fall back to the
|
||||
// folder_token so the picker still renders something usable.
|
||||
folderName := rootFolderToken
|
||||
if meta, mErr := client.GetDriveFolderMeta(ctx, rootFolderToken); mErr != nil {
|
||||
logger.Warnf(ctx, "[FeishuDrive] resolve root folder name failed: %v (falling back to token)", mErr)
|
||||
} else if meta.Data.Name != "" {
|
||||
folderName = meta.Data.Name
|
||||
}
|
||||
return []types.Resource{c.driveFolderToResource(rootFolderToken, "", rootFolderToken, folderName)}, nil
|
||||
}
|
||||
|
||||
// Lazy load: list only the direct children of the given folder.
|
||||
rootFolderToken, folderToken := parseDriveResourceID(parentID)
|
||||
if folderToken == "" {
|
||||
// parentID is a bare root folder token -> list its children.
|
||||
folderToken = rootFolderToken
|
||||
}
|
||||
files, err := client.ListDriveFilesAllPages(ctx, folderToken)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("list feishu drive files under %s: %w", parentID, err)
|
||||
}
|
||||
|
||||
resources := make([]types.Resource, 0, len(files))
|
||||
for _, f := range files {
|
||||
resources = append(resources, c.driveFileToResource(rootFolderToken, f))
|
||||
}
|
||||
return resources, nil
|
||||
}
|
||||
|
||||
// ResolveResourceAncestors returns the resource IDs of every parent folder that
|
||||
// has to be expanded so the lazily-loaded picker can reveal each selection.
|
||||
//
|
||||
// The wiki connector walks up via GetWikiNode (parent_node_token) in O(depth)
|
||||
// single-node queries. Drive has no single-file parent query API (verified -
|
||||
// metas/batch_query does not return parent), so we walk top-down from the root
|
||||
// folder with ListDriveFiles and share the traversal across all selections in
|
||||
// the same root. Best-effort: a broken path just stays collapsed. See ADR-0003.
|
||||
func (c *DriveConnector) ResolveResourceAncestors(
|
||||
ctx context.Context, config *types.DataSourceConfig, resourceIDs []string,
|
||||
) ([]string, error) {
|
||||
feishuConfig, err := core.ParseFeishuConfig(config, c.region)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
client := core.NewClient(feishuConfig)
|
||||
|
||||
seen := make(map[string]bool)
|
||||
ancestors := make([]string, 0)
|
||||
add := func(id string) {
|
||||
if id != "" && !seen[id] {
|
||||
seen[id] = true
|
||||
ancestors = append(ancestors, id)
|
||||
}
|
||||
}
|
||||
|
||||
// Group selections by root folder so one shared traversal covers them all.
|
||||
type selection struct {
|
||||
fileToken string
|
||||
}
|
||||
rootSelections := make(map[string][]selection)
|
||||
for _, rid := range resourceIDs {
|
||||
rootFolderToken, fileToken := parseDriveResourceID(rid)
|
||||
if fileToken == "" {
|
||||
// A root-level selection is already a top-level node; nothing to reveal.
|
||||
continue
|
||||
}
|
||||
add(rootFolderToken)
|
||||
rootSelections[rootFolderToken] = append(rootSelections[rootFolderToken], selection{fileToken})
|
||||
}
|
||||
|
||||
// For each root, BFS from the root: at each folder, ListDriveFiles and check
|
||||
// which selections are direct children (record their parent chain) and which
|
||||
// sub-folders may still contain selections (enqueue). Shared traversal means
|
||||
// selections in the same subtree reuse list calls.
|
||||
for rootFolderToken, sels := range rootSelections {
|
||||
remaining := make(map[string]bool, len(sels))
|
||||
for _, s := range sels {
|
||||
remaining[s.fileToken] = true
|
||||
}
|
||||
// parentChain[fileToken] = resourceID of its parent folder
|
||||
parentChain := make(map[string]string)
|
||||
|
||||
queue := []string{rootFolderToken}
|
||||
for len(queue) > 0 && len(remaining) > 0 {
|
||||
cur := queue[0]
|
||||
queue = queue[1:]
|
||||
|
||||
files, err := client.ListDriveFilesAllPages(ctx, cur)
|
||||
if err != nil {
|
||||
logger.Warnf(ctx, "[FeishuDrive] resolve ancestors: list %s: %v", cur, err)
|
||||
break // best-effort: stop this root's traversal
|
||||
}
|
||||
for _, f := range files {
|
||||
if remaining[f.Token] {
|
||||
delete(remaining, f.Token)
|
||||
// Record the parent chain from root down to this file's parent.
|
||||
chain := buildDriveAncestorChain(rootFolderToken, cur, parentChain)
|
||||
for _, a := range chain {
|
||||
add(a)
|
||||
}
|
||||
}
|
||||
if f.Type == "folder" {
|
||||
parentChain[f.Token] = makeDriveResourceID(rootFolderToken, cur)
|
||||
queue = append(queue, f.Token)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return ancestors, nil
|
||||
}
|
||||
|
||||
// buildDriveAncestorChain walks the parentChain map from cur up to root,
|
||||
// returning the resourceIDs (root, ... , cur's parent) in root-first order.
|
||||
func buildDriveAncestorChain(rootFolderToken, cur string, parentChain map[string]string) []string {
|
||||
var chain []string
|
||||
node := cur
|
||||
for node != "" && node != rootFolderToken {
|
||||
parent, ok := parentChain[node]
|
||||
if !ok {
|
||||
break
|
||||
}
|
||||
chain = append([]string{parent}, chain...)
|
||||
_, parentFolderToken := parseDriveResourceID(parent)
|
||||
node = parentFolderToken
|
||||
}
|
||||
return chain
|
||||
}
|
||||
|
||||
// FetchAll performs a full sync of all documents from the selected Drive
|
||||
// folders. Defensive fallback path - the service prefers FetchStream when the
|
||||
// connector implements StreamingConnector.
|
||||
func (c *DriveConnector) FetchAll(
|
||||
ctx context.Context, config *types.DataSourceConfig, resourceIDs []string,
|
||||
) ([]types.FetchedItem, error) {
|
||||
feishuConfig, err := core.ParseFeishuConfig(config, c.region)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
client := core.NewClient(feishuConfig)
|
||||
return core.FetchAllEngine(ctx, client, config, resourceIDs, driveOps{region: c.region})
|
||||
}
|
||||
|
||||
// FetchIncremental performs an incremental sync by comparing file modified_time
|
||||
// against the previously recorded state. Defensive fallback path - the service
|
||||
// prefers FetchStream. Routed through the same engine as FetchStream, so the
|
||||
// #2136 failure-doesn't-advance-cursor semantics apply here too.
|
||||
func (c *DriveConnector) FetchIncremental(
|
||||
ctx context.Context, config *types.DataSourceConfig, cursor *types.SyncCursor,
|
||||
) ([]types.FetchedItem, *types.SyncCursor, error) {
|
||||
feishuConfig, err := core.ParseFeishuConfig(config, c.region)
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
client := core.NewClient(feishuConfig)
|
||||
ops := driveOps{region: c.region}
|
||||
if len(config.ResourceIDs) == 0 {
|
||||
return nil, nil, errors.New(ops.EmptyResourceIDsError())
|
||||
}
|
||||
return core.FetchIncrementalEngine(ctx, client, config, cursor, ops)
|
||||
}
|
||||
|
||||
// FetchStream performs a resumable, memory-bounded sync. It unifies the full
|
||||
// and incremental paths: with cursor == nil it fetches everything, and with a
|
||||
// cursor it skips files whose recorded modified_time is unchanged - the same
|
||||
// mechanism that lets a sync which timed out mid-traversal resume from the last
|
||||
// checkpoint instead of restarting (Tencent/WeKnora#2136).
|
||||
//
|
||||
// The per-node loop lives in the shared engine (engine.go); this shell only
|
||||
// wires the Drive NodeOps adapter.
|
||||
func (c *DriveConnector) FetchStream(
|
||||
ctx context.Context, config *types.DataSourceConfig,
|
||||
cursor *types.SyncCursor, h datasource.StreamHandler,
|
||||
) (*types.SyncCursor, error) {
|
||||
feishuConfig, err := core.ParseFeishuConfig(config, c.region)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
client := core.NewClient(feishuConfig)
|
||||
ops := driveOps{region: c.region}
|
||||
if len(config.ResourceIDs) == 0 {
|
||||
return nil, errors.New(ops.EmptyResourceIDsError())
|
||||
}
|
||||
return core.FetchStreamEngine(ctx, client, config, cursor, h, ops)
|
||||
}
|
||||
|
||||
// driveOps adapts the Drive DriveConnector to the generic sync engine. It
|
||||
// carries the region (for channel + URL) and encodes/decodes the Drive cursor
|
||||
// wire format (core.FeishuDriveCursor / file_times) so the engine stays format-agnostic.
|
||||
type driveOps struct {
|
||||
region core.Region
|
||||
}
|
||||
|
||||
func (o driveOps) List(ctx context.Context, client *core.Client, resourceID string) ([]core.DriveFile, error, error) {
|
||||
files, err := listDriveFilesForResource(ctx, client, resourceID)
|
||||
if err == nil {
|
||||
return files, nil, nil
|
||||
}
|
||||
var partial *core.PartialDriveFileListError
|
||||
if errors.As(err, &partial) {
|
||||
return files, err, nil
|
||||
}
|
||||
return files, nil, err
|
||||
}
|
||||
|
||||
func (o driveOps) Token(n core.DriveFile) string { return n.Token }
|
||||
func (o driveOps) Title(n core.DriveFile) string { return n.Name }
|
||||
func (o driveOps) ObjType(n core.DriveFile) string { return n.Type }
|
||||
func (o driveOps) EditTime(n core.DriveFile) string { return n.ModifiedTime }
|
||||
|
||||
func (o driveOps) Fetch(ctx context.Context, client *core.Client, n core.DriveFile, resourceID string, multimodal bool) ([]*types.FetchedItem, error) {
|
||||
return fetchDriveFileContent(ctx, client, n, resourceID, multimodal, o.region)
|
||||
}
|
||||
|
||||
func (o driveOps) ListFailureItems(resourceID string, partial error) []types.FetchedItem {
|
||||
var pe *core.PartialDriveFileListError
|
||||
if errors.As(partial, &pe) {
|
||||
return appendDriveFileListFailureItems(nil, resourceID, o.channel(), pe.Failures)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (o driveOps) channel() string {
|
||||
if o.region.ConnectorType == types.ConnectorTypeLarkDrive {
|
||||
return types.ChannelLarkDrive
|
||||
}
|
||||
return types.ChannelFeishuDrive
|
||||
}
|
||||
|
||||
func (o driveOps) ResourceNoun() string { return "files" }
|
||||
func (o driveOps) EmptyResourceIDsError() string {
|
||||
return "no resource IDs (Drive folder tokens) configured"
|
||||
}
|
||||
func (o driveOps) LogTag() string { return "[FeishuDrive]" }
|
||||
|
||||
func (o driveOps) DecodeCursorTimes(m map[string]interface{}) map[string]map[string]string {
|
||||
var prev core.FeishuDriveCursor
|
||||
b, _ := json.Marshal(m)
|
||||
_ = json.Unmarshal(b, &prev)
|
||||
return prev.FileTimes
|
||||
}
|
||||
|
||||
func (o driveOps) EncodeCursor(times map[string]map[string]string, lastSync time.Time) *types.SyncCursor {
|
||||
fc := core.FeishuDriveCursor{LastSyncTime: lastSync, FileTimes: times}
|
||||
m := make(map[string]interface{})
|
||||
b, _ := json.Marshal(fc)
|
||||
_ = json.Unmarshal(b, &m)
|
||||
return &types.SyncCursor{LastSyncTime: lastSync, ConnectorCursor: m}
|
||||
}
|
||||
|
||||
// fetchDriveFileContent fetches the content of a single Drive file and converts
|
||||
// it to FetchedItems. Dispatches by file.Type, mirroring the wiki
|
||||
// fetchNodeContent. Shortcuts have already been expanded to their target by
|
||||
// ListDriveFilesRecursiveFrom, so this only sees the target type.
|
||||
//
|
||||
// - docx -> blocks API (Markdown) with export fallback; may return attachments/images
|
||||
// - doc/sheet/bitable -> ExportAndDownload -> docx/xlsx
|
||||
// - file -> DownloadDriveFile -> original file
|
||||
// - mindnote/slides/board -> Skip (no API), returns (nil, nil)
|
||||
func fetchDriveFileContent(
|
||||
ctx context.Context, client *core.Client, file core.DriveFile, resourceID string, multimodalEnabled bool, region core.Region,
|
||||
) ([]*types.FetchedItem, error) {
|
||||
if !core.IsSupportedDocType(file.Type) {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
editTime := core.ParseFeishuTimestamp(file.ModifiedTime)
|
||||
// Channel marks the knowledge "source" label. Drive uses its own channel
|
||||
// (feishu_drive / lark_drive) so Drive docs show "飞书云盘" / "Lark 云盘"
|
||||
// distinct from the wiki connector's "飞书".
|
||||
channel := types.ChannelFeishuDrive
|
||||
if region.ConnectorType == types.ConnectorTypeLarkDrive {
|
||||
channel = types.ChannelLarkDrive
|
||||
}
|
||||
baseMeta := map[string]string{
|
||||
"obj_token": file.Token,
|
||||
"obj_type": file.Type,
|
||||
"file_token": file.Token,
|
||||
"folder_token": file.ParentToken,
|
||||
"channel": channel,
|
||||
}
|
||||
|
||||
switch file.Type {
|
||||
case "docx":
|
||||
return core.FetchDocxWithBlocks(ctx, client, core.DocxFetchInput{
|
||||
DocToken: file.Token,
|
||||
ObjToken: file.Token,
|
||||
Title: file.Name,
|
||||
URL: file.URL,
|
||||
ResourceID: resourceID,
|
||||
EditTime: editTime,
|
||||
BaseMeta: baseMeta,
|
||||
MultimodalEnabled: multimodalEnabled,
|
||||
})
|
||||
|
||||
case "doc", "sheet", "bitable":
|
||||
data, fileName, err := client.ExportAndDownload(ctx, file.Token, file.Type)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("export %s (%s): %w", file.Name, file.Type, err)
|
||||
}
|
||||
|
||||
ext := core.ExportFileExtToSuffix[core.ObjTypeToExportFileExtension[file.Type]]
|
||||
if fileName == "" {
|
||||
fileName = core.SanitizeFileName(file.Name) + ext
|
||||
} else if !strings.HasSuffix(strings.ToLower(fileName), ext) {
|
||||
fileName = core.SanitizeFileName(fileName) + ext
|
||||
}
|
||||
|
||||
return []*types.FetchedItem{{
|
||||
ExternalID: file.Token,
|
||||
Title: file.Name,
|
||||
Content: data,
|
||||
ContentType: "application/octet-stream",
|
||||
FileName: fileName,
|
||||
URL: file.URL,
|
||||
UpdatedAt: editTime,
|
||||
SourceResourceID: resourceID,
|
||||
Metadata: baseMeta,
|
||||
}}, nil
|
||||
|
||||
case "file":
|
||||
data, err := client.DownloadDriveFile(ctx, file.Token)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("download file %s (%s): %w", file.Name, file.Token, err)
|
||||
}
|
||||
|
||||
fileName := file.Name
|
||||
if fileName == "" {
|
||||
fileName = file.Token
|
||||
}
|
||||
|
||||
return []*types.FetchedItem{{
|
||||
ExternalID: file.Token,
|
||||
Title: file.Name,
|
||||
Content: data,
|
||||
ContentType: "application/octet-stream",
|
||||
FileName: fileName,
|
||||
URL: file.URL,
|
||||
UpdatedAt: editTime,
|
||||
SourceResourceID: resourceID,
|
||||
Metadata: baseMeta,
|
||||
}}, nil
|
||||
|
||||
default:
|
||||
return nil, nil
|
||||
}
|
||||
}
|
||||
|
||||
// --- Helpers ---
|
||||
|
||||
// makeDriveResourceID encodes a Drive ResourceID: "folderToken" (root) or
|
||||
// "folderToken:fileToken" (child). Reuses core.FeishuWikiNodeResourceSeparator.
|
||||
func makeDriveResourceID(rootFolderToken, fileToken string) string {
|
||||
if fileToken == "" {
|
||||
return rootFolderToken
|
||||
}
|
||||
return rootFolderToken + core.FeishuWikiNodeResourceSeparator + fileToken
|
||||
}
|
||||
|
||||
// parseDriveResourceID splits a Drive resourceID into (rootFolderToken, fileToken).
|
||||
// Mirrors parseWikiResourceID.
|
||||
func parseDriveResourceID(resourceID string) (rootFolderToken, fileToken string) {
|
||||
rootFolderToken, fileToken, _ = strings.Cut(resourceID, core.FeishuWikiNodeResourceSeparator)
|
||||
return rootFolderToken, fileToken
|
||||
}
|
||||
|
||||
// listDriveFilesForResource lists the files to sync for a given resourceID.
|
||||
// A resourceID is either a bare root folderToken (sync the whole subtree) or
|
||||
// "rootFolderToken:fileToken" (sync a single selected file or sub-folder).
|
||||
//
|
||||
// For a single-file selection we cannot pass the fileToken to
|
||||
// ListDriveFilesRecursiveFrom - that API expects a folder and returns 1061002
|
||||
// (params error) for a file token. Instead we walk the root folder subtree (the
|
||||
// file's parent) and filter to just the selected fileToken. This mirrors the
|
||||
// wiki connector, which resolves a single selected node via GetWikiNode; Drive
|
||||
// has no single-file meta API, so filtering the subtree walk is the equivalent.
|
||||
//
|
||||
// A sub-folder selection (fileToken is itself a folder) is handled by walking
|
||||
// that sub-folder's subtree directly - ListDriveFilesRecursiveFrom accepts a
|
||||
// folder token, so no filtering is needed there.
|
||||
func listDriveFilesForResource(
|
||||
ctx context.Context, client *core.Client, resourceID string,
|
||||
) ([]core.DriveFile, error) {
|
||||
rootFolderToken, fileToken := parseDriveResourceID(resourceID)
|
||||
if fileToken == "" {
|
||||
return client.ListDriveFilesRecursiveFrom(ctx, rootFolderToken)
|
||||
}
|
||||
files, err := client.ListDriveFilesRecursiveFrom(ctx, fileToken)
|
||||
if err == nil {
|
||||
return files, nil
|
||||
}
|
||||
if !isDriveNotFolderError(err) {
|
||||
return nil, err
|
||||
}
|
||||
all, walkErr := client.ListDriveFilesRecursiveFrom(ctx, rootFolderToken)
|
||||
if walkErr != nil {
|
||||
var partialErr *core.PartialDriveFileListError
|
||||
if !errors.As(walkErr, &partialErr) {
|
||||
return nil, walkErr
|
||||
}
|
||||
all = filterDriveFileByToken(all, fileToken)
|
||||
if len(all) == 0 {
|
||||
return nil, walkErr
|
||||
}
|
||||
return all, walkErr
|
||||
}
|
||||
return filterDriveFileByToken(all, fileToken), nil
|
||||
}
|
||||
|
||||
// isDriveNotFolderError reports whether err indicates the token was not a
|
||||
// folder (1061002 params error from the list API when a file token is passed).
|
||||
func isDriveNotFolderError(err error) bool {
|
||||
s := strings.ToLower(err.Error())
|
||||
return strings.Contains(s, "1061002") || strings.Contains(s, "params error")
|
||||
}
|
||||
|
||||
// filterDriveFileByToken returns only the entries whose Token matches token.
|
||||
func filterDriveFileByToken(files []core.DriveFile, token string) []core.DriveFile {
|
||||
var out []core.DriveFile
|
||||
for _, f := range files {
|
||||
if f.Token == token {
|
||||
out = append(out, f)
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// driveRootFolderToken extracts the user-supplied root folder_token from the
|
||||
// data source config (ResourceIDs[0]).
|
||||
func driveRootFolderToken(config *types.DataSourceConfig) string {
|
||||
if config == nil || len(config.ResourceIDs) == 0 {
|
||||
return ""
|
||||
}
|
||||
root, _ := parseDriveResourceID(config.ResourceIDs[0])
|
||||
return root
|
||||
}
|
||||
|
||||
// driveFolderToResource builds the root Resource for a Drive folder. The root
|
||||
// folder's name is resolved via GetDriveFolderMeta by the caller (best-effort,
|
||||
// falling back to the token). For sub-folders, use driveFileToResource instead -
|
||||
// the list API returns each child folder's Name.
|
||||
//
|
||||
// The root folder's ExternalID is the bare rootFolderToken (no ":fileToken"
|
||||
// suffix) so it matches the resource_id the user saved in
|
||||
// form.config.resource_ids = [folderToken]. A "token:token" encoding would
|
||||
// break selection matching on edit.
|
||||
func (c *DriveConnector) driveFolderToResource(rootFolderToken, parentToken, folderToken, name string) types.Resource {
|
||||
if name == "" {
|
||||
name = folderToken
|
||||
}
|
||||
return types.Resource{
|
||||
ExternalID: rootFolderToken,
|
||||
Name: name,
|
||||
Type: "drive_folder",
|
||||
URL: c.region.DriveFolderURL(folderToken),
|
||||
HasChildren: true,
|
||||
Metadata: map[string]interface{}{
|
||||
"folder_token": folderToken,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// driveFileToResource converts a core.DriveFile (list result) into a picker Resource.
|
||||
// The ParentID must match the parent folder's ExternalID: the root folder's
|
||||
// ExternalID is the bare rootFolderToken (see driveFolderToResource), while any
|
||||
// sub-folder's ExternalID is "rootFolderToken:folderToken". Direct children of
|
||||
// the root have file.ParentToken == rootFolderToken, so their ParentID is the
|
||||
// bare rootFolderToken; deeper descendants use the encoded form.
|
||||
func (c *DriveConnector) driveFileToResource(rootFolderToken string, file core.DriveFile) types.Resource {
|
||||
name := file.Name
|
||||
if name == "" {
|
||||
name = file.Token
|
||||
}
|
||||
|
||||
modifiedAt := core.ParseFeishuTimestamp(file.ModifiedTime)
|
||||
|
||||
parentID := makeDriveResourceID(rootFolderToken, file.ParentToken)
|
||||
if file.ParentToken == rootFolderToken || file.ParentToken == "" {
|
||||
// Direct child of the root folder: parent is the root, whose
|
||||
// ExternalID is the bare rootFolderToken (no ":token" suffix).
|
||||
parentID = rootFolderToken
|
||||
}
|
||||
|
||||
return types.Resource{
|
||||
ExternalID: makeDriveResourceID(rootFolderToken, file.Token),
|
||||
Name: name,
|
||||
Type: file.Type,
|
||||
URL: file.URL,
|
||||
ParentID: parentID,
|
||||
HasChildren: file.Type == "folder",
|
||||
ModifiedAt: modifiedAt,
|
||||
Metadata: map[string]interface{}{
|
||||
"file_token": file.Token,
|
||||
"obj_type": file.Type,
|
||||
"folder_token": file.ParentToken,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// appendDriveFileListFailureItems converts Drive listing failures into error
|
||||
// FetchedItems so the sync log surfaces which sub-folders could not be listed.
|
||||
// Mirrors appendWikiNodeListFailureItems.
|
||||
func appendDriveFileListFailureItems(items []types.FetchedItem, resourceID, channel string, failures []core.DriveFileListFailure) []types.FetchedItem {
|
||||
for _, failure := range failures {
|
||||
items = append(items, types.FetchedItem{
|
||||
ExternalID: failure.FolderToken,
|
||||
Title: failure.FolderToken,
|
||||
SourceResourceID: resourceID,
|
||||
Metadata: core.FeishuErrorItemMeta(failure.Err, map[string]string{
|
||||
"channel": channel,
|
||||
"folder_token": failure.FolderToken,
|
||||
"failure_stage": "list_children",
|
||||
}),
|
||||
})
|
||||
}
|
||||
return items
|
||||
}
|
||||
@@ -0,0 +1,42 @@
|
||||
package drive
|
||||
|
||||
import (
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/datasource/connector/feishu/core"
|
||||
)
|
||||
|
||||
// TestDriveCursorRoundTrip verifies driveOps.EncodeCursor/DecodeCursorTimes
|
||||
// survive a JSON marshal/unmarshal cycle (snapshot isolation), mirroring the
|
||||
// wiki TestFeishuCursorRoundTrip. Design §3.2: driveOps is new code that
|
||||
// previously had no direct test.
|
||||
func TestDriveCursorRoundTrip(t *testing.T) {
|
||||
ops := driveOps{region: core.RegionFeishuDrive}
|
||||
times := map[string]map[string]string{
|
||||
"folder1": {"fdoc1": "100", "fdoc2": "200"},
|
||||
"folder1:fdoc1": {"fdoc1": "100"},
|
||||
"folder2:subA": {"fdoc3": "300"},
|
||||
}
|
||||
cur := ops.EncodeCursor(times, time.Unix(1000, 0))
|
||||
if cur == nil {
|
||||
t.Fatal("EncodeCursor returned nil")
|
||||
}
|
||||
|
||||
got := ops.DecodeCursorTimes(cur.ConnectorCursor)
|
||||
if len(got) != len(times) {
|
||||
t.Fatalf("decoded %d resources, want %d", len(got), len(times))
|
||||
}
|
||||
for rid, files := range times {
|
||||
gotFiles, ok := got[rid]
|
||||
if !ok {
|
||||
t.Errorf("resource %q missing from decoded cursor", rid)
|
||||
continue
|
||||
}
|
||||
for tok, mt := range files {
|
||||
if gotFiles[tok] != mt {
|
||||
t.Errorf("decoded[%q][%q] = %q, want %q", rid, tok, gotFiles[tok], mt)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,277 @@
|
||||
package drive
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/datasource/connector/feishu/core"
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
)
|
||||
|
||||
// makeDriveConfig builds a DataSourceConfig for the Drive connector, mirroring
|
||||
// makeConfig but with the drive connector type and a multimodal toggle.
|
||||
func makeDriveConfig(cfg *core.Config, resourceIDs []string, multimodal bool) *types.DataSourceConfig {
|
||||
c := makeConfig(cfg, resourceIDs)
|
||||
c.Type = types.ConnectorTypeFeishuDrive
|
||||
c.MultimodalEnabled = multimodal
|
||||
return c
|
||||
}
|
||||
|
||||
// fakeFeishuDriveDocx serves the Drive list endpoint plus the docx blocks /
|
||||
// media / export endpoints. blocksMode controls the blocks API behaviour:
|
||||
// - "ok": serve the given blocks
|
||||
// - "fail": HTTP 500 (missing scope) → exercises the export fallback
|
||||
// - "empty": serve only a page block → empty Markdown → export fallback
|
||||
//
|
||||
// The export trio is always registered so both fallback modes can complete.
|
||||
func fakeFeishuDriveDocx(t *testing.T, files []core.DriveFile, docToken string,
|
||||
blocks []core.DocxBlock, blocksMode string, mediaContent []byte,
|
||||
) (*httptest.Server, *core.Config) {
|
||||
t.Helper()
|
||||
mux := http.NewServeMux()
|
||||
|
||||
mux.HandleFunc("/open-apis/auth/v3/tenant_access_token/internal", func(w http.ResponseWriter, _ *http.Request) {
|
||||
writeJSON(w, core.TokenResponse{ApiResponse: core.ApiResponse{Code: 0}, TenantAccessToken: "fake-token", Expire: 7200})
|
||||
})
|
||||
|
||||
// Drive file listing (single page).
|
||||
mux.HandleFunc("/open-apis/drive/v1/files", func(w http.ResponseWriter, _ *http.Request) {
|
||||
writeJSON(w, core.DriveFileListResponse{
|
||||
ApiResponse: core.ApiResponse{Code: 0},
|
||||
Data: core.DriveFileListData{Files: files},
|
||||
})
|
||||
})
|
||||
|
||||
// Blocks API for the given docx document.
|
||||
blocksPath := "/open-apis/docx/v1/documents/" + docToken + "/blocks"
|
||||
switch blocksMode {
|
||||
case "fail":
|
||||
mux.HandleFunc(blocksPath, func(w http.ResponseWriter, _ *http.Request) {
|
||||
w.WriteHeader(http.StatusInternalServerError)
|
||||
_, _ = w.Write([]byte(`{"code":99991400,"msg":"insufficient scope"}`))
|
||||
})
|
||||
case "empty":
|
||||
mux.HandleFunc(blocksPath, func(w http.ResponseWriter, r *http.Request) {
|
||||
writeJSON(w, core.DocxBlocksResponse{
|
||||
ApiResponse: core.ApiResponse{Code: 0},
|
||||
Data: core.DocxBlocksData{Items: []core.DocxBlock{{BlockID: "b1", BlockType: core.BlockTypePage}}},
|
||||
})
|
||||
})
|
||||
default: // "ok"
|
||||
mux.HandleFunc(blocksPath, func(w http.ResponseWriter, r *http.Request) {
|
||||
writeJSON(w, core.DocxBlocksResponse{
|
||||
ApiResponse: core.ApiResponse{Code: 0},
|
||||
Data: core.DocxBlocksData{Items: blocks},
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
// Media download for File/Image block tokens.
|
||||
mux.HandleFunc("/open-apis/drive/v1/medias/", func(w http.ResponseWriter, r *http.Request) {
|
||||
if strings.HasSuffix(r.URL.Path, "/download") {
|
||||
w.Header().Set("Content-Type", "application/octet-stream")
|
||||
_, _ = w.Write(mediaContent)
|
||||
return
|
||||
}
|
||||
http.NotFound(w, r)
|
||||
})
|
||||
|
||||
// Export trio for the fallback path.
|
||||
mux.HandleFunc("/open-apis/drive/v1/export_tasks", func(w http.ResponseWriter, r *http.Request) {
|
||||
writeJSON(w, core.ExportTaskCreateResponse{
|
||||
ApiResponse: core.ApiResponse{Code: 0},
|
||||
Data: core.ExportTaskCreateData{Ticket: "ticket-drv"},
|
||||
})
|
||||
})
|
||||
mux.HandleFunc("/open-apis/drive/v1/export_tasks/ticket-drv", func(w http.ResponseWriter, r *http.Request) {
|
||||
writeJSON(w, core.ExportTaskStatusResponse{
|
||||
ApiResponse: core.ApiResponse{Code: 0},
|
||||
Data: core.ExportTaskStatusData{
|
||||
Result: core.ExportTaskResult{FileToken: "ft-export-drv", FileSize: 512, JobStatus: 0, FileName: "drive-fallback.docx"},
|
||||
},
|
||||
})
|
||||
})
|
||||
mux.HandleFunc("/open-apis/drive/v1/export_tasks/file/ft-export-drv/download", func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "application/octet-stream")
|
||||
_, _ = w.Write([]byte("fake-drive-export-binary"))
|
||||
})
|
||||
|
||||
ts := httptest.NewServer(mux)
|
||||
t.Cleanup(ts.Close)
|
||||
return ts, &core.Config{AppID: "test-app-id", AppSecret: "test-app-secret", BaseURL: ts.URL}
|
||||
}
|
||||
|
||||
func driveDocxBlocks(attToken, attName string) []core.DocxBlock {
|
||||
return []core.DocxBlock{
|
||||
{BlockID: "b1", BlockType: core.BlockTypePage},
|
||||
{BlockID: "b2", BlockType: core.BlockTypeText, Text: &core.BlockText{
|
||||
Elements: []core.TextElement{{TextRun: &core.TextRun{Content: "Hello drive"}}},
|
||||
}},
|
||||
{BlockID: "b3", BlockType: core.BlockTypeFile, File: &core.BlockFileRef{Token: attToken, Name: attName}},
|
||||
}
|
||||
}
|
||||
|
||||
const driveDocxFileToken = "fdoc1"
|
||||
|
||||
func driveDocxFile() core.DriveFile {
|
||||
return core.DriveFile{
|
||||
Token: driveDocxFileToken,
|
||||
Type: "docx",
|
||||
Name: "Drive Doc",
|
||||
ModifiedTime: "500",
|
||||
ParentToken: "folder1",
|
||||
URL: "https://example.feishu.cn/file/" + driveDocxFileToken,
|
||||
}
|
||||
}
|
||||
|
||||
// A drive docx goes through the blocks path: main Markdown item (external_id =
|
||||
// file token, channel = feishu_drive) plus the attachment sub-item.
|
||||
func TestDriveFetchStream_DocxBlocksMultiItem(t *testing.T) {
|
||||
t.Setenv("FEISHU_DOCX_PARSE_MODE", "blocks")
|
||||
const (
|
||||
attToken = "ft-drv-att"
|
||||
attName = "report.pdf"
|
||||
)
|
||||
attContent := bytes.Repeat([]byte("x"), core.MinAttachmentBytes+1)
|
||||
|
||||
_, cfg := fakeFeishuDriveDocx(t, []core.DriveFile{driveDocxFile()}, driveDocxFileToken,
|
||||
driveDocxBlocks(attToken, attName), "ok", attContent)
|
||||
|
||||
c := NewDriveConnector(core.RegionFeishuDrive)
|
||||
h := &recordingHandler{}
|
||||
_, err := c.FetchStream(context.Background(), makeDriveConfig(cfg, []string{"folder1"}, false), nil, h)
|
||||
if err != nil {
|
||||
t.Fatalf("FetchStream() error: %v", err)
|
||||
}
|
||||
|
||||
if len(h.emitted) != 2 {
|
||||
t.Fatalf("expected 2 emitted items (main doc + attachment), got %d: %+v", len(h.emitted), h.emitted)
|
||||
}
|
||||
|
||||
main := h.emitted[0]
|
||||
if main.ExternalID != driveDocxFileToken {
|
||||
t.Errorf("items[0].ExternalID = %q, want %q", main.ExternalID, driveDocxFileToken)
|
||||
}
|
||||
if main.ContentType != "text/markdown" {
|
||||
t.Errorf("items[0].ContentType = %q, want text/markdown", main.ContentType)
|
||||
}
|
||||
if !main.ReplacesSubtree {
|
||||
t.Errorf("items[0].ReplacesSubtree = false, want true")
|
||||
}
|
||||
if main.Metadata["channel"] != types.ChannelFeishuDrive {
|
||||
t.Errorf("items[0].Metadata[channel] = %q, want %q", main.Metadata["channel"], types.ChannelFeishuDrive)
|
||||
}
|
||||
if main.URL != "https://example.feishu.cn/file/"+driveDocxFileToken {
|
||||
t.Errorf("items[0].URL = %q, want drive file URL passthrough", main.URL)
|
||||
}
|
||||
if !strings.Contains(string(main.Content), "Hello drive") {
|
||||
t.Errorf("items[0].Content missing expected text; got %q", string(main.Content))
|
||||
}
|
||||
|
||||
att := h.emitted[1]
|
||||
wantAttID := driveDocxFileToken + "#file#" + attToken
|
||||
if att.ExternalID != wantAttID {
|
||||
t.Errorf("items[1].ExternalID = %q, want %q", att.ExternalID, wantAttID)
|
||||
}
|
||||
if att.Metadata["attachment"] != "true" {
|
||||
t.Errorf("items[1].Metadata[attachment] = %q, want \"true\"", att.Metadata["attachment"])
|
||||
}
|
||||
}
|
||||
|
||||
// Blocks API failure falls back to export: exactly one octet-stream item, no
|
||||
// ReplacesSubtree (must not sweep good prior children on a transient failure).
|
||||
func TestDriveFetchStream_DocxBlocksFailFallsBackToExport(t *testing.T) {
|
||||
t.Setenv("FEISHU_DOCX_PARSE_MODE", "blocks")
|
||||
_, cfg := fakeFeishuDriveDocx(t, []core.DriveFile{driveDocxFile()}, driveDocxFileToken, nil, "fail", nil)
|
||||
|
||||
c := NewDriveConnector(core.RegionFeishuDrive)
|
||||
h := &recordingHandler{}
|
||||
_, err := c.FetchStream(context.Background(), makeDriveConfig(cfg, []string{"folder1"}, false), nil, h)
|
||||
if err != nil {
|
||||
t.Fatalf("FetchStream() error: %v", err)
|
||||
}
|
||||
|
||||
if len(h.emitted) != 1 {
|
||||
t.Fatalf("expected 1 emitted item (export fallback), got %d: %+v", len(h.emitted), h.emitted)
|
||||
}
|
||||
item := h.emitted[0]
|
||||
if item.ExternalID != driveDocxFileToken {
|
||||
t.Errorf("item.ExternalID = %q, want %q", item.ExternalID, driveDocxFileToken)
|
||||
}
|
||||
if item.ContentType != "application/octet-stream" {
|
||||
t.Errorf("item.ContentType = %q, want application/octet-stream", item.ContentType)
|
||||
}
|
||||
if !strings.HasSuffix(item.FileName, ".docx") {
|
||||
t.Errorf("item.FileName = %q, want .docx suffix", item.FileName)
|
||||
}
|
||||
if item.ReplacesSubtree {
|
||||
t.Error("export-fallback item must not set ReplacesSubtree")
|
||||
}
|
||||
}
|
||||
|
||||
// Blocks rendering to empty Markdown also falls back to export (a blank page
|
||||
// would otherwise ingest as a login-gated URL core.Fetch and fail).
|
||||
func TestDriveFetchStream_DocxBlocksEmptyFallsBackToExport(t *testing.T) {
|
||||
t.Setenv("FEISHU_DOCX_PARSE_MODE", "blocks")
|
||||
_, cfg := fakeFeishuDriveDocx(t, []core.DriveFile{driveDocxFile()}, driveDocxFileToken, nil, "empty", nil)
|
||||
|
||||
c := NewDriveConnector(core.RegionFeishuDrive)
|
||||
h := &recordingHandler{}
|
||||
_, err := c.FetchStream(context.Background(), makeDriveConfig(cfg, []string{"folder1"}, false), nil, h)
|
||||
if err != nil {
|
||||
t.Fatalf("FetchStream() error: %v", err)
|
||||
}
|
||||
|
||||
if len(h.emitted) != 1 {
|
||||
t.Fatalf("expected 1 emitted item (export fallback), got %d: %+v", len(h.emitted), h.emitted)
|
||||
}
|
||||
item := h.emitted[0]
|
||||
if item.ContentType != "application/octet-stream" {
|
||||
t.Errorf("item.ContentType = %q, want application/octet-stream", item.ContentType)
|
||||
}
|
||||
if item.ReplacesSubtree {
|
||||
t.Error("export-fallback item must not set ReplacesSubtree")
|
||||
}
|
||||
}
|
||||
|
||||
// With multimodal disabled, an embedded image is not downloaded or emitted,
|
||||
// but its external_id is still in SubtreeKeep so a later toggle-on does not
|
||||
// sweep it, and toggling VLM off later does not delete previously OCR'd images.
|
||||
func TestDriveFetchStream_DocxImageMultimodalOff(t *testing.T) {
|
||||
t.Setenv("FEISHU_DOCX_PARSE_MODE", "blocks")
|
||||
const imgToken = "img-drv-1"
|
||||
blocks := []core.DocxBlock{
|
||||
{BlockID: "b1", BlockType: core.BlockTypePage},
|
||||
{BlockID: "b2", BlockType: core.BlockTypeText, Text: &core.BlockText{
|
||||
Elements: []core.TextElement{{TextRun: &core.TextRun{Content: "has image"}}},
|
||||
}},
|
||||
{BlockID: "b3", BlockType: core.BlockTypeImage, Image: &core.BlockTokenRef{Token: imgToken}},
|
||||
}
|
||||
_, cfg := fakeFeishuDriveDocx(t, []core.DriveFile{driveDocxFile()}, driveDocxFileToken, blocks, "ok", nil)
|
||||
|
||||
c := NewDriveConnector(core.RegionFeishuDrive)
|
||||
h := &recordingHandler{}
|
||||
_, err := c.FetchStream(context.Background(), makeDriveConfig(cfg, []string{"folder1"}, false), nil, h)
|
||||
if err != nil {
|
||||
t.Fatalf("FetchStream() error: %v", err)
|
||||
}
|
||||
|
||||
if len(h.emitted) != 1 {
|
||||
t.Fatalf("expected 1 emitted item (main only, image core.Skipped), got %d: %+v", len(h.emitted), h.emitted)
|
||||
}
|
||||
main := h.emitted[0]
|
||||
wantKeep := driveDocxFileToken + "#image#" + imgToken
|
||||
found := false
|
||||
for _, k := range main.SubtreeKeep {
|
||||
if k == wantKeep {
|
||||
found = true
|
||||
}
|
||||
}
|
||||
if !found {
|
||||
t.Errorf("SubtreeKeep = %v, want it to contain %q", main.SubtreeKeep, wantKeep)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,61 @@
|
||||
package drive
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"os"
|
||||
"testing"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/datasource/connector/feishu/core"
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
secutils "github.com/Tencent/WeKnora/internal/utils"
|
||||
)
|
||||
|
||||
func TestMain(m *testing.M) {
|
||||
os.Setenv("SSRF_WHITELIST", "127.0.0.1,localhost")
|
||||
secutils.ResetSSRFWhitelistForTest()
|
||||
os.Exit(m.Run())
|
||||
}
|
||||
|
||||
func writeJSON(w http.ResponseWriter, v interface{}) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
json.NewEncoder(w).Encode(v)
|
||||
}
|
||||
|
||||
func makeConfig(cfg *core.Config, resourceIDs []string) *types.DataSourceConfig {
|
||||
creds := map[string]interface{}{
|
||||
"app_id": cfg.AppID,
|
||||
"app_secret": cfg.AppSecret,
|
||||
"base_url": cfg.BaseURL,
|
||||
}
|
||||
return &types.DataSourceConfig{
|
||||
Type: types.ConnectorTypeFeishu,
|
||||
Credentials: creds,
|
||||
ResourceIDs: resourceIDs,
|
||||
}
|
||||
}
|
||||
|
||||
type recordingHandler struct {
|
||||
emitted []types.FetchedItem
|
||||
checkpoints []core.FeishuDriveCursor
|
||||
emitErr func(item types.FetchedItem) error
|
||||
}
|
||||
|
||||
func (h *recordingHandler) Emit(_ context.Context, item types.FetchedItem) error {
|
||||
if h.emitErr != nil {
|
||||
if err := h.emitErr(item); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
h.emitted = append(h.emitted, item)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (h *recordingHandler) Checkpoint(_ context.Context, cursor *types.SyncCursor) error {
|
||||
var fc core.FeishuDriveCursor
|
||||
b, _ := json.Marshal(cursor.ConnectorCursor)
|
||||
_ = json.Unmarshal(b, &fc)
|
||||
h.checkpoints = append(h.checkpoints, fc)
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,467 @@
|
||||
package wiki
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/datasource"
|
||||
"github.com/Tencent/WeKnora/internal/datasource/connector/feishu/core"
|
||||
"github.com/Tencent/WeKnora/internal/logger"
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
)
|
||||
|
||||
// Connector implements the datasource.Connector interface for Feishu and, with
|
||||
// the same code, for Lark: the two clouds expose an identical wiki/docx/drive
|
||||
// API surface. A core.Region picks the cloud — see region.go.
|
||||
type Connector struct {
|
||||
region core.Region
|
||||
}
|
||||
|
||||
// NewConnector creates a connector for the given region (RegionFeishu or RegionLark).
|
||||
func NewConnector(region core.Region) *Connector {
|
||||
return &Connector{region: region}
|
||||
}
|
||||
|
||||
// Feishu supports resumable streaming sync; the service prefers FetchStream over
|
||||
// FetchAll/FetchIncremental when a connector implements StreamingConnector.
|
||||
var _ datasource.StreamingConnector = (*Connector)(nil)
|
||||
|
||||
// Type returns the connector type identifier.
|
||||
func (c *Connector) Type() string {
|
||||
return c.region.ConnectorType
|
||||
}
|
||||
|
||||
// Validate verifies that the Feishu configuration is valid by testing connectivity.
|
||||
func (c *Connector) Validate(ctx context.Context, config *types.DataSourceConfig) error {
|
||||
feishuConfig, err := core.ParseFeishuConfig(config, c.region)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
client := core.NewClient(feishuConfig)
|
||||
if err := client.Ping(ctx); err != nil {
|
||||
return fmt.Errorf("feishu connection failed: %w", err)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// ListResources lists Feishu Wiki resources for selection, loading the tree
|
||||
// lazily one level at a time to avoid traversing the entire wiki up front.
|
||||
//
|
||||
// - parentID == "" → list all accessible wiki spaces.
|
||||
// - parentID == spaceID → list the top-level nodes of that space.
|
||||
// - parentID == "spaceID:nodeToken" → list the direct children of that node.
|
||||
//
|
||||
// Eagerly recursing the whole tree here used to time out for large wikis
|
||||
// (Tencent/WeKnora#1672); the recursive walk now happens only at sync time.
|
||||
func (c *Connector) ListResources(
|
||||
ctx context.Context, config *types.DataSourceConfig, parentID string,
|
||||
) ([]types.Resource, error) {
|
||||
feishuConfig, err := core.ParseFeishuConfig(config, c.region)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
client := core.NewClient(feishuConfig)
|
||||
|
||||
if parentID == "" {
|
||||
spaces, err := client.ListWikiSpaces(ctx)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("list feishu wiki spaces: %w", err)
|
||||
}
|
||||
|
||||
resources := make([]types.Resource, 0, len(spaces))
|
||||
for _, space := range spaces {
|
||||
resources = append(resources, types.Resource{
|
||||
ExternalID: space.SpaceID,
|
||||
Name: space.Name,
|
||||
Type: "wiki_space",
|
||||
Description: space.Description,
|
||||
URL: c.region.WikiURL(space.SpaceID),
|
||||
HasChildren: true,
|
||||
Metadata: map[string]interface{}{
|
||||
"visibility": space.Visibility,
|
||||
"space_id": space.SpaceID,
|
||||
},
|
||||
})
|
||||
}
|
||||
return resources, nil
|
||||
}
|
||||
|
||||
// Lazy load: list only the direct children of the given space / node.
|
||||
spaceID, nodeToken := parseWikiResourceID(parentID)
|
||||
nodes, err := client.ListWikiNodes(ctx, spaceID, nodeToken)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("list feishu wiki nodes under %s: %w", parentID, err)
|
||||
}
|
||||
|
||||
resources := make([]types.Resource, 0, len(nodes))
|
||||
for _, node := range nodes {
|
||||
resources = append(resources, c.wikiNodeToResource(spaceID, node))
|
||||
}
|
||||
return resources, nil
|
||||
}
|
||||
|
||||
// ResolveResourceAncestors returns the resource IDs of every parent that has to
|
||||
// be expanded so the lazily-loaded picker can reveal each given selection. For a
|
||||
// selected node "spaceID:nodeToken" that is its space plus every intermediate
|
||||
// node up the tree; the walk uses GetWikiNode (parent_node_token) and is O(depth)
|
||||
// per selection, so it never re-traverses the whole wiki.
|
||||
func (c *Connector) ResolveResourceAncestors(
|
||||
ctx context.Context, config *types.DataSourceConfig, resourceIDs []string,
|
||||
) ([]string, error) {
|
||||
feishuConfig, err := core.ParseFeishuConfig(config, c.region)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
client := core.NewClient(feishuConfig)
|
||||
|
||||
seen := make(map[string]bool)
|
||||
ancestors := make([]string, 0)
|
||||
add := func(id string) {
|
||||
if id != "" && !seen[id] {
|
||||
seen[id] = true
|
||||
ancestors = append(ancestors, id)
|
||||
}
|
||||
}
|
||||
|
||||
for _, rid := range resourceIDs {
|
||||
spaceID, nodeToken := parseWikiResourceID(rid)
|
||||
if spaceID == "" || nodeToken == "" {
|
||||
// A space-level selection is already a top-level node in the picker;
|
||||
// there is nothing above it to reveal.
|
||||
continue
|
||||
}
|
||||
// The space's direct children must be loaded to reveal the top-level node.
|
||||
add(spaceID)
|
||||
|
||||
// Walk up from the selection to the top, loading each intermediate
|
||||
// parent so the path down to the selection becomes visible.
|
||||
current := nodeToken
|
||||
for current != "" {
|
||||
node, err := client.GetWikiNode(ctx, spaceID, current)
|
||||
if err != nil {
|
||||
// Best-effort: a broken path just stays collapsed, the rest of
|
||||
// the selections are still revealed.
|
||||
logger.Warnf(ctx, "[Feishu] resolve ancestors: get node %s:%s: %v", spaceID, current, err)
|
||||
break
|
||||
}
|
||||
if node.ParentNodeID == "" {
|
||||
break
|
||||
}
|
||||
add(makeWikiNodeResourceID(spaceID, node.ParentNodeID))
|
||||
current = node.ParentNodeID
|
||||
}
|
||||
}
|
||||
|
||||
return ancestors, nil
|
||||
}
|
||||
|
||||
// FetchAll performs a full sync of all documents from the specified wiki spaces.
|
||||
// Defensive fallback path - the service prefers FetchStream when the connector
|
||||
// implements StreamingConnector.
|
||||
func (c *Connector) FetchAll(ctx context.Context, config *types.DataSourceConfig, resourceIDs []string) ([]types.FetchedItem, error) {
|
||||
feishuConfig, err := core.ParseFeishuConfig(config, c.region)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
client := core.NewClient(feishuConfig)
|
||||
return core.FetchAllEngine(ctx, client, config, resourceIDs, wikiOps{region: c.region})
|
||||
}
|
||||
|
||||
// FetchIncremental performs an incremental sync by comparing node edit times
|
||||
// against the previously recorded state. Defensive fallback path - the service
|
||||
// prefers FetchStream. Routed through the same engine, so the #2136
|
||||
// failure-doesn't-advance-cursor semantics apply here too (previously this path
|
||||
// advanced the cursor before fetching, a latent #2136 bug).
|
||||
func (c *Connector) FetchIncremental(ctx context.Context, config *types.DataSourceConfig, cursor *types.SyncCursor) ([]types.FetchedItem, *types.SyncCursor, error) {
|
||||
feishuConfig, err := core.ParseFeishuConfig(config, c.region)
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
client := core.NewClient(feishuConfig)
|
||||
ops := wikiOps{region: c.region}
|
||||
if len(config.ResourceIDs) == 0 {
|
||||
return nil, nil, errors.New(ops.EmptyResourceIDsError())
|
||||
}
|
||||
return core.FetchIncrementalEngine(ctx, client, config, cursor, ops)
|
||||
}
|
||||
|
||||
// FetchStream performs a resumable, memory-bounded sync. It unifies the full
|
||||
// and incremental paths: with cursor == nil it fetches everything, and with a
|
||||
// cursor it skips nodes whose recorded edit time is unchanged - the same
|
||||
// mechanism that lets a sync which timed out mid-traversal resume from the last
|
||||
// checkpoint instead of restarting (Tencent/WeKnora#2136).
|
||||
//
|
||||
// The per-node loop lives in the shared engine (engine.go); this shell only
|
||||
// wires the wiki NodeOps adapter.
|
||||
func (c *Connector) FetchStream(
|
||||
ctx context.Context, config *types.DataSourceConfig,
|
||||
cursor *types.SyncCursor, h datasource.StreamHandler,
|
||||
) (*types.SyncCursor, error) {
|
||||
feishuConfig, err := core.ParseFeishuConfig(config, c.region)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
client := core.NewClient(feishuConfig)
|
||||
ops := wikiOps{region: c.region}
|
||||
if len(config.ResourceIDs) == 0 {
|
||||
return nil, errors.New(ops.EmptyResourceIDsError())
|
||||
}
|
||||
return core.FetchStreamEngine(ctx, client, config, cursor, h, ops)
|
||||
}
|
||||
|
||||
// wikiOps adapts the wiki Connector to the generic sync engine. It carries the
|
||||
// region (for URL rendering) and encodes/decodes the wiki cursor wire format
|
||||
// (core.FeishuCursor / space_node_times) so the engine can stay format-agnostic.
|
||||
type wikiOps struct {
|
||||
region core.Region
|
||||
}
|
||||
|
||||
func (o wikiOps) List(ctx context.Context, client *core.Client, resourceID string) ([]core.WikiNode, error, error) {
|
||||
spaceID, nodeToken := parseWikiResourceID(resourceID)
|
||||
nodes, err := client.ListWikiNodesRecursiveFrom(ctx, spaceID, nodeToken)
|
||||
if err == nil {
|
||||
return nodes, nil, nil
|
||||
}
|
||||
var partial *core.PartialWikiNodeListError
|
||||
if errors.As(err, &partial) {
|
||||
// Partial listing: nodes are still usable; the failed sub-trees are
|
||||
// surfaced via ListFailureItems, and the sync continues.
|
||||
return nodes, err, nil
|
||||
}
|
||||
return nodes, nil, err
|
||||
}
|
||||
|
||||
func (o wikiOps) Token(n core.WikiNode) string { return n.NodeToken }
|
||||
func (o wikiOps) Title(n core.WikiNode) string { return n.Title }
|
||||
func (o wikiOps) ObjType(n core.WikiNode) string { return n.ObjType }
|
||||
|
||||
// EditTime is the change-detection timestamp: ObjEditTime (document content)
|
||||
// with a NodeEditTime fallback for nodes that lack obj_edit_time. This drives
|
||||
// the cursor comparison, NOT FetchedItem.UpdatedAt (which uses NodeEditTime).
|
||||
func (o wikiOps) EditTime(n core.WikiNode) string {
|
||||
if n.ObjEditTime != "" {
|
||||
return n.ObjEditTime
|
||||
}
|
||||
return n.NodeEditTime
|
||||
}
|
||||
|
||||
func (o wikiOps) Fetch(ctx context.Context, client *core.Client, n core.WikiNode, resourceID string, multimodal bool) ([]*types.FetchedItem, error) {
|
||||
spaceID, _ := parseWikiResourceID(resourceID)
|
||||
return fetchNodeContent(ctx, client, n, spaceID, resourceID, multimodal, o.region)
|
||||
}
|
||||
|
||||
func (o wikiOps) ListFailureItems(resourceID string, partial error) []types.FetchedItem {
|
||||
spaceID, _ := parseWikiResourceID(resourceID)
|
||||
var pe *core.PartialWikiNodeListError
|
||||
if errors.As(partial, &pe) {
|
||||
return appendWikiNodeListFailureItems(nil, spaceID, resourceID, pe.Failures)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (o wikiOps) ResourceNoun() string { return "nodes" }
|
||||
func (o wikiOps) EmptyResourceIDsError() string {
|
||||
return "no resource IDs (wiki space IDs or wiki node IDs) configured"
|
||||
}
|
||||
func (o wikiOps) LogTag() string { return "[Feishu]" }
|
||||
|
||||
func (o wikiOps) DecodeCursorTimes(m map[string]interface{}) map[string]map[string]string {
|
||||
var prev core.FeishuCursor
|
||||
b, _ := json.Marshal(m)
|
||||
_ = json.Unmarshal(b, &prev)
|
||||
return prev.SpaceNodeTimes
|
||||
}
|
||||
|
||||
func (o wikiOps) EncodeCursor(times map[string]map[string]string, lastSync time.Time) *types.SyncCursor {
|
||||
fc := core.FeishuCursor{LastSyncTime: lastSync, SpaceNodeTimes: times}
|
||||
m := make(map[string]interface{})
|
||||
b, _ := json.Marshal(fc)
|
||||
_ = json.Unmarshal(b, &m)
|
||||
return &types.SyncCursor{LastSyncTime: lastSync, ConnectorCursor: m}
|
||||
}
|
||||
|
||||
func appendWikiNodeListFailureItems(items []types.FetchedItem, spaceID string, resourceID string, failures []core.WikiNodeListFailure) []types.FetchedItem {
|
||||
for _, failure := range failures {
|
||||
node := failure.Node
|
||||
title := node.Title
|
||||
if title == "" {
|
||||
title = node.NodeToken
|
||||
}
|
||||
items = append(items, types.FetchedItem{
|
||||
ExternalID: node.NodeToken,
|
||||
Title: title,
|
||||
SourceResourceID: resourceID,
|
||||
Metadata: core.FeishuErrorItemMeta(failure.Err, map[string]string{
|
||||
"channel": types.ChannelFeishu,
|
||||
"node_token": node.NodeToken,
|
||||
"space_id": spaceID,
|
||||
"failure_stage": "list_children",
|
||||
}),
|
||||
})
|
||||
}
|
||||
return items
|
||||
}
|
||||
|
||||
// fetchNodeContent fetches the content of a single wiki node and converts it to a
|
||||
// slice of FetchedItems. For docx nodes it fans out into a main Markdown document
|
||||
// plus optional attachment sub-items. Dispatches to different retrieval strategies
|
||||
// based on obj_type:
|
||||
// - docx → blocks API (Markdown) with export fallback; may return attachments
|
||||
// - doc/sheet/bitable → export API → binary file
|
||||
// - file → drive download → original file (PDF/Word/image/etc.)
|
||||
// - mindnote → Skip (no API)
|
||||
// - slides → Skip (no API)
|
||||
func fetchNodeContent(ctx context.Context, client *core.Client, node core.WikiNode, spaceID string, resourceID string, multimodalEnabled bool, region core.Region) ([]*types.FetchedItem, error) {
|
||||
if !core.IsSupportedDocType(node.ObjType) {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
editTime := core.ParseFeishuTimestamp(node.NodeEditTime)
|
||||
baseMeta := map[string]string{
|
||||
"obj_token": node.ObjToken,
|
||||
"obj_type": node.ObjType,
|
||||
"node_token": node.NodeToken,
|
||||
"space_id": spaceID,
|
||||
"creator": node.Creator,
|
||||
"owner": node.Owner,
|
||||
"channel": types.ChannelFeishu,
|
||||
}
|
||||
|
||||
switch node.ObjType {
|
||||
case "docx":
|
||||
return core.FetchDocxWithBlocks(ctx, client, core.DocxFetchInput{
|
||||
DocToken: node.NodeToken,
|
||||
ObjToken: node.ObjToken,
|
||||
Title: node.Title,
|
||||
URL: region.WikiURL(node.NodeToken),
|
||||
ResourceID: resourceID,
|
||||
EditTime: editTime,
|
||||
BaseMeta: baseMeta,
|
||||
MultimodalEnabled: multimodalEnabled,
|
||||
})
|
||||
case "doc", "sheet", "bitable":
|
||||
item, err := fetchViaExport(ctx, client, node, resourceID, editTime, baseMeta, region)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return []*types.FetchedItem{item}, nil
|
||||
case "file":
|
||||
item, err := fetchDriveFile(ctx, client, node, resourceID, editTime, baseMeta, region)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return []*types.FetchedItem{item}, nil
|
||||
default:
|
||||
return nil, nil
|
||||
}
|
||||
}
|
||||
|
||||
// fetchViaExport exports a doc/sheet/bitable node via the async export API and
|
||||
// returns a single FetchedItem containing the exported binary.
|
||||
func fetchViaExport(ctx context.Context, client *core.Client, node core.WikiNode, resourceID string, editTime time.Time, baseMeta map[string]string, region core.Region) (*types.FetchedItem, error) {
|
||||
// Export as a file via the async export API
|
||||
data, fileName, err := client.ExportAndDownload(ctx, node.ObjToken, node.ObjType)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("export %s (%s): %w", node.Title, node.ObjType, err)
|
||||
}
|
||||
|
||||
// Ensure a reasonable file name with correct extension
|
||||
ext := core.ExportFileExtToSuffix[core.ObjTypeToExportFileExtension[node.ObjType]]
|
||||
if fileName == "" {
|
||||
fileName = core.SanitizeFileName(node.Title) + ext
|
||||
} else if !strings.HasSuffix(strings.ToLower(fileName), ext) {
|
||||
// Feishu often returns the doc title without extension - append it
|
||||
fileName = core.SanitizeFileName(fileName) + ext
|
||||
}
|
||||
|
||||
return &types.FetchedItem{
|
||||
ExternalID: node.NodeToken,
|
||||
Title: node.Title,
|
||||
Content: data,
|
||||
ContentType: "application/octet-stream",
|
||||
FileName: fileName,
|
||||
URL: region.WikiURL(node.NodeToken),
|
||||
UpdatedAt: editTime,
|
||||
SourceResourceID: resourceID,
|
||||
Metadata: baseMeta,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// fetchDriveFile downloads an original uploaded file from Drive and returns a
|
||||
// single FetchedItem containing the raw bytes.
|
||||
func fetchDriveFile(ctx context.Context, client *core.Client, node core.WikiNode, resourceID string, editTime time.Time, baseMeta map[string]string, region core.Region) (*types.FetchedItem, error) {
|
||||
// Download the original uploaded file from Drive
|
||||
data, err := client.DownloadDriveFile(ctx, node.ObjToken)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("download file %s (%s): %w", node.Title, node.ObjToken, err)
|
||||
}
|
||||
|
||||
// Use the node title as file name; it usually preserves the original extension
|
||||
fileName := node.Title
|
||||
if fileName == "" {
|
||||
fileName = node.ObjToken
|
||||
}
|
||||
|
||||
return &types.FetchedItem{
|
||||
ExternalID: node.NodeToken,
|
||||
Title: node.Title,
|
||||
Content: data,
|
||||
ContentType: "application/octet-stream",
|
||||
FileName: fileName,
|
||||
URL: region.WikiURL(node.NodeToken),
|
||||
UpdatedAt: editTime,
|
||||
SourceResourceID: resourceID,
|
||||
Metadata: baseMeta,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// --- Helper functions ---
|
||||
|
||||
func makeWikiNodeResourceID(spaceID, nodeToken string) string {
|
||||
return spaceID + core.FeishuWikiNodeResourceSeparator + nodeToken
|
||||
}
|
||||
|
||||
func parseWikiResourceID(resourceID string) (spaceID string, nodeToken string) {
|
||||
spaceID, nodeToken, _ = strings.Cut(resourceID, core.FeishuWikiNodeResourceSeparator)
|
||||
return spaceID, nodeToken
|
||||
}
|
||||
|
||||
func (c *Connector) wikiNodeToResource(spaceID string, node core.WikiNode) types.Resource {
|
||||
parentID := spaceID
|
||||
if node.ParentNodeID != "" {
|
||||
parentID = makeWikiNodeResourceID(spaceID, node.ParentNodeID)
|
||||
}
|
||||
|
||||
name := node.Title
|
||||
if name == "" {
|
||||
name = node.NodeToken
|
||||
}
|
||||
|
||||
modifiedAt := core.ParseFeishuTimestamp(node.ObjEditTime)
|
||||
if modifiedAt.IsZero() {
|
||||
modifiedAt = core.ParseFeishuTimestamp(node.NodeEditTime)
|
||||
}
|
||||
|
||||
return types.Resource{
|
||||
ExternalID: makeWikiNodeResourceID(spaceID, node.NodeToken),
|
||||
Name: name,
|
||||
Type: "wiki_node",
|
||||
URL: c.region.WikiURL(node.NodeToken),
|
||||
ParentID: parentID,
|
||||
HasChildren: node.HasChild,
|
||||
ModifiedAt: modifiedAt,
|
||||
Metadata: map[string]interface{}{
|
||||
"space_id": spaceID,
|
||||
"node_token": node.NodeToken,
|
||||
"obj_token": node.ObjToken,
|
||||
"obj_type": node.ObjType,
|
||||
},
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,267 @@
|
||||
package wiki
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/datasource/connector/feishu/core"
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
)
|
||||
|
||||
// statefulFeishu is a fake Feishu server whose per-node export outcome can be
|
||||
// changed between sync passes. A node whose obj_token is in failTokens has its
|
||||
// export status poll return a failed job, so fetchNodeContent errors for it —
|
||||
// modelling a transient rate-limit / export failure that later recovers.
|
||||
type statefulFeishu struct {
|
||||
mu sync.Mutex
|
||||
failTokens map[string]bool
|
||||
}
|
||||
|
||||
func (s *statefulFeishu) setFail(tokens ...string) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
s.failTokens = map[string]bool{}
|
||||
for _, t := range tokens {
|
||||
s.failTokens[t] = true
|
||||
}
|
||||
}
|
||||
|
||||
func (s *statefulFeishu) shouldFail(token string) bool {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
return s.failTokens[token]
|
||||
}
|
||||
|
||||
func newStatefulFeishu(nodes []core.WikiNode) (*httptest.Server, *core.Config, *statefulFeishu) {
|
||||
s := &statefulFeishu{failTokens: map[string]bool{}}
|
||||
mux := http.NewServeMux()
|
||||
|
||||
mux.HandleFunc("/open-apis/auth/v3/tenant_access_token/internal", func(w http.ResponseWriter, r *http.Request) {
|
||||
writeJSON(w, core.TokenResponse{ApiResponse: core.ApiResponse{Code: 0}, TenantAccessToken: "fake-token", Expire: 7200})
|
||||
})
|
||||
mux.HandleFunc("/open-apis/wiki/v2/spaces", func(w http.ResponseWriter, r *http.Request) {
|
||||
writeJSON(w, core.WikiSpaceListResponse{
|
||||
ApiResponse: core.ApiResponse{Code: 0},
|
||||
Data: core.WikiSpaceListData{Items: []core.WikiSpace{{SpaceID: "space1", Name: "Test Space"}}},
|
||||
})
|
||||
})
|
||||
mux.HandleFunc("/open-apis/wiki/v2/spaces/space1/nodes", func(w http.ResponseWriter, r *http.Request) {
|
||||
writeJSON(w, core.WikiNodeListResponse{
|
||||
ApiResponse: core.ApiResponse{Code: 0},
|
||||
Data: core.WikiNodeListData{Items: nodes},
|
||||
})
|
||||
})
|
||||
// Export create: ticket == obj_token so the poll below can key on it.
|
||||
mux.HandleFunc("/open-apis/drive/v1/export_tasks", func(w http.ResponseWriter, r *http.Request) {
|
||||
var body struct {
|
||||
Token string `json:"token"`
|
||||
}
|
||||
_ = json.NewDecoder(r.Body).Decode(&body)
|
||||
writeJSON(w, core.ExportTaskCreateResponse{
|
||||
ApiResponse: core.ApiResponse{Code: 0},
|
||||
Data: core.ExportTaskCreateData{Ticket: body.Token},
|
||||
})
|
||||
})
|
||||
// Export status poll: /open-apis/drive/v1/export_tasks/<ticket>?token=<objToken>
|
||||
mux.HandleFunc("/open-apis/drive/v1/export_tasks/", func(w http.ResponseWriter, r *http.Request) {
|
||||
if strings.Contains(r.URL.Path, "/file/") && strings.HasSuffix(r.URL.Path, "/download") {
|
||||
w.Header().Set("Content-Type", "application/octet-stream")
|
||||
_, _ = w.Write([]byte("fake-docx-content"))
|
||||
return
|
||||
}
|
||||
token := r.URL.Query().Get("token")
|
||||
status := 0
|
||||
if s.shouldFail(token) {
|
||||
status = 3 // failed job
|
||||
}
|
||||
resp := core.ExportTaskStatusResponse{ApiResponse: core.ApiResponse{Code: 0}}
|
||||
resp.Data.Result.FileToken = "file-" + token
|
||||
resp.Data.Result.FileName = "exported.docx"
|
||||
resp.Data.Result.JobStatus = status
|
||||
resp.Data.Result.JobErrorMsg = "rate limited"
|
||||
writeJSON(w, resp)
|
||||
})
|
||||
|
||||
ts := httptest.NewServer(mux)
|
||||
return ts, &core.Config{AppID: "test-app-id", AppSecret: "test-app-secret", BaseURL: ts.URL}, s
|
||||
}
|
||||
|
||||
// convergenceHandler models the service side across a resumable sync:
|
||||
// - Emit ingests content items (mirroring the service's ctx check + apply),
|
||||
// records failure items separately, and can simulate a mid-tree task timeout
|
||||
// by canceling the run's context after cancelAfterCalls Emit calls.
|
||||
// - Checkpoint snapshots the connector cursor (the DB round-trip).
|
||||
type convergenceHandler struct {
|
||||
ingested []string // ExternalIDs of successfully ingested content items
|
||||
failed []string // ExternalIDs surfaced as failures (Metadata["error"])
|
||||
exportedTokens map[string]int
|
||||
checkpoints []core.FeishuCursor
|
||||
calls int
|
||||
cancelAfterCall int // >0: cancel ctx and abort on the Nth Emit call
|
||||
cancel context.CancelFunc
|
||||
}
|
||||
|
||||
func (h *convergenceHandler) Emit(ctx context.Context, item types.FetchedItem) error {
|
||||
if err := ctx.Err(); err != nil { // mirror service Emit
|
||||
return err
|
||||
}
|
||||
h.calls++
|
||||
if h.cancelAfterCall > 0 && h.calls >= h.cancelAfterCall {
|
||||
if h.cancel != nil {
|
||||
h.cancel()
|
||||
}
|
||||
return context.Canceled // service returns ctx err → connector aborts
|
||||
}
|
||||
if item.Metadata["error"] != "" {
|
||||
h.failed = append(h.failed, item.ExternalID)
|
||||
return nil
|
||||
}
|
||||
h.ingested = append(h.ingested, item.ExternalID)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (h *convergenceHandler) Checkpoint(ctx context.Context, cursor *types.SyncCursor) error {
|
||||
var fc core.FeishuCursor
|
||||
b, _ := json.Marshal(cursor.ConnectorCursor)
|
||||
_ = json.Unmarshal(b, &fc)
|
||||
h.checkpoints = append(h.checkpoints, fc)
|
||||
return nil
|
||||
}
|
||||
|
||||
func lastCheckpointCursor(h *convergenceHandler, t *testing.T) *types.SyncCursor {
|
||||
t.Helper()
|
||||
if len(h.checkpoints) == 0 {
|
||||
return nil
|
||||
}
|
||||
last := h.checkpoints[len(h.checkpoints)-1]
|
||||
b, _ := json.Marshal(last)
|
||||
var m map[string]interface{}
|
||||
if err := json.Unmarshal(b, &m); err != nil {
|
||||
t.Fatalf("rebuild persisted cursor: %v", err)
|
||||
}
|
||||
return &types.SyncCursor{ConnectorCursor: m}
|
||||
}
|
||||
|
||||
// TestFetchStream_ResumeConvergesAfterTimeoutAndTransientFailure is the
|
||||
// end-to-end proof for Tencent/WeKnora#2136: a large-ish wiki that (a) hits a
|
||||
// transient per-node export failure and (b) is killed by the 2h task timeout
|
||||
// mid-traversal must, on the asynq retry, resume from the last Checkpoint,
|
||||
// re-core.Fetch only what is outstanding, retry the transiently-failed node, and end
|
||||
// with EVERY document synced exactly once — no permanent core.Skip, no full restart,
|
||||
// no redundant re-export of already-done nodes.
|
||||
func TestFetchStream_ResumeConvergesAfterTimeoutAndTransientFailure(t *testing.T) {
|
||||
// Checkpoint on every processed node so the persisted cursor is precise.
|
||||
prevN := core.FeishuStreamCheckpointInterval
|
||||
core.FeishuStreamCheckpointInterval = 1
|
||||
defer func() { core.FeishuStreamCheckpointInterval = prevN }()
|
||||
|
||||
nodes := []core.WikiNode{
|
||||
{NodeToken: "nt1", ObjToken: "obj1", ObjType: "docx", Title: "Doc1", ObjEditTime: "100"},
|
||||
{NodeToken: "nt2", ObjToken: "obj2", ObjType: "docx", Title: "Doc2", ObjEditTime: "200"},
|
||||
{NodeToken: "nt3", ObjToken: "obj3", ObjType: "docx", Title: "Doc3", ObjEditTime: "300"},
|
||||
{NodeToken: "nt4", ObjToken: "obj4", ObjType: "docx", Title: "Doc4", ObjEditTime: "400"},
|
||||
{NodeToken: "nt5", ObjToken: "obj5", ObjType: "docx", Title: "Doc5", ObjEditTime: "500"},
|
||||
}
|
||||
ts, cfg, srv := newStatefulFeishu(nodes)
|
||||
defer ts.Close()
|
||||
cfgDS := makeConfig(cfg, []string{"space1"})
|
||||
c := NewConnector(core.RegionFeishu)
|
||||
|
||||
// ---- Pass 1: obj2 export fails transiently; task "times out" on the 3rd
|
||||
// Emit call (nt1 success, nt2 failure-item, then cancel as nt3 is emitted).
|
||||
srv.setFail("obj2")
|
||||
ctx1, cancel1 := context.WithCancel(context.Background())
|
||||
h1 := &convergenceHandler{cancelAfterCall: 3, cancel: cancel1}
|
||||
_, err1 := c.FetchStream(ctx1, cfgDS, nil, h1)
|
||||
cancel1()
|
||||
if err1 == nil {
|
||||
t.Fatalf("pass 1 expected a timeout/abort error, got nil")
|
||||
}
|
||||
|
||||
// Pass 1 ingested only nt1; nt2 surfaced as a failure; nt3+ never reached.
|
||||
if got := strings.Join(h1.ingested, ","); got != "nt1" {
|
||||
t.Fatalf("pass 1 ingested = %q, want just nt1", got)
|
||||
}
|
||||
|
||||
// The persisted cursor (last Checkpoint) must contain nt1 but NOT the
|
||||
// aborted nt3, nor the transiently-failed nt2 — otherwise resume would core.Skip
|
||||
// them forever.
|
||||
persisted := lastCheckpointCursor(h1, t)
|
||||
if persisted == nil {
|
||||
t.Fatalf("pass 1 wrote no Checkpoint — resume would restart from scratch")
|
||||
}
|
||||
var pc core.FeishuCursor
|
||||
pb, _ := json.Marshal(persisted.ConnectorCursor)
|
||||
_ = json.Unmarshal(pb, &pc)
|
||||
p := pc.SpaceNodeTimes["space1"]
|
||||
if _, ok := p["nt1"]; !ok {
|
||||
t.Errorf("persisted cursor missing nt1 (done work lost on resume)")
|
||||
}
|
||||
if _, ok := p["nt2"]; ok {
|
||||
t.Errorf("persisted cursor recorded transiently-failed nt2 — it would be core.Skipped forever")
|
||||
}
|
||||
if _, ok := p["nt3"]; ok {
|
||||
t.Errorf("persisted cursor recorded aborted nt3 — it would be core.Skipped forever")
|
||||
}
|
||||
|
||||
// ---- Pass 2: asynq retry. obj2 has recovered; run to completion resuming
|
||||
// from the persisted cursor.
|
||||
srv.setFail() // clear all failures
|
||||
h2 := &convergenceHandler{}
|
||||
next2, err2 := c.FetchStream(context.Background(), cfgDS, persisted, h2)
|
||||
if err2 != nil {
|
||||
t.Fatalf("pass 2 error: %v", err2)
|
||||
}
|
||||
|
||||
// nt1 was already done → must be core.Skipped (not re-exported, not re-ingested).
|
||||
for _, id := range h2.ingested {
|
||||
if id == "nt1" {
|
||||
t.Errorf("pass 2 re-ingested nt1 — redundant re-export of completed work")
|
||||
}
|
||||
}
|
||||
// Everything still outstanding (nt2 transiently-failed, nt3 aborted, nt4, nt5)
|
||||
// must be ingested in pass 2.
|
||||
want := map[string]bool{"nt2": true, "nt3": true, "nt4": true, "nt5": true}
|
||||
got2 := map[string]bool{}
|
||||
for _, id := range h2.ingested {
|
||||
got2[id] = true
|
||||
}
|
||||
for id := range want {
|
||||
if !got2[id] {
|
||||
t.Errorf("pass 2 did not ingest %s (want all outstanding nodes)", id)
|
||||
}
|
||||
}
|
||||
|
||||
// ---- Convergence: union across both passes covers ALL five nodes exactly.
|
||||
union := map[string]bool{}
|
||||
for _, id := range h1.ingested {
|
||||
union[id] = true
|
||||
}
|
||||
for _, id := range h2.ingested {
|
||||
union[id] = true
|
||||
}
|
||||
for _, n := range nodes {
|
||||
if !union[n.NodeToken] {
|
||||
t.Errorf("node %s was NEVER synced across both passes — data loss (#2136)", n.NodeToken)
|
||||
}
|
||||
}
|
||||
if len(union) != len(nodes) {
|
||||
t.Errorf("synced %d distinct nodes, want %d", len(union), len(nodes))
|
||||
}
|
||||
|
||||
// Final cursor is a complete snapshot of every node (next incremental sync
|
||||
// starts clean).
|
||||
var fc core.FeishuCursor
|
||||
nb, _ := json.Marshal(next2.ConnectorCursor)
|
||||
_ = json.Unmarshal(nb, &fc)
|
||||
for _, n := range nodes {
|
||||
if _, ok := fc.SpaceNodeTimes["space1"][n.NodeToken]; !ok {
|
||||
t.Errorf("final cursor missing %s — incremental sync would re-core.Fetch it", n.NodeToken)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,266 @@
|
||||
package wiki
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/datasource/connector/feishu/core"
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
)
|
||||
|
||||
// ──────────────────────────────────────────────────────────────────────
|
||||
// Golden end-to-end test: a single rich docx node exercising EVERY new
|
||||
// capability at once, driven through the real Connector.FetchAll → real
|
||||
// core.Client → Feishu-shaped fake endpoints (blocks / sheets-v2 / bitable-v1 /
|
||||
// medias). This is the glue test the per-function unit tests can't be:
|
||||
// it proves the whole pipeline produces correct, ordered Markdown plus the
|
||||
// correctly filtered attachment sub-items.
|
||||
// ──────────────────────────────────────────────────────────────────────
|
||||
|
||||
// blk constructors kept local to this file to keep the fixture readable.
|
||||
func sheetBlk(id, token string) core.DocxBlock {
|
||||
return core.DocxBlock{BlockID: id, BlockType: core.BlockTypeSheet, Sheet: &core.BlockTokenRef{Token: token}}
|
||||
}
|
||||
|
||||
func bitableBlk(id, token string) core.DocxBlock {
|
||||
return core.DocxBlock{BlockID: id, BlockType: core.BlockTypeBitable, Bitable: &core.BlockTokenRef{Token: token}}
|
||||
}
|
||||
|
||||
func imageBlk(id, token string) core.DocxBlock {
|
||||
return core.DocxBlock{BlockID: id, BlockType: core.BlockTypeImage, Image: &core.BlockTokenRef{Token: token}}
|
||||
}
|
||||
|
||||
func fileBlk(id, token, name string) core.DocxBlock {
|
||||
return core.DocxBlock{BlockID: id, BlockType: core.BlockTypeFile, File: &core.BlockFileRef{Token: token, Name: name}}
|
||||
}
|
||||
|
||||
// cellBlk builds a table_cell container. A real Feishu table_cell holds no
|
||||
// inline text of its own — its text lives in a child block, so the cell only
|
||||
// references the child id (see cellTextBlk).
|
||||
func cellBlk(id string) core.DocxBlock {
|
||||
return core.DocxBlock{BlockID: id, BlockType: core.BlockTypeTableCell, Children: []string{id + "_txt"}}
|
||||
}
|
||||
|
||||
// cellTextBlk builds the child text block a table cell points at.
|
||||
func cellTextBlk(id, text string) core.DocxBlock {
|
||||
return core.DocxBlock{BlockID: id + "_txt", BlockType: core.BlockTypeText, Text: txt(text)}
|
||||
}
|
||||
|
||||
func tableBlk(id string, cols int, cellIDs ...string) core.DocxBlock {
|
||||
b := core.DocxBlock{BlockID: id, BlockType: core.BlockTypeTable}
|
||||
b.Table = &core.BlockTable{Cells: cellIDs}
|
||||
b.Table.Property = &core.BlockTableProperty{ColumnSize: cols}
|
||||
return b
|
||||
}
|
||||
|
||||
// fakeFeishuGolden serves the full API surface a rich docx node needs.
|
||||
func fakeFeishuGolden(nodes []core.WikiNode, docToken string, blocks []core.DocxBlock,
|
||||
mediaByToken map[string][]byte,
|
||||
) (*httptest.Server, *core.Config) {
|
||||
mux := http.NewServeMux()
|
||||
|
||||
mux.HandleFunc("/open-apis/auth/v3/tenant_access_token/internal", func(w http.ResponseWriter, r *http.Request) {
|
||||
writeJSON(w, core.TokenResponse{ApiResponse: core.ApiResponse{Code: 0}, TenantAccessToken: "fake-token", Expire: 7200})
|
||||
})
|
||||
mux.HandleFunc("/open-apis/wiki/v2/spaces", func(w http.ResponseWriter, r *http.Request) {
|
||||
writeJSON(w, core.WikiSpaceListResponse{ApiResponse: core.ApiResponse{Code: 0}, Data: core.WikiSpaceListData{Items: []core.WikiSpace{{SpaceID: "space1", Name: "Test Space"}}}})
|
||||
})
|
||||
mux.HandleFunc("/open-apis/wiki/v2/spaces/space1/nodes", func(w http.ResponseWriter, r *http.Request) {
|
||||
writeJSON(w, core.WikiNodeListResponse{ApiResponse: core.ApiResponse{Code: 0}, Data: core.WikiNodeListData{Items: nodes}})
|
||||
})
|
||||
|
||||
// docx blocks — paginated across two pages to exercise the paging glue.
|
||||
mux.HandleFunc("/open-apis/docx/v1/documents/"+docToken+"/blocks", func(w http.ResponseWriter, r *http.Request) {
|
||||
half := len(blocks) / 2
|
||||
page := blocks[:half]
|
||||
hasMore := true
|
||||
if r.URL.Query().Get("page_token") == "p2" {
|
||||
page = blocks[half:]
|
||||
hasMore = false
|
||||
}
|
||||
nextTok := "p2"
|
||||
if !hasMore {
|
||||
nextTok = ""
|
||||
}
|
||||
writeJSON(w, core.DocxBlocksResponse{ApiResponse: core.ApiResponse{Code: 0}, Data: core.DocxBlocksData{Items: page, HasMore: hasMore, PageToken: nextTok}})
|
||||
})
|
||||
|
||||
// sheets-v2 values: sht_spread_0 → spreadsheet "sht_spread", sheet "0".
|
||||
// Raw JSON keeps the fake honest to the real wire shape without wrestling
|
||||
// nested anonymous-struct literals.
|
||||
mux.HandleFunc("/open-apis/sheets/v2/spreadsheets/sht_spread/values/0", func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.Write([]byte(`{"code":0,"msg":"","data":{"valueRange":{"values":[["名称","数量"],["苹果",3]]}}}`))
|
||||
})
|
||||
|
||||
// bitable-v1 fields + records: bascApp_tblMain → app "bascApp", table "tblMain".
|
||||
mux.HandleFunc("/open-apis/bitable/v1/apps/bascApp/tables/tblMain/fields", func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.Write([]byte(`{"code":0,"msg":"","data":{"items":[{"field_name":"任务"},{"field_name":"状态"}]}}`))
|
||||
})
|
||||
mux.HandleFunc("/open-apis/bitable/v1/apps/bascApp/tables/tblMain/records/search", func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.Write([]byte(`{"code":0,"msg":"","data":{"has_more":false,"page_token":"","items":[{"fields":{"任务":"写码","状态":"完成"}}]}}`))
|
||||
})
|
||||
|
||||
// medias download, routed by token embedded in the path.
|
||||
mux.HandleFunc("/open-apis/drive/v1/medias/", func(w http.ResponseWriter, r *http.Request) {
|
||||
if !strings.HasSuffix(r.URL.Path, "/download") {
|
||||
http.NotFound(w, r)
|
||||
return
|
||||
}
|
||||
// path: /open-apis/drive/v1/medias/<token>/download
|
||||
p := strings.TrimPrefix(r.URL.Path, "/open-apis/drive/v1/medias/")
|
||||
tok := strings.TrimSuffix(p, "/download")
|
||||
data, ok := mediaByToken[tok]
|
||||
if !ok {
|
||||
http.NotFound(w, r)
|
||||
return
|
||||
}
|
||||
w.Header().Set("Content-Type", "application/octet-stream")
|
||||
w.Write(data)
|
||||
})
|
||||
|
||||
ts := httptest.NewServer(mux)
|
||||
return ts, &core.Config{AppID: "test-app-id", AppSecret: "test-app-secret", BaseURL: ts.URL}
|
||||
}
|
||||
|
||||
func TestGolden_RichDocxAllCapabilities(t *testing.T) {
|
||||
t.Setenv("FEISHU_DOCX_PARSE_MODE", "blocks")
|
||||
const docToken = "obj-golden"
|
||||
bigPDF := bytes.Repeat([]byte("A"), core.MinAttachmentBytes+512) // kept
|
||||
tinyPDF := bytes.Repeat([]byte("B"), 100) // whitelisted but too small → dropped
|
||||
|
||||
blocks := []core.DocxBlock{
|
||||
{BlockID: "root", BlockType: core.BlockTypePage},
|
||||
{BlockID: "h1", BlockType: core.BlockTypeHeading1, Heading1: txt("季度报告")},
|
||||
{BlockID: "p1", BlockType: core.BlockTypeText, Text: txt("本季度概览。")},
|
||||
{BlockID: "h2", BlockType: core.BlockTypeHeading1 + 1, Heading2: txt("关键指标")},
|
||||
{BlockID: "b1", BlockType: core.BlockTypeBullet, Bullet: txt("收入增长")},
|
||||
{BlockID: "o1", BlockType: core.BlockTypeOrdered, Ordered: txt("第一步立项")},
|
||||
{BlockID: "code1", BlockType: core.BlockTypeCode, Code: txt("SELECT 1")},
|
||||
{BlockID: "q1", BlockType: core.BlockTypeQuote, Quote: txt("重要提示")},
|
||||
{BlockID: "todo1", BlockType: core.BlockTypeTodo, Todo: txt("完成复盘")},
|
||||
{BlockID: "call1", BlockType: core.BlockTypeCallout, Callout: txt("注意风险")},
|
||||
{BlockID: "div1", BlockType: core.BlockTypeDivider},
|
||||
tableBlk("tbl1", 2, "c1", "c2", "c3", "c4"),
|
||||
sheetBlk("sh1", "sht_spread_0"),
|
||||
bitableBlk("bt1", "bascApp_tblMain"),
|
||||
imageBlk("im1", "img-tok-SECRET"),
|
||||
fileBlk("fbig", "tok-big", "手册.pdf"),
|
||||
fileBlk("flogo", "tok-logo", "logo.png"), // non-whitelisted ext → no sub-item
|
||||
fileBlk("fsmall", "tok-small", "small.pdf"), // whitelisted but too small → no sub-item
|
||||
// table cells (containers) + their child text blocks — both are core.Skipped by
|
||||
// the main loop and read by the table renderer via the cells' Children.
|
||||
cellBlk("c1"), cellBlk("c2"), cellBlk("c3"), cellBlk("c4"),
|
||||
cellTextBlk("c1", "列A"), cellTextBlk("c2", "列B"),
|
||||
cellTextBlk("c3", "1"), cellTextBlk("c4", "2"),
|
||||
}
|
||||
|
||||
nodes := []core.WikiNode{{
|
||||
NodeToken: "nt-golden", ObjToken: docToken, ObjType: "docx",
|
||||
Title: "季度报告文档", NodeEditTime: "1711468800",
|
||||
}}
|
||||
media := map[string][]byte{"tok-big": bigPDF, "tok-small": tinyPDF}
|
||||
|
||||
ts, cfg := fakeFeishuGolden(nodes, docToken, blocks, media)
|
||||
defer ts.Close()
|
||||
|
||||
c := NewConnector(core.RegionFeishu)
|
||||
items, err := c.FetchAll(context.Background(), makeConfig(cfg, []string{"space1"}), []string{"space1"})
|
||||
if err != nil {
|
||||
t.Fatalf("FetchAll error: %v", err)
|
||||
}
|
||||
|
||||
// ── main markdown item + exactly one attachment sub-item ──
|
||||
if len(items) != 2 {
|
||||
t.Fatalf("want 2 items (main + 1 kept attachment), got %d:\n%+v", len(items), items)
|
||||
}
|
||||
var main, att *types.FetchedItem
|
||||
for i := range items {
|
||||
if items[i].ExternalID == "nt-golden" {
|
||||
main = &items[i]
|
||||
} else {
|
||||
att = &items[i]
|
||||
}
|
||||
}
|
||||
if main == nil {
|
||||
t.Fatal("main markdown item missing")
|
||||
}
|
||||
if main.ContentType != "text/markdown" {
|
||||
t.Errorf("main ContentType = %q, want text/markdown", main.ContentType)
|
||||
}
|
||||
if !main.ReplacesSubtree {
|
||||
t.Error("main item must set ReplacesSubtree for orphan sweeping")
|
||||
}
|
||||
if !strings.HasSuffix(main.FileName, ".md") {
|
||||
t.Errorf("main FileName = %q, want .md", main.FileName)
|
||||
}
|
||||
md := string(main.Content)
|
||||
|
||||
// ── every text construct rendered correctly ──
|
||||
fragments := []string{
|
||||
"# 季度报告",
|
||||
"本季度概览。",
|
||||
"## 关键指标",
|
||||
"- 收入增长",
|
||||
"1. 第一步立项",
|
||||
"```\nSELECT 1\n```",
|
||||
"> 重要提示",
|
||||
"- [ ] 完成复盘",
|
||||
"> 注意风险",
|
||||
"---",
|
||||
"| 列A | 列B |", // native docx table header
|
||||
"| 1 | 2 |", // native docx table row
|
||||
"| 名称 | 数量 |", // embedded sheet header
|
||||
"| 苹果 | 3 |", // embedded sheet row
|
||||
"| 任务 | 状态 |", // embedded bitable header
|
||||
"| 写码 | 完成 |", // embedded bitable row
|
||||
"![图片]()", // token-free image placeholder
|
||||
"📎 附件:手册.pdf",
|
||||
"📎 附件:logo.png",
|
||||
"📎 附件:small.pdf",
|
||||
}
|
||||
for _, f := range fragments {
|
||||
if !strings.Contains(md, f) {
|
||||
t.Errorf("markdown missing fragment %q\n--- full markdown ---\n%s", f, md)
|
||||
}
|
||||
}
|
||||
|
||||
// ── document order preserved (heading before table before attachments) ──
|
||||
order := []string{"# 季度报告", "## 关键指标", "| 列A | 列B |", "| 名称 | 数量 |", "| 任务 | 状态 |", "📎 附件:手册.pdf"}
|
||||
last := -1
|
||||
for _, f := range order {
|
||||
idx := strings.Index(md, f)
|
||||
if idx <= last {
|
||||
t.Errorf("fragment %q out of order (idx=%d, prev=%d)\n%s", f, idx, last, md)
|
||||
}
|
||||
last = idx
|
||||
}
|
||||
|
||||
// ── the internal image media token must NEVER leak into embeddings ──
|
||||
if strings.Contains(md, "img-tok-SECRET") {
|
||||
t.Errorf("internal image token leaked into markdown:\n%s", md)
|
||||
}
|
||||
|
||||
// ── attachment filtering: only the big whitelisted PDF becomes a sub-item ──
|
||||
if att == nil {
|
||||
t.Fatal("kept attachment sub-item missing")
|
||||
}
|
||||
if att.ExternalID != "nt-golden#file#tok-big" {
|
||||
t.Errorf("attachment ExternalID = %q, want nt-golden#file#tok-big", att.ExternalID)
|
||||
}
|
||||
if att.Title != "手册.pdf" {
|
||||
t.Errorf("attachment Title = %q", att.Title)
|
||||
}
|
||||
if !bytes.Equal(att.Content, bigPDF) {
|
||||
t.Errorf("attachment content mismatch: got %d bytes want %d", len(att.Content), len(bigPDF))
|
||||
}
|
||||
if att.Metadata["attachment"] != "true" || att.Metadata["parent_node_token"] != "nt-golden" {
|
||||
t.Errorf("attachment metadata wrong: %+v", att.Metadata)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,234 @@
|
||||
//go:build feishu_integration
|
||||
|
||||
// Real-API integration test for FetchStream against a live Feishu tenant.
|
||||
//
|
||||
// It is OPT-IN and never runs in the normal suite or CI: it is guarded by the
|
||||
// `feishu_integration` build tag AND skips unless the required env vars are set.
|
||||
//
|
||||
// Run it with:
|
||||
//
|
||||
// export CPLUS_INCLUDE_PATH="$(xcrun --show-sdk-path)/usr/include/c++/v1"
|
||||
// FEISHU_APP_ID=... FEISHU_APP_SECRET=... FEISHU_TEST_SPACE_ID=... \
|
||||
// go test -tags feishu_integration -run TestRealAPI -v \
|
||||
// ./internal/datasource/connector/feishu/
|
||||
//
|
||||
// The test is strictly READ-ONLY: it lists and exports documents from the given
|
||||
// wiki space but never creates, edits or deletes anything. It whitelists
|
||||
// *.feishu.cn / *.larksuite.com for SSRF in-process so it works from a dev
|
||||
// machine behind a fake-ip proxy without any production code change.
|
||||
package wiki
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"os"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/datasource/connector/feishu/core"
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
"github.com/Tencent/WeKnora/internal/utils"
|
||||
)
|
||||
|
||||
// collectHandler records everything a real FetchStream emits so the test can
|
||||
// assert coverage and resume behaviour. cancelAfter>0 simulates a task timeout
|
||||
// by cancelling the run after that many successful content emits.
|
||||
type collectHandler struct {
|
||||
ingested []string
|
||||
failed []string
|
||||
checkpoints []*types.SyncCursor
|
||||
calls int
|
||||
cancelAfter int
|
||||
cancel context.CancelFunc
|
||||
}
|
||||
|
||||
func (h *collectHandler) Emit(ctx context.Context, item types.FetchedItem) error {
|
||||
if err := ctx.Err(); err != nil {
|
||||
return err
|
||||
}
|
||||
if item.Metadata["error"] != "" {
|
||||
h.failed = append(h.failed, item.ExternalID)
|
||||
return nil
|
||||
}
|
||||
h.calls++
|
||||
if h.cancelAfter > 0 && h.calls > h.cancelAfter {
|
||||
if h.cancel != nil {
|
||||
h.cancel()
|
||||
}
|
||||
return context.Canceled
|
||||
}
|
||||
h.ingested = append(h.ingested, item.ExternalID)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (h *collectHandler) Checkpoint(_ context.Context, cursor *types.SyncCursor) error {
|
||||
// Snapshot the cursor JSON (mirrors the service serialising synchronously).
|
||||
b, _ := json.Marshal(cursor.ConnectorCursor)
|
||||
var m map[string]interface{}
|
||||
_ = json.Unmarshal(b, &m)
|
||||
h.checkpoints = append(h.checkpoints, &types.SyncCursor{ConnectorCursor: m})
|
||||
return nil
|
||||
}
|
||||
|
||||
func nodeTimes(t *testing.T, cur *types.SyncCursor) map[string]map[string]string {
|
||||
t.Helper()
|
||||
if cur == nil {
|
||||
return nil
|
||||
}
|
||||
var fc core.FeishuCursor
|
||||
b, _ := json.Marshal(cur.ConnectorCursor)
|
||||
_ = json.Unmarshal(b, &fc)
|
||||
return fc.SpaceNodeTimes
|
||||
}
|
||||
|
||||
func countNodes(m map[string]map[string]string) int {
|
||||
n := 0
|
||||
for _, v := range m {
|
||||
n += len(v)
|
||||
}
|
||||
return n
|
||||
}
|
||||
|
||||
// TestRealAPI_ListSpaces discovers which wiki spaces the app can access, so the
|
||||
// space ID for the other tests can be found without guessing. Read-only.
|
||||
func TestRealAPI_ListSpaces(t *testing.T) {
|
||||
appID := os.Getenv("FEISHU_APP_ID")
|
||||
appSecret := os.Getenv("FEISHU_APP_SECRET")
|
||||
if appID == "" || appSecret == "" {
|
||||
t.Skip("set FEISHU_APP_ID, FEISHU_APP_SECRET to run")
|
||||
}
|
||||
utils.SetSSRFWhitelistFromRaw("*.feishu.cn,*.larksuite.com")
|
||||
|
||||
client := core.NewClient(&core.Config{AppID: appID, AppSecret: appSecret, BaseURL: os.Getenv("FEISHU_BASE_URL")})
|
||||
spaces, err := client.ListWikiSpaces(context.Background())
|
||||
if err != nil {
|
||||
t.Fatalf("ListWikiSpaces: %v", err)
|
||||
}
|
||||
if len(spaces) == 0 {
|
||||
t.Fatalf("app can access 0 wiki spaces — add the app as a member of a knowledge base")
|
||||
}
|
||||
for _, s := range spaces {
|
||||
t.Logf("SPACE space_id=%s name=%q", s.SpaceID, s.Name)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRealAPI_FetchStreamResumeConverges(t *testing.T) {
|
||||
appID := os.Getenv("FEISHU_APP_ID")
|
||||
appSecret := os.Getenv("FEISHU_APP_SECRET")
|
||||
spaceID := os.Getenv("FEISHU_TEST_SPACE_ID")
|
||||
if appID == "" || appSecret == "" || spaceID == "" {
|
||||
t.Skip("set FEISHU_APP_ID, FEISHU_APP_SECRET, FEISHU_TEST_SPACE_ID to run")
|
||||
}
|
||||
|
||||
// Whitelist the Feishu/Lark hosts for SSRF so the dev-machine fake-ip proxy
|
||||
// (198.18.x) is allowed. No production code change; scoped to these hosts.
|
||||
utils.SetSSRFWhitelistFromRaw("*.feishu.cn,*.larksuite.com")
|
||||
|
||||
baseURL := os.Getenv("FEISHU_BASE_URL") // optional; defaults to region base
|
||||
cfg := &types.DataSourceConfig{
|
||||
Type: types.ConnectorTypeFeishu,
|
||||
Credentials: map[string]interface{}{
|
||||
"app_id": appID,
|
||||
"app_secret": appSecret,
|
||||
"base_url": baseURL,
|
||||
},
|
||||
ResourceIDs: []string{spaceID},
|
||||
}
|
||||
c := NewConnector(core.RegionFeishu)
|
||||
|
||||
// ---- Pass 1: full sync against the real space.
|
||||
h1 := &collectHandler{}
|
||||
cur1, err := c.FetchStream(context.Background(), cfg, nil, h1)
|
||||
if err != nil {
|
||||
t.Fatalf("pass 1 (full) FetchStream error: %v", err)
|
||||
}
|
||||
n := len(h1.ingested)
|
||||
t.Logf("pass 1: ingested=%d failed=%d cursor_nodes=%d", n, len(h1.failed), countNodes(nodeTimes(t, cur1)))
|
||||
if n == 0 {
|
||||
t.Fatalf("pass 1 ingested 0 documents — create a few docs in the space, or check app scopes (wiki/docx/drive read)")
|
||||
}
|
||||
// Every ingested *node* must be recorded in the cursor. A multi-item node (a
|
||||
// docx that fans out into a main document plus attachment/image sub-items)
|
||||
// also ingests child items whose external_id is "<node>#file#..." or
|
||||
// "<node>#image#...". Those are not wiki nodes: they ride their parent node's
|
||||
// edit-time and intentionally carry no independent cursor entry (the parent's
|
||||
// edit-time gates re-core.Fetch of the whole subtree). Only node-level ids are
|
||||
// required to be present.
|
||||
nt1 := nodeTimes(t, cur1)[spaceID]
|
||||
for _, id := range h1.ingested {
|
||||
if strings.Contains(id, "#") {
|
||||
continue // attachment/image sub-item, not a top-level cursor node
|
||||
}
|
||||
if _, ok := nt1[id]; !ok {
|
||||
t.Errorf("pass 1 cursor missing edit time for ingested node %s", id)
|
||||
}
|
||||
}
|
||||
|
||||
// ---- Pass 2: incremental with pass-1 cursor. Against the REAL API this is
|
||||
// the load-bearing check a fake can't make: real obj_edit_time must be
|
||||
// stable for unchanged docs, so NOTHING is re-ingested.
|
||||
h2 := &collectHandler{}
|
||||
cur2, err := c.FetchStream(context.Background(), cfg, cur1, h2)
|
||||
if err != nil {
|
||||
t.Fatalf("pass 2 (incremental) FetchStream error: %v", err)
|
||||
}
|
||||
t.Logf("pass 2: re-ingested=%d (want 0)", len(h2.ingested))
|
||||
if len(h2.ingested) != 0 {
|
||||
t.Errorf("pass 2 re-ingested %v — real edit-time is unstable or core.Skip logic is broken", h2.ingested)
|
||||
}
|
||||
if got := countNodes(nodeTimes(t, cur2)); got != countNodes(nt1AsMap(nt1)) {
|
||||
t.Logf("note: cursor node count changed between passes (%d→%d) — space may have been edited concurrently", countNodes(nt1AsMap(nt1)), got)
|
||||
}
|
||||
|
||||
// ---- Pass 3: interrupt mid-sync then resume, against the real API.
|
||||
// Only meaningful with at least 2 documents.
|
||||
if n < 2 {
|
||||
t.Logf("only %d doc(s); skipping interrupt/resume convergence sub-test (needs >=2)", n)
|
||||
return
|
||||
}
|
||||
prevN := core.FeishuStreamCheckpointInterval
|
||||
core.FeishuStreamCheckpointInterval = 1 // Checkpoint every node so resume is precise
|
||||
defer func() { core.FeishuStreamCheckpointInterval = prevN }()
|
||||
|
||||
ctx3, cancel3 := context.WithCancel(context.Background())
|
||||
h3 := &collectHandler{cancelAfter: 1, cancel: cancel3} // abort after 1 success
|
||||
_, err = c.FetchStream(ctx3, cfg, nil, h3)
|
||||
cancel3()
|
||||
if err == nil {
|
||||
t.Fatalf("pass 3 expected an abort error from the simulated timeout")
|
||||
}
|
||||
if len(h3.checkpoints) == 0 {
|
||||
t.Fatalf("pass 3 wrote no Checkpoint — resume would restart from scratch")
|
||||
}
|
||||
persisted := h3.checkpoints[len(h3.checkpoints)-1]
|
||||
t.Logf("pass 3: ingested=%d before abort, persisted cursor nodes=%d", len(h3.ingested), countNodes(nodeTimes(t, persisted)))
|
||||
|
||||
// Resume from the persisted Checkpoint; must converge to full coverage.
|
||||
h4 := &collectHandler{}
|
||||
_, err = c.FetchStream(context.Background(), cfg, persisted, h4)
|
||||
if err != nil {
|
||||
t.Fatalf("pass 4 (resume) FetchStream error: %v", err)
|
||||
}
|
||||
union := map[string]bool{}
|
||||
for _, id := range h3.ingested {
|
||||
union[id] = true
|
||||
}
|
||||
for _, id := range h4.ingested {
|
||||
union[id] = true
|
||||
}
|
||||
t.Logf("pass 3+4 union ingested=%d, pass-1 full=%d", len(union), n)
|
||||
if len(union) < n {
|
||||
t.Errorf("resume converged to %d docs but full sync had %d — a document was lost across interrupt+resume", len(union), n)
|
||||
}
|
||||
// Resume must not redo already-checkpointed work.
|
||||
done := nodeTimes(t, persisted)[spaceID]
|
||||
for _, id := range h4.ingested {
|
||||
if _, ok := done[id]; ok {
|
||||
t.Errorf("pass 4 re-ingested already-checkpointed node %s (redundant re-export)", id)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func nt1AsMap(m map[string]string) map[string]map[string]string {
|
||||
return map[string]map[string]string{"_": m}
|
||||
}
|
||||
@@ -0,0 +1,152 @@
|
||||
//go:build feishu_integration
|
||||
|
||||
// Seeder for the real-API integration test. OPT-IN (feishu_integration tag +
|
||||
// env vars) so it never runs in the normal suite or CI.
|
||||
//
|
||||
// It WRITES to your Feishu tenant: it creates a handful of docx documents under
|
||||
// the given wiki space, each titled "WeKnora sync test doc N" so you can find
|
||||
// and delete them afterwards. Requires app write scopes:
|
||||
// - wiki:wiki (create wiki nodes)
|
||||
// - docx:document (write docx content)
|
||||
//
|
||||
// plus the app being a member of the space.
|
||||
//
|
||||
// Run:
|
||||
//
|
||||
// export CPLUS_INCLUDE_PATH="$(xcrun --show-sdk-path)/usr/include/c++/v1"
|
||||
// FEISHU_APP_ID=... FEISHU_APP_SECRET=... FEISHU_TEST_SPACE_ID=... FEISHU_SEED_COUNT=3 \
|
||||
// go test -tags feishu_integration -run TestRealAPI_SeedDocs -v \
|
||||
// ./internal/datasource/connector/feishu/
|
||||
package wiki
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"os"
|
||||
"strconv"
|
||||
"testing"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/datasource/connector/feishu/core"
|
||||
"github.com/Tencent/WeKnora/internal/utils"
|
||||
)
|
||||
|
||||
// wikiNodeCreateResp is the subset of POST /wiki/v2/spaces/:id/nodes we need.
|
||||
type wikiNodeCreateResp struct {
|
||||
Code int `json:"code"`
|
||||
Msg string `json:"msg"`
|
||||
Data struct {
|
||||
Node struct {
|
||||
NodeToken string `json:"node_token"`
|
||||
ObjToken string `json:"obj_token"`
|
||||
ObjType string `json:"obj_type"`
|
||||
Title string `json:"title"`
|
||||
} `json:"node"`
|
||||
} `json:"data"`
|
||||
}
|
||||
|
||||
// TestRealAPI_CreateSpace creates a fresh knowledge base owned by the app, so
|
||||
// the app can seed and read it without a human adding it to an existing space.
|
||||
// It prints the new space_id for use as FEISHU_TEST_SPACE_ID.
|
||||
func TestRealAPI_CreateSpace(t *testing.T) {
|
||||
appID := os.Getenv("FEISHU_APP_ID")
|
||||
appSecret := os.Getenv("FEISHU_APP_SECRET")
|
||||
if appID == "" || appSecret == "" {
|
||||
t.Skip("set FEISHU_APP_ID, FEISHU_APP_SECRET to run")
|
||||
}
|
||||
utils.SetSSRFWhitelistFromRaw("*.feishu.cn,*.larksuite.com")
|
||||
client := core.NewClient(&core.Config{AppID: appID, AppSecret: appSecret, BaseURL: os.Getenv("FEISHU_BASE_URL")})
|
||||
|
||||
var resp struct {
|
||||
Code int `json:"code"`
|
||||
Msg string `json:"msg"`
|
||||
Data struct {
|
||||
Space struct {
|
||||
SpaceID string `json:"space_id"`
|
||||
Name string `json:"name"`
|
||||
} `json:"space"`
|
||||
} `json:"data"`
|
||||
}
|
||||
err := client.DoRequest(context.Background(), "POST", "/open-apis/wiki/v2/spaces",
|
||||
map[string]interface{}{
|
||||
"name": "WeKnora Resilience Test KB",
|
||||
"description": "Auto-created for connector resilience testing (#2136). Safe to delete.",
|
||||
}, &resp)
|
||||
if err != nil {
|
||||
t.Fatalf("create space: %v", err)
|
||||
}
|
||||
if resp.Code != 0 {
|
||||
t.Fatalf("create space: code=%d msg=%s (needs wiki:wiki write scope)", resp.Code, resp.Msg)
|
||||
}
|
||||
t.Logf("CREATED_SPACE space_id=%s name=%q", resp.Data.Space.SpaceID, resp.Data.Space.Name)
|
||||
}
|
||||
|
||||
func TestRealAPI_SeedDocs(t *testing.T) {
|
||||
appID := os.Getenv("FEISHU_APP_ID")
|
||||
appSecret := os.Getenv("FEISHU_APP_SECRET")
|
||||
spaceID := os.Getenv("FEISHU_TEST_SPACE_ID")
|
||||
if appID == "" || appSecret == "" || spaceID == "" {
|
||||
t.Skip("set FEISHU_APP_ID, FEISHU_APP_SECRET, FEISHU_TEST_SPACE_ID to run")
|
||||
}
|
||||
count := 3
|
||||
if v := os.Getenv("FEISHU_SEED_COUNT"); v != "" {
|
||||
if n, err := strconv.Atoi(v); err == nil && n > 0 {
|
||||
count = n
|
||||
}
|
||||
}
|
||||
|
||||
utils.SetSSRFWhitelistFromRaw("*.feishu.cn,*.larksuite.com")
|
||||
|
||||
baseURL := os.Getenv("FEISHU_BASE_URL")
|
||||
client := core.NewClient(&core.Config{AppID: appID, AppSecret: appSecret, BaseURL: baseURL})
|
||||
ctx := context.Background()
|
||||
|
||||
for i := 1; i <= count; i++ {
|
||||
title := fmt.Sprintf("WeKnora sync test doc %d", i)
|
||||
|
||||
// 1) Create a docx wiki node.
|
||||
var created wikiNodeCreateResp
|
||||
err := client.DoRequest(ctx, "POST",
|
||||
fmt.Sprintf("/open-apis/wiki/v2/spaces/%s/nodes", spaceID),
|
||||
map[string]interface{}{
|
||||
"obj_type": "docx",
|
||||
"node_type": "origin",
|
||||
"title": title,
|
||||
}, &created)
|
||||
if err != nil {
|
||||
t.Fatalf("create wiki node %d: %v", i, err)
|
||||
}
|
||||
if created.Code != 0 {
|
||||
t.Fatalf("create wiki node %d: code=%d msg=%s (check wiki:wiki write scope + space membership)", i, created.Code, created.Msg)
|
||||
}
|
||||
docID := created.Data.Node.ObjToken // for docx, obj_token == document_id
|
||||
|
||||
// 2) Add a text block so the export is non-empty.
|
||||
var blockResp struct {
|
||||
Code int `json:"code"`
|
||||
Msg string `json:"msg"`
|
||||
}
|
||||
berr := client.DoRequest(ctx, "POST",
|
||||
fmt.Sprintf("/open-apis/docx/v1/documents/%s/blocks/%s/children", docID, docID),
|
||||
map[string]interface{}{
|
||||
"index": 0,
|
||||
"children": []map[string]interface{}{{
|
||||
"block_type": 2, // text
|
||||
"text": map[string]interface{}{
|
||||
"elements": []map[string]interface{}{{
|
||||
"text_run": map[string]interface{}{
|
||||
"content": fmt.Sprintf("This is WeKnora sync test document number %d. It exists to validate connector resilience (issue #2136).", i),
|
||||
},
|
||||
}},
|
||||
},
|
||||
}},
|
||||
}, &blockResp)
|
||||
if berr != nil {
|
||||
t.Logf("warn: add content to doc %d failed (%v) — doc still created, export may be empty", i, berr)
|
||||
} else if blockResp.Code != 0 {
|
||||
t.Logf("warn: add content to doc %d: code=%d msg=%s (needs docx:document write scope) — doc still created", i, blockResp.Code, blockResp.Msg)
|
||||
}
|
||||
|
||||
t.Logf("created doc %d: title=%q node_token=%s obj_token=%s", i, title, created.Data.Node.NodeToken, docID)
|
||||
}
|
||||
t.Logf("seeded %d docs under space %s. Now run TestRealAPI_FetchStreamResumeConverges.", count, spaceID)
|
||||
}
|
||||
@@ -0,0 +1,490 @@
|
||||
package wiki
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/datasource/connector/feishu/core"
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
)
|
||||
|
||||
// fakeFeishuFailingExport serves the auth/spaces/nodes endpoints normally but
|
||||
// fails every document export (code != 0), so fetchNodeContent returns an error
|
||||
// for each supported node — modelling a rate-limited / broken export.
|
||||
func fakeFeishuFailingExport(nodes []core.WikiNode) (*httptest.Server, *core.Config) {
|
||||
mux := http.NewServeMux()
|
||||
mux.HandleFunc("/open-apis/auth/v3/tenant_access_token/internal", func(w http.ResponseWriter, r *http.Request) {
|
||||
writeJSON(w, core.TokenResponse{ApiResponse: core.ApiResponse{Code: 0}, TenantAccessToken: "fake-token", Expire: 7200})
|
||||
})
|
||||
mux.HandleFunc("/open-apis/wiki/v2/spaces", func(w http.ResponseWriter, r *http.Request) {
|
||||
writeJSON(w, core.WikiSpaceListResponse{
|
||||
ApiResponse: core.ApiResponse{Code: 0},
|
||||
Data: core.WikiSpaceListData{Items: []core.WikiSpace{{SpaceID: "space1", Name: "Test Space"}}},
|
||||
})
|
||||
})
|
||||
mux.HandleFunc("/open-apis/wiki/v2/spaces/space1/nodes", func(w http.ResponseWriter, r *http.Request) {
|
||||
writeJSON(w, core.WikiNodeListResponse{
|
||||
ApiResponse: core.ApiResponse{Code: 0},
|
||||
Data: core.WikiNodeListData{Items: nodes},
|
||||
})
|
||||
})
|
||||
// Export creation fails for every document.
|
||||
mux.HandleFunc("/open-apis/drive/v1/export_tasks", func(w http.ResponseWriter, r *http.Request) {
|
||||
writeJSON(w, core.ApiResponse{Code: 1, Msg: "export unavailable"})
|
||||
})
|
||||
ts := httptest.NewServer(mux)
|
||||
return ts, &core.Config{AppID: "test-app-id", AppSecret: "test-app-secret", BaseURL: ts.URL}
|
||||
}
|
||||
|
||||
// A node whose core.Fetch fails must NOT have its new edit time recorded in the
|
||||
// returned cursor: recording it would make the next sync's unchanged fast-path
|
||||
// core.Skip it forever, silently dropping a document on a transient export failure
|
||||
// (Tencent/WeKnora#2136). With a prior edit time known, the prior value is
|
||||
// retained so prev != current next run and the node is retried.
|
||||
func TestFetchStream_FailedFetchRetainsPriorCursor(t *testing.T) {
|
||||
nodes := []core.WikiNode{{NodeToken: "nt1", ObjToken: "obj1", ObjType: "docx", Title: "Doc", ObjEditTime: "100"}}
|
||||
ts, cfg := fakeFeishuFailingExport(nodes)
|
||||
defer ts.Close()
|
||||
|
||||
cursor := makeStreamCursor(t, map[string]map[string]string{"space1": {"nt1": "50"}}) // prior, older
|
||||
|
||||
c := NewConnector(core.RegionFeishu)
|
||||
h := &recordingHandler{}
|
||||
next, err := c.FetchStream(context.Background(), makeConfig(cfg, []string{"space1"}), cursor, h)
|
||||
if err != nil {
|
||||
t.Fatalf("FetchStream() error: %v", err)
|
||||
}
|
||||
|
||||
// A failure item must be surfaced, not silently dropped.
|
||||
if len(h.emitted) != 1 || h.emitted[0].Metadata["error"] == "" {
|
||||
t.Fatalf("expected 1 emitted failure item with error metadata, got %+v", h.emitted)
|
||||
}
|
||||
|
||||
var fc core.FeishuCursor
|
||||
b, _ := json.Marshal(next.ConnectorCursor)
|
||||
_ = json.Unmarshal(b, &fc)
|
||||
got := fc.SpaceNodeTimes["space1"]["nt1"]
|
||||
if got == "100" {
|
||||
t.Fatalf("failed node advanced to current edit time %q — it will be core.Skipped forever", got)
|
||||
}
|
||||
if got != "50" {
|
||||
t.Errorf("failed node cursor = %q, want prior value \"50\" (retry next run)", got)
|
||||
}
|
||||
}
|
||||
|
||||
// With no prior cursor entry, a failed core.Fetch must leave the node out of the
|
||||
// returned cursor entirely, so the next run treats it as new and retries it.
|
||||
func TestFetchStream_FailedFetchNoPriorOmitsFromCursor(t *testing.T) {
|
||||
nodes := []core.WikiNode{{NodeToken: "nt1", ObjToken: "obj1", ObjType: "docx", Title: "Doc", ObjEditTime: "100"}}
|
||||
ts, cfg := fakeFeishuFailingExport(nodes)
|
||||
defer ts.Close()
|
||||
|
||||
c := NewConnector(core.RegionFeishu)
|
||||
h := &recordingHandler{}
|
||||
next, err := c.FetchStream(context.Background(), makeConfig(cfg, []string{"space1"}), nil, h)
|
||||
if err != nil {
|
||||
t.Fatalf("FetchStream() error: %v", err)
|
||||
}
|
||||
|
||||
var fc core.FeishuCursor
|
||||
b, _ := json.Marshal(next.ConnectorCursor)
|
||||
_ = json.Unmarshal(b, &fc)
|
||||
if v, ok := fc.SpaceNodeTimes["space1"]["nt1"]; ok {
|
||||
t.Errorf("failed node recorded in cursor as %q; want absent so it is retried next run", v)
|
||||
}
|
||||
}
|
||||
|
||||
// recordingHandler captures the items and checkpoints a streaming core.Fetch emits.
|
||||
// Checkpoints are snapshotted (JSON-encoded) at call time — mirroring the
|
||||
// service, which serializes the cursor synchronously inside Checkpoint — so the
|
||||
// test observes the cursor state as it was when Checkpoint was called, not the
|
||||
// connector's later-mutated map.
|
||||
type recordingHandler struct {
|
||||
emitted []types.FetchedItem
|
||||
checkpoints []core.FeishuCursor
|
||||
emitErr func(item types.FetchedItem) error
|
||||
}
|
||||
|
||||
func (h *recordingHandler) Emit(ctx context.Context, item types.FetchedItem) error {
|
||||
if h.emitErr != nil {
|
||||
if err := h.emitErr(item); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
h.emitted = append(h.emitted, item)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (h *recordingHandler) Checkpoint(ctx context.Context, cursor *types.SyncCursor) error {
|
||||
var fc core.FeishuCursor
|
||||
b, _ := json.Marshal(cursor.ConnectorCursor)
|
||||
_ = json.Unmarshal(b, &fc)
|
||||
h.checkpoints = append(h.checkpoints, fc)
|
||||
return nil
|
||||
}
|
||||
|
||||
func makeStreamCursor(t *testing.T, spaceNodeTimes map[string]map[string]string) *types.SyncCursor {
|
||||
t.Helper()
|
||||
prev := core.FeishuCursor{SpaceNodeTimes: spaceNodeTimes}
|
||||
b, _ := json.Marshal(prev)
|
||||
var m map[string]interface{}
|
||||
if err := json.Unmarshal(b, &m); err != nil {
|
||||
t.Fatalf("build cursor: %v", err)
|
||||
}
|
||||
return &types.SyncCursor{ConnectorCursor: m}
|
||||
}
|
||||
|
||||
// A nil-cursor stream emits every supported node once and skips unsupported
|
||||
// types, and the returned cursor records the edit time of every discovered node
|
||||
// (so the next incremental run can detect changes).
|
||||
func TestFetchStream_EmitsSupportedSkipsUnsupported(t *testing.T) {
|
||||
nodes := []core.WikiNode{
|
||||
{NodeToken: "nt1", ObjToken: "obj1", ObjType: "docx", Title: "Doc", ObjEditTime: "100"},
|
||||
{NodeToken: "nt2", ObjToken: "obj2", ObjType: "mindnote", Title: "Brain", ObjEditTime: "200"},
|
||||
{NodeToken: "nt3", ObjToken: "obj3", ObjType: "docx", Title: "Doc3", ObjEditTime: "300"},
|
||||
}
|
||||
ts, cfg := fakeFeishu(nodes)
|
||||
defer ts.Close()
|
||||
|
||||
c := NewConnector(core.RegionFeishu)
|
||||
h := &recordingHandler{}
|
||||
next, err := c.FetchStream(context.Background(), makeConfig(cfg, []string{"space1"}), nil, h)
|
||||
if err != nil {
|
||||
t.Fatalf("FetchStream() error: %v", err)
|
||||
}
|
||||
|
||||
if len(h.emitted) != 2 {
|
||||
t.Fatalf("emitted %d items, want 2 (nt1, nt3)", len(h.emitted))
|
||||
}
|
||||
if h.emitted[0].ExternalID != "nt1" || h.emitted[1].ExternalID != "nt3" {
|
||||
t.Errorf("emitted ids = %q,%q; want nt1,nt3", h.emitted[0].ExternalID, h.emitted[1].ExternalID)
|
||||
}
|
||||
|
||||
var fc core.FeishuCursor
|
||||
b, _ := json.Marshal(next.ConnectorCursor)
|
||||
_ = json.Unmarshal(b, &fc)
|
||||
times := fc.SpaceNodeTimes["space1"]
|
||||
for _, tok := range []string{"nt1", "nt2", "nt3"} {
|
||||
if _, ok := times[tok]; !ok {
|
||||
t.Errorf("returned cursor missing edit time for %s", tok)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// When a cursor already records a node at its current edit time, that node is
|
||||
// unchanged and must not be re-emitted; only new/changed nodes stream through.
|
||||
// This is the resume/incremental-core.Skip behavior that lets a timed-out sync
|
||||
// converge across retries instead of re-exporting everything.
|
||||
func TestFetchStream_SkipsUnchangedNodesFromCursor(t *testing.T) {
|
||||
nodes := []core.WikiNode{
|
||||
{NodeToken: "nt1", ObjToken: "obj1", ObjType: "docx", Title: "Doc", ObjEditTime: "100"},
|
||||
{NodeToken: "nt3", ObjToken: "obj3", ObjType: "docx", Title: "Doc3", ObjEditTime: "300"},
|
||||
}
|
||||
ts, cfg := fakeFeishu(nodes)
|
||||
defer ts.Close()
|
||||
|
||||
cursor := makeStreamCursor(t, map[string]map[string]string{
|
||||
"space1": {"nt1": "100"}, // nt1 unchanged; nt3 unknown → new
|
||||
})
|
||||
|
||||
c := NewConnector(core.RegionFeishu)
|
||||
h := &recordingHandler{}
|
||||
if _, err := c.FetchStream(context.Background(), makeConfig(cfg, []string{"space1"}), cursor, h); err != nil {
|
||||
t.Fatalf("FetchStream() error: %v", err)
|
||||
}
|
||||
|
||||
if len(h.emitted) != 1 {
|
||||
t.Fatalf("emitted %d items, want 1 (only changed nt3)", len(h.emitted))
|
||||
}
|
||||
if h.emitted[0].ExternalID != "nt3" {
|
||||
t.Errorf("emitted id = %q, want nt3", h.emitted[0].ExternalID)
|
||||
}
|
||||
}
|
||||
|
||||
// Checkpoints must persist progress at page boundaries so a crash mid-sync
|
||||
// resumes from the last Checkpoint. With the interval set to 1, each emitted
|
||||
// item triggers a Checkpoint, and the first Checkpoint must already contain the
|
||||
// first node's edit time.
|
||||
func TestFetchStream_CheckpointsProgress(t *testing.T) {
|
||||
prev := core.FeishuStreamCheckpointInterval
|
||||
core.FeishuStreamCheckpointInterval = 1
|
||||
defer func() { core.FeishuStreamCheckpointInterval = prev }()
|
||||
|
||||
nodes := []core.WikiNode{
|
||||
{NodeToken: "nt1", ObjToken: "obj1", ObjType: "docx", Title: "Doc", ObjEditTime: "100"},
|
||||
{NodeToken: "nt3", ObjToken: "obj3", ObjType: "docx", Title: "Doc3", ObjEditTime: "300"},
|
||||
}
|
||||
ts, cfg := fakeFeishu(nodes)
|
||||
defer ts.Close()
|
||||
|
||||
c := NewConnector(core.RegionFeishu)
|
||||
h := &recordingHandler{}
|
||||
if _, err := c.FetchStream(context.Background(), makeConfig(cfg, []string{"space1"}), nil, h); err != nil {
|
||||
t.Fatalf("FetchStream() error: %v", err)
|
||||
}
|
||||
|
||||
if len(h.checkpoints) == 0 {
|
||||
t.Fatalf("expected at least one Checkpoint")
|
||||
}
|
||||
first := h.checkpoints[0]
|
||||
if _, ok := first.SpaceNodeTimes["space1"]["nt1"]; !ok {
|
||||
t.Errorf("first Checkpoint missing nt1 progress: %+v", first.SpaceNodeTimes)
|
||||
}
|
||||
}
|
||||
|
||||
// Checkpoints must ALSO fire on elapsed time, not only every N nodes.
|
||||
// Otherwise a sync with fewer than the node interval of slow (rate-limited)
|
||||
// exports reaches the 2h task timeout having never checkpointed, and resumes
|
||||
// from scratch forever — exactly the #2136 "never fully syncs" case. With the
|
||||
// node interval effectively disabled and the time interval at 0, every
|
||||
// processed node must still produce a Checkpoint.
|
||||
func TestFetchStream_CheckpointsOnElapsedTime(t *testing.T) {
|
||||
prevN := core.FeishuStreamCheckpointInterval
|
||||
prevT := core.FeishuStreamCheckpointMaxInterval
|
||||
core.FeishuStreamCheckpointInterval = 1 << 30 // never fires by count
|
||||
core.FeishuStreamCheckpointMaxInterval = 0 // fires by elapsed time every node
|
||||
defer func() {
|
||||
core.FeishuStreamCheckpointInterval = prevN
|
||||
core.FeishuStreamCheckpointMaxInterval = prevT
|
||||
}()
|
||||
|
||||
nodes := []core.WikiNode{
|
||||
{NodeToken: "nt1", ObjToken: "obj1", ObjType: "docx", Title: "Doc", ObjEditTime: "100"},
|
||||
{NodeToken: "nt3", ObjToken: "obj3", ObjType: "docx", Title: "Doc3", ObjEditTime: "300"},
|
||||
}
|
||||
ts, cfg := fakeFeishu(nodes)
|
||||
defer ts.Close()
|
||||
|
||||
c := NewConnector(core.RegionFeishu)
|
||||
h := &recordingHandler{}
|
||||
if _, err := c.FetchStream(context.Background(), makeConfig(cfg, []string{"space1"}), nil, h); err != nil {
|
||||
t.Fatalf("FetchStream() error: %v", err)
|
||||
}
|
||||
if len(h.checkpoints) == 0 {
|
||||
t.Fatalf("expected time-based checkpoints even though node interval never fires")
|
||||
}
|
||||
}
|
||||
|
||||
// An Emit error aborts the stream immediately — the connector must return the
|
||||
// error and stop fetching further nodes (the sync is failing; do not burn API
|
||||
// budget on the rest of the tree).
|
||||
func TestFetchStream_EmitErrorAborts(t *testing.T) {
|
||||
nodes := []core.WikiNode{
|
||||
{NodeToken: "nt1", ObjToken: "obj1", ObjType: "docx", Title: "Doc", ObjEditTime: "100"},
|
||||
{NodeToken: "nt3", ObjToken: "obj3", ObjType: "docx", Title: "Doc3", ObjEditTime: "300"},
|
||||
}
|
||||
ts, cfg := fakeFeishu(nodes)
|
||||
defer ts.Close()
|
||||
|
||||
boom := errors.New("ingest failed")
|
||||
c := NewConnector(core.RegionFeishu)
|
||||
h := &recordingHandler{emitErr: func(item types.FetchedItem) error { return boom }}
|
||||
_, err := c.FetchStream(context.Background(), makeConfig(cfg, []string{"space1"}), nil, h)
|
||||
if !errors.Is(err, boom) {
|
||||
t.Fatalf("FetchStream() error = %v, want %v", err, boom)
|
||||
}
|
||||
if len(h.emitted) != 0 {
|
||||
t.Errorf("emitted %d items, want 0 (aborted on first Emit)", len(h.emitted))
|
||||
}
|
||||
}
|
||||
|
||||
// ──────────────────────────────────────────────────────────────────────
|
||||
// Stream integration: docx blocks fan-out (multi-item) through FetchStream
|
||||
// ──────────────────────────────────────────────────────────────────────
|
||||
|
||||
// TestFetchStream_DocxMultiItem is a stream-level integration test proving that a
|
||||
// docx node fans out to a main Markdown item + an attachment item through the real
|
||||
// FetchStream → fetchNodeContent → core.FetchDocxWithBlocks path. The fake server
|
||||
// serves the blocks API (one text block + one file block) and the drive download;
|
||||
// no export endpoint is registered, so any accidental fall-through to the export
|
||||
// path would 404 and surface as a core.Fetch error.
|
||||
func TestFetchStream_DocxMultiItem(t *testing.T) {
|
||||
t.Setenv("FEISHU_DOCX_PARSE_MODE", "blocks")
|
||||
const (
|
||||
nodeToken = "nt-blocks"
|
||||
objToken = "obj-blocks"
|
||||
attToken = "ft-stream-att"
|
||||
attName = "slides.pdf"
|
||||
)
|
||||
// Attachment content must exceed core.MinAttachmentBytes (2 KiB).
|
||||
attContent := bytes.Repeat([]byte("a"), core.MinAttachmentBytes+1)
|
||||
|
||||
nodes := []core.WikiNode{{
|
||||
NodeToken: nodeToken,
|
||||
ObjToken: objToken,
|
||||
ObjType: "docx",
|
||||
Title: "Stream Blocks Doc",
|
||||
ObjEditTime: "500",
|
||||
}}
|
||||
ts, cfg := fakeFeishuWithBlocks(nodes, objToken, attToken, attName, attContent)
|
||||
defer ts.Close()
|
||||
|
||||
c := NewConnector(core.RegionFeishu)
|
||||
h := &recordingHandler{}
|
||||
_, err := c.FetchStream(context.Background(), makeConfig(cfg, []string{"space1"}), nil, h)
|
||||
if err != nil {
|
||||
t.Fatalf("FetchStream() error: %v", err)
|
||||
}
|
||||
|
||||
if len(h.emitted) != 2 {
|
||||
t.Fatalf("expected 2 emitted items (main doc + attachment), got %d: %+v", len(h.emitted), h.emitted)
|
||||
}
|
||||
|
||||
main := h.emitted[0]
|
||||
if main.ExternalID != nodeToken {
|
||||
t.Errorf("items[0].ExternalID = %q, want %q", main.ExternalID, nodeToken)
|
||||
}
|
||||
if main.ContentType != "text/markdown" {
|
||||
t.Errorf("items[0].ContentType = %q, want text/markdown", main.ContentType)
|
||||
}
|
||||
if !main.ReplacesSubtree {
|
||||
t.Errorf("items[0].ReplacesSubtree = false, want true")
|
||||
}
|
||||
|
||||
att := h.emitted[1]
|
||||
wantAttID := nodeToken + "#file#" + attToken
|
||||
if att.ExternalID != wantAttID {
|
||||
t.Errorf("items[1].ExternalID = %q, want %q", att.ExternalID, wantAttID)
|
||||
}
|
||||
if att.Metadata["attachment"] != "true" {
|
||||
t.Errorf("items[1].Metadata[attachment] = %q, want \"true\"", att.Metadata["attachment"])
|
||||
}
|
||||
}
|
||||
|
||||
// fakeFeishuWithBlocksFallback returns a server where the blocks API returns HTTP 500
|
||||
// (simulating a permission/scope error) and the export task path returns success.
|
||||
// This wires the fallback path: core.FetchDocxWithBlocks → blocks error → fetchViaExport.
|
||||
func fakeFeishuWithBlocksFallback(nodes []core.WikiNode, docToken string) (*httptest.Server, *core.Config) {
|
||||
mux := http.NewServeMux()
|
||||
|
||||
mux.HandleFunc("/open-apis/auth/v3/tenant_access_token/internal", func(w http.ResponseWriter, r *http.Request) {
|
||||
writeJSON(w, core.TokenResponse{
|
||||
ApiResponse: core.ApiResponse{Code: 0},
|
||||
TenantAccessToken: "fake-token",
|
||||
Expire: 7200,
|
||||
})
|
||||
})
|
||||
mux.HandleFunc("/open-apis/wiki/v2/spaces", func(w http.ResponseWriter, r *http.Request) {
|
||||
writeJSON(w, core.WikiSpaceListResponse{
|
||||
ApiResponse: core.ApiResponse{Code: 0},
|
||||
Data: core.WikiSpaceListData{Items: []core.WikiSpace{{SpaceID: "space1", Name: "Test Space"}}},
|
||||
})
|
||||
})
|
||||
mux.HandleFunc("/open-apis/wiki/v2/spaces/space1/nodes", func(w http.ResponseWriter, r *http.Request) {
|
||||
writeJSON(w, core.WikiNodeListResponse{
|
||||
ApiResponse: core.ApiResponse{Code: 0},
|
||||
Data: core.WikiNodeListData{Items: nodes},
|
||||
})
|
||||
})
|
||||
|
||||
// Blocks API — always returns HTTP 500 (missing scope / permission denied).
|
||||
blocksPath := "/open-apis/docx/v1/documents/" + docToken + "/blocks"
|
||||
mux.HandleFunc(blocksPath, func(w http.ResponseWriter, r *http.Request) {
|
||||
w.WriteHeader(http.StatusInternalServerError)
|
||||
_, _ = w.Write([]byte(`{"code":99991400,"msg":"insufficient scope"}`))
|
||||
})
|
||||
|
||||
// Export task creation → returns ticket.
|
||||
mux.HandleFunc("/open-apis/drive/v1/export_tasks", func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.Method == http.MethodPost {
|
||||
writeJSON(w, core.ExportTaskCreateResponse{
|
||||
ApiResponse: core.ApiResponse{Code: 0},
|
||||
Data: core.ExportTaskCreateData{Ticket: "ticket-fb"},
|
||||
})
|
||||
return
|
||||
}
|
||||
http.NotFound(w, r)
|
||||
})
|
||||
// Export task status polling.
|
||||
mux.HandleFunc("/open-apis/drive/v1/export_tasks/ticket-fb", func(w http.ResponseWriter, r *http.Request) {
|
||||
writeJSON(w, core.ExportTaskStatusResponse{
|
||||
ApiResponse: core.ApiResponse{Code: 0},
|
||||
Data: core.ExportTaskStatusData{
|
||||
Result: core.ExportTaskResult{
|
||||
FileToken: "ft-export-fallback",
|
||||
FileSize: 512,
|
||||
JobStatus: 0, // done
|
||||
FileName: "fallback.docx",
|
||||
},
|
||||
},
|
||||
})
|
||||
})
|
||||
// Export file download.
|
||||
mux.HandleFunc("/open-apis/drive/v1/export_tasks/file/ft-export-fallback/download", func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "application/octet-stream")
|
||||
_, _ = w.Write([]byte("fake-exported-fallback-binary"))
|
||||
})
|
||||
// Drive file download endpoint (not expected to be called in the fallback path,
|
||||
// but registered to avoid 404 panic if the mux catches a stray request).
|
||||
mux.HandleFunc("/open-apis/drive/v1/files/", func(w http.ResponseWriter, r *http.Request) {
|
||||
if strings.HasSuffix(r.URL.Path, "/download") {
|
||||
w.Header().Set("Content-Type", "application/octet-stream")
|
||||
_, _ = w.Write([]byte("should-not-be-called"))
|
||||
return
|
||||
}
|
||||
http.NotFound(w, r)
|
||||
})
|
||||
|
||||
ts := httptest.NewServer(mux)
|
||||
return ts, &core.Config{AppID: "test-app-id", AppSecret: "test-app-secret", BaseURL: ts.URL}
|
||||
}
|
||||
|
||||
// TestFetchStream_DocxBlocksFallback proves that when the blocks API returns HTTP 500
|
||||
// (e.g. missing docx:document:readonly scope), the connector automatically falls back
|
||||
// to the export-API path and emits exactly one item with ContentType
|
||||
// "application/octet-stream", not "text/markdown". No hard failure occurs —
|
||||
// the sync completes successfully with the exported binary.
|
||||
func TestFetchStream_DocxBlocksFallback(t *testing.T) {
|
||||
t.Setenv("FEISHU_DOCX_PARSE_MODE", "blocks")
|
||||
const (
|
||||
nodeToken = "nt-fallback"
|
||||
objToken = "obj-fallback"
|
||||
)
|
||||
nodes := []core.WikiNode{{
|
||||
NodeToken: nodeToken,
|
||||
ObjToken: objToken,
|
||||
ObjType: "docx",
|
||||
Title: "Fallback Doc",
|
||||
ObjEditTime: "600",
|
||||
}}
|
||||
ts, cfg := fakeFeishuWithBlocksFallback(nodes, objToken)
|
||||
defer ts.Close()
|
||||
|
||||
c := NewConnector(core.RegionFeishu)
|
||||
h := &recordingHandler{}
|
||||
_, err := c.FetchStream(context.Background(), makeConfig(cfg, []string{"space1"}), nil, h)
|
||||
if err != nil {
|
||||
t.Fatalf("FetchStream() error: %v", err)
|
||||
}
|
||||
|
||||
// The fallback path must Emit exactly one item (exported binary, not multi-item).
|
||||
if len(h.emitted) != 1 {
|
||||
t.Fatalf("expected 1 emitted item (export fallback), got %d: %+v", len(h.emitted), h.emitted)
|
||||
}
|
||||
item := h.emitted[0]
|
||||
if item.ExternalID != nodeToken {
|
||||
t.Errorf("item.ExternalID = %q, want %q", item.ExternalID, nodeToken)
|
||||
}
|
||||
if item.ContentType == "text/markdown" {
|
||||
t.Errorf("item.ContentType = %q; want application/octet-stream (export fallback, not blocks path)", item.ContentType)
|
||||
}
|
||||
if item.ContentType != "application/octet-stream" {
|
||||
t.Errorf("item.ContentType = %q, want application/octet-stream", item.ContentType)
|
||||
}
|
||||
if item.Metadata["error"] != "" {
|
||||
t.Errorf("fallback item should have no error metadata, got %q", item.Metadata["error"])
|
||||
}
|
||||
// The export fallback cannot re-enumerate the doc's attachments, so it must
|
||||
// NOT set ReplacesSubtree — otherwise a transient blocks-API failure would
|
||||
// sweep and permanently delete the good attachment children from the prior
|
||||
// blocks-path sync with nothing to replace them.
|
||||
if item.ReplacesSubtree {
|
||||
t.Error("export-fallback item must not set ReplacesSubtree (would delete good prior attachments on a transient failure)")
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,7 @@
|
||||
package wiki
|
||||
|
||||
import "github.com/Tencent/WeKnora/internal/datasource/connector/feishu/core"
|
||||
|
||||
func txt(s string) *core.BlockText {
|
||||
return &core.BlockText{Elements: []core.TextElement{{TextRun: &core.TextRun{Content: s}}}}
|
||||
}
|
||||
@@ -0,0 +1,86 @@
|
||||
package wiki
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/datasource/connector/feishu/core"
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
)
|
||||
|
||||
// A Lark connector must surface wiki links on the Lark host. Before Lark became
|
||||
// its own connector type, every resource URL was hardcoded to feishu.cn, so Lark
|
||||
// users got links into a cloud their content does not live on.
|
||||
//
|
||||
// The fake server stands in for the Open Platform (base_url points at it), but
|
||||
// the resource URL is built from the region's web host, so this exercises the
|
||||
// real ListResources → wikiNodeToResource → region.WikiURL path.
|
||||
func TestListResources_ResourceURLFollowsRegion(t *testing.T) {
|
||||
cases := []struct {
|
||||
region core.Region
|
||||
wantURL string
|
||||
forbiddenIn string
|
||||
}{
|
||||
{core.RegionFeishu, "https://feishu.cn/wiki/space1", "larksuite"},
|
||||
{core.RegionLark, "https://larksuite.com/wiki/space1", "feishu"},
|
||||
}
|
||||
|
||||
for _, c := range cases {
|
||||
t.Run(c.region.Label, func(t *testing.T) {
|
||||
ts, cfg := fakeFeishu(nil)
|
||||
defer ts.Close()
|
||||
|
||||
conn := NewConnector(c.region)
|
||||
resources, err := conn.ListResources(context.Background(), makeConfig(cfg, nil), "")
|
||||
if err != nil {
|
||||
t.Fatalf("ListResources: %v", err)
|
||||
}
|
||||
if len(resources) != 1 {
|
||||
t.Fatalf("got %d resources, want 1", len(resources))
|
||||
}
|
||||
if got := resources[0].URL; got != c.wantURL {
|
||||
t.Errorf("resource URL = %q, want %q", got, c.wantURL)
|
||||
}
|
||||
if strings.Contains(resources[0].URL, c.forbiddenIn) {
|
||||
t.Errorf("%s resource URL leaks the other cloud's host: %q", c.region.Label, resources[0].URL)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// A Lark connector with no base_url override must reach open.larksuite.com,
|
||||
// never open.feishu.cn — the SDK-less client has no default of its own, so the
|
||||
// region is the only thing standing between a Lark app and the wrong cloud.
|
||||
func TestClient_ResolvesRegionHostWithoutOverride(t *testing.T) {
|
||||
cases := []struct {
|
||||
region core.Region
|
||||
want string
|
||||
}{
|
||||
{core.RegionFeishu, "https://open.feishu.cn"},
|
||||
{core.RegionLark, "https://open.larksuite.com"},
|
||||
}
|
||||
|
||||
for _, c := range cases {
|
||||
t.Run(c.region.Label, func(t *testing.T) {
|
||||
cfg, err := core.ParseFeishuConfig(makeConfigNoBaseURL(), c.region)
|
||||
if err != nil {
|
||||
t.Fatalf("ParseFeishuConfig: %v", err)
|
||||
}
|
||||
if cfg.GetBaseURL() != c.want {
|
||||
t.Errorf("baseURL = %q, want %q", cfg.GetBaseURL(), c.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// makeConfigNoBaseURL builds credentials without base_url, so the region default
|
||||
// is what resolves.
|
||||
func makeConfigNoBaseURL() *types.DataSourceConfig {
|
||||
return &types.DataSourceConfig{
|
||||
Credentials: map[string]interface{}{
|
||||
"app_id": "cli_x",
|
||||
"app_secret": "s",
|
||||
},
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user