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

This commit is contained in:
yawan
2026-09-22 16:19:06 +08:00
commit f7fc992f57
3929 changed files with 1006607 additions and 0 deletions
@@ -0,0 +1,497 @@
# Adding a New Connector - Step-by-Step Guide
This guide walks you through implementing a new connector (e.g., Feishu, Notion, Confluence) for the data source sync framework.
## Overview
A connector is an adapter that translates between WeKnora's data model and an external platform's API. It handles:
- Connection validation (credentials, connectivity)
- Resource listing (documents, spaces, folders the user can choose from)
- Full sync (fetch all items from selected resources)
- Incremental sync (fetch only changed items since last sync)
## Step 1: Create Connector Package Structure
```bash
mkdir -p internal/datasource/connector/yourtype/
```
Create three files:
- `client.go` - API client wrapper
- `connector.go` - Implements Connector interface
- `types.go` - Platform-specific data structures
## Step 2: Define Platform Types (types.go)
```go
package yourtype
import "time"
// Platform-specific configuration
type Config struct {
BaseURL string `json:"base_url"`
APIToken string `json:"api_token"`
// Or OAuth fields:
AccessToken string `json:"access_token"`
RefreshToken string `json:"refresh_token"`
ExpiresAt time.Time `json:"expires_at"`
}
// Platform-specific resource representation
type YourResource struct {
ID string
Name string
Type string // "document", "folder", "space", etc.
ModifiedAt time.Time
URL string
}
// Platform-specific item representation
type YourItem struct {
ID string
Title string
Content string
ContentHTML string
ModifiedAt time.Time
URL string
CreatedBy string
}
// Platform-specific pagination/cursor
type YourCursor struct {
Offset int `json:"offset,omitempty"`
LastModified time.Time `json:"last_modified,omitempty"`
PageToken string `json:"page_token,omitempty"`
}
```
## Step 3: Implement API Client (client.go)
```go
package yourtype
import (
"context"
"fmt"
"net/http"
"encoding/json"
)
type Client struct {
baseURL string
apiToken string
httpClient *http.Client
}
// NewClient creates a new API client
func NewClient(config *Config) *Client {
return &Client{
baseURL: config.BaseURL,
apiToken: config.APIToken,
httpClient: &http.Client{Timeout: 30 * time.Second},
}
}
// Example methods
func (c *Client) GetResources(ctx context.Context) ([]YourResource, error) {
// Call platform API
// Parse response
// Return resources
}
func (c *Client) GetDocument(ctx context.Context, docID string) (*YourItem, error) {
// Fetch single document
}
func (c *Client) GetDocumentsModifiedSince(ctx context.Context, since time.Time) ([]YourItem, error) {
// Fetch documents modified since timestamp
}
```
## Step 4: Implement Connector Interface (connector.go)
```go
package yourtype
import (
"context"
"fmt"
"github.com/Tencent/WeKnora/internal/types"
)
type YourConnector struct {
client *Client
}
// NewConnector creates a new connector
func NewConnector() *YourConnector {
return &YourConnector{}
}
// Type returns the connector type identifier
func (c *YourConnector) Type() string {
return types.ConnectorTypeYourType // Must match constant in types/datasource.go
}
// Validate verifies that the configuration is valid
func (c *YourConnector) Validate(ctx context.Context, config *types.DataSourceConfig) error {
if config == nil {
return fmt.Errorf("config is nil")
}
// Parse your type-specific config
yourConfig := &Config{}
if err := parseConfig(config, yourConfig); err != nil {
return fmt.Errorf("invalid config: %w", err)
}
// Create client
client := NewClient(yourConfig)
// Test connection
_, err := client.GetResources(ctx)
if err != nil {
return fmt.Errorf("connection failed: %w", err)
}
return nil
}
// ListResources lists available resources (documents, spaces, folders)
func (c *YourConnector) ListResources(ctx context.Context, config *types.DataSourceConfig) ([]types.Resource, error) {
yourConfig := &Config{}
if err := parseConfig(config, yourConfig); err != nil {
return nil, err
}
client := NewClient(yourConfig)
yourResources, err := client.GetResources(ctx)
if err != nil {
return nil, err
}
// Convert to WeKnora Resource format
resources := make([]types.Resource, len(yourResources))
for i, yr := range yourResources {
resources[i] = types.Resource{
ExternalID: yr.ID,
Name: yr.Name,
Type: yr.Type,
URL: yr.URL,
ModifiedAt: yr.ModifiedAt,
}
}
return resources, nil
}
// FetchAll performs a full sync
func (c *YourConnector) FetchAll(ctx context.Context, config *types.DataSourceConfig, resourceIDs []string) ([]types.FetchedItem, error) {
yourConfig := &Config{}
if err := parseConfig(config, yourConfig); err != nil {
return nil, err
}
client := NewClient(yourConfig)
var allItems []types.FetchedItem
// Fetch all documents from specified resources
for _, resourceID := range resourceIDs {
// Get documents from this resource (implementation depends on platform)
yourItems, err := client.GetDocumentsFromResource(ctx, resourceID)
if err != nil {
return nil, fmt.Errorf("failed to fetch resource %s: %w", resourceID, err)
}
// Convert to FetchedItem format
for _, yi := range yourItems {
item := types.FetchedItem{
ExternalID: yi.ID,
Title: yi.Title,
Content: []byte(yi.Content),
ContentType: "text/markdown",
FileName: fmt.Sprintf("%s.md", yi.Title),
URL: yi.URL,
UpdatedAt: yi.ModifiedAt,
SourceResourceID: resourceID,
Metadata: map[string]string{
"created_by": yi.CreatedBy,
"platform": "yourtype",
},
}
allItems = append(allItems, item)
}
}
return allItems, nil
}
// FetchIncremental performs an incremental sync
func (c *YourConnector) FetchIncremental(ctx context.Context, config *types.DataSourceConfig, cursor *types.SyncCursor) ([]types.FetchedItem, *types.SyncCursor, error) {
yourConfig := &Config{}
if err := parseConfig(config, yourConfig); err != nil {
return nil, nil, err
}
client := NewClient(yourConfig)
// Determine start time for incremental fetch
var sinceTime time.Time
if cursor != nil && !cursor.LastSyncTime.IsZero() {
sinceTime = cursor.LastSyncTime
} else {
sinceTime = time.Now().AddDate(0, 0, -7) // Default: last 7 days
}
// Fetch changed items
yourItems, err := client.GetDocumentsModifiedSince(ctx, sinceTime)
if err != nil {
return nil, nil, fmt.Errorf("incremental fetch failed: %w", err)
}
// Convert to FetchedItem format
items := make([]types.FetchedItem, len(yourItems))
for i, yi := range yourItems {
items[i] = types.FetchedItem{
ExternalID: yi.ID,
Title: yi.Title,
Content: []byte(yi.Content),
ContentType: "text/markdown",
FileName: fmt.Sprintf("%s.md", yi.Title),
URL: yi.URL,
UpdatedAt: yi.ModifiedAt,
Metadata: map[string]string{
"created_by": yi.CreatedBy,
"platform": "yourtype",
},
}
}
// Create new cursor for next sync
nextCursor := &types.SyncCursor{
LastSyncTime: time.Now(),
ConnectorCursor: map[string]interface{}{
"last_modified": time.Now(),
},
}
return items, nextCursor, nil
}
// Helper function to parse config
func parseConfig(config *types.DataSourceConfig, target interface{}) error {
data, err := json.Marshal(config.Credentials)
if err != nil {
return err
}
return json.Unmarshal(data, target)
}
```
## Step 5: Register in Container
Edit `internal/container/container.go`:
```go
import (
// ... existing imports
yourconnector "github.com/Tencent/WeKnora/internal/datasource/connector/yourtype"
)
func setupContainer() (*dig.Container, error) {
container := dig.New()
// ... existing registrations ...
// Register your connector
must(container.Provide(func() *datasource.Connector {
var c datasource.Connector = yourconnector.NewConnector()
return &c
}))
// ... rest of setup ...
}
```
Better yet, register through the registry in the connector setup:
```go
// In the service initialization section
connectorRegistry := datasource.NewConnectorRegistry()
connectorRegistry.Register(yourconnector.NewConnector())
connectorRegistry.Register(feishuconnector.NewConnector())
// ... etc
container.Provide(func() *datasource.ConnectorRegistry {
return connectorRegistry
})
```
## Step 6: Add Connector Type Constant
Edit `internal/types/datasource.go`:
```go
const (
// ... existing types ...
ConnectorTypeYourType = "yourtype"
)
```
## Step 7: Add Metadata
Edit `internal/datasource/connector.go`:
```go
var ConnectorMetadataRegistry = map[string]ConnectorMetadata{
// ... existing entries ...
types.ConnectorTypeYourType: {
Type: types.ConnectorTypeYourType,
Name: "Your Platform Name",
Description: "Sync documents from Your Platform",
Priority: X, // Lower number = higher priority in UI
AuthType: "oauth2", // or "api_key", "token", "password"
Capabilities: []string{"incremental", "webhook", "deletion_sync"},
},
}
```
## Step 8: Test Your Connector
```go
// Example test
func TestYourConnectorValidate(t *testing.T) {
connector := NewConnector()
config := &types.DataSourceConfig{
Type: types.ConnectorTypeYourType,
Credentials: map[string]interface{}{
"api_token": "test_token",
},
}
err := connector.Validate(context.Background(), config)
// assert no error
}
func TestYourConnectorFetchAll(t *testing.T) {
connector := NewConnector()
config := &types.DataSourceConfig{
Type: types.ConnectorTypeYourType,
Credentials: map[string]interface{}{
"api_token": "test_token",
},
ResourceIDs: []string{"resource_1"},
}
items, err := connector.FetchAll(context.Background(), config, []string{"resource_1"})
// assert results
}
```
## Checklist
- [ ] Created `internal/datasource/connector/yourtype/` package
- [ ] Implemented `types.go` with platform data structures
- [ ] Implemented `client.go` with API wrapper
- [ ] Implemented `connector.go` with Connector interface
- [ ] Added connector type constant
- [ ] Registered in container
- [ ] Added metadata entry
- [ ] Added unit tests
- [ ] Tested manually with real API
- [ ] Documented any special requirements
## Common Patterns
### OAuth Flow
If using OAuth, store tokens in config:
```go
type Config struct {
AccessToken string
RefreshToken string
ExpiresAt time.Time
}
// Refresh tokens when expired
func (c *Client) ensureValidToken(ctx context.Context) error {
if time.Now().After(c.config.ExpiresAt) {
return c.refreshToken(ctx)
}
return nil
}
```
### Pagination
For platforms with pagination:
```go
func (c *Client) GetDocumentsPage(ctx context.Context, pageToken string) (*Page, error) {
// Returns {Items, NextPageToken}
}
```
### Incremental Sync
For timestamp-based incremental sync:
```go
func (c *Client) GetModifiedSince(ctx context.Context, since time.Time) ([]Item, error) {
// Uses API parameter like &modified_after=2026-03-26T10:00:00Z
}
```
### Deletion Tracking
For platforms that report deletions:
```go
type Item struct {
IsDeleted bool // Set when item is deleted
}
```
## Testing with Real API
1. Set up test credentials
2. Create small test resource (e.g., single document)
3. Test each method:
```go
connector := NewConnector()
config := &types.DataSourceConfig{...}
// Test Validate
err := connector.Validate(ctx, config)
// Test ListResources
resources, err := connector.ListResources(ctx, config)
// Test FetchAll
items, err := connector.FetchAll(ctx, config, []string{resources[0].ExternalID})
// Test FetchIncremental
items, cursor, err := connector.FetchIncremental(ctx, config, nil)
```
---
## Example: Feishu Connector Reference
The Feishu connector would be a good first implementation since:
1. Feishu API is well-documented
2. WeKnora already has `internal/im/feishu/` for reference
3. Popular in China (key market)
4. Has webhook support for real-time sync
Structure:
```
internal/datasource/connector/feishu/
├── client.go (Feishu API client)
├── connector.go (Implements Connector)
└── types.go (Feishu types)
```
Key files to reference:
- `internal/im/feishu/adapter.go` - Feishu API patterns
- `internal/im/feishu/longconn.go` - Connection handling
This will be a good model for other connectors.
+260
View File
@@ -0,0 +1,260 @@
# WeKnora Data Source Sync Framework
## Overview
The data source sync framework enables WeKnora to automatically import and synchronize content from external platforms (Feishu, Notion, Confluence, etc.) into knowledge bases. This is the foundational layer upon which all specific connectors are built.
## Architecture
### Core Components
```
┌─────────────────────────────────────────────────────────────┐
│ External Data Sources │
│ (Feishu, Notion, Confluence, GitHub, Google Drive...) │
└─────────────┬───────────────────────────────────────────────┘
│
▼
┌─────────────────────────────────────────────────────────────┐
│ Connector Registry & Adapters │
│ Each platform (feishu/, notion/, confluence/, etc) │
│ implements the Connector interface │
└─────────────┬───────────────────────────────────────────────┘
│
▼
┌─────────────────────────────────────────────────────────────┐
│ DataSourceService (Business Logic) │
│ ├─ CreateDataSource / UpdateDataSource │
│ ├─ ManualSync / ListAvailableResources │
│ ├─ ValidateConnection / PauseDataSource │
│ └─ ProcessSync (asynq task handler) │
└─────────────┬───────────────────────────────────────────────┘
│
▼
┌─────────────────────────────────────────────────────────────┐
│ HTTP Handler & API Routes (/api/v1/datasource) │
│ REST endpoints for UI/programmatic access │
└─────────────┬───────────────────────────────────────────────┘
│
▼
┌─────────────────────────────────────────────────────────────┐
│ Database (MySQL / PostgreSQL) │
│ ├─ data_sources (configuration & state) │
│ └─ sync_logs (history & metadata) │
└─────────────────────────────────────────────────────────────┘
```
### Task Flow
```
Manual Trigger or Scheduled Job (cron)
↓
Scheduler enqueues asynq Task
↓
ProcessSync (task handler) executes
↓
Connector.Fetch* methods pull data
↓
Diff engine identifies changes (new/update/delete)
↓
KnowledgeService creates/updates Knowledge entries
↓
Documents parsed → chunks → vectors → indexed
```
## File Structure
```
internal/datasource/
├── connector.go # Connector interface & registry
├── errors.go # Error definitions
└── README.md # This file
internal/types/
├── datasource.go # Data models (DataSource, SyncLog, etc)
└── interfaces/datasource.go # Service interfaces
internal/application/
├── repository/datasource_repo.go # Database access layer
└── service/datasource_service.go # Business logic
internal/handler/
└── datasource.go # HTTP request handlers
internal/router/
├── router.go # Route registration
├── task.go # Asynq task registration
└── sync_task.go # Lite mode task registration
migrations/versioned/
└── 000028_datasource_tables.up/down.sql # Database schema
```
## Data Models
### DataSource
Represents a configured external data source:
- **ID**: Unique identifier
- **Type**: Connector type (feishu, notion, confluence, etc)
- **Config**: Encrypted credentials and configuration (JSON)
- **SyncSchedule**: Cron expression (e.g., "0 */6 * * *")
- **SyncMode**: "incremental" or "full"
- **Status**: active | paused | error
- **ConflictStrategy**: overwrite | skip
- **SyncDeletions**: Whether to sync deletions from source
- **LastSyncAt**: Timestamp of last successful sync
- **LastSyncCursor**: State for incremental sync
- **LastSyncResult**: Summary of last sync
### SyncLog
Tracks execution of each sync operation:
- **Status**: running | success | partial | failed | canceled
- **StartedAt/FinishedAt**: Execution timestamps
- **Items***: Counters (total, created, updated, deleted, skipped, failed)
- **ErrorMessage**: Error details if failed
- **Result**: Detailed sync result (JSON)
### Connector Interface
All external data sources implement this interface:
```go
type Connector interface {
Type() string
Validate(ctx, config) error
ListResources(ctx, config) ([]Resource, error)
FetchAll(ctx, config, resourceIDs) ([]FetchedItem, error)
FetchIncremental(ctx, config, cursor) ([]FetchedItem, *SyncCursor, error)
}
```
## API Endpoints
### Data Source Management
```
POST /api/v1/datasource # Create
GET /api/v1/datasource # List (by kb_id)
GET /api/v1/datasource/:id # Get
PUT /api/v1/datasource/:id # Update
DELETE /api/v1/datasource/:id # Delete
```
### Operations
```
POST /api/v1/datasource/:id/validate # Test connection
GET /api/v1/datasource/:id/resources # List available resources
POST /api/v1/datasource/:id/sync # Trigger manual sync
POST /api/v1/datasource/:id/pause # Pause
POST /api/v1/datasource/:id/resume # Resume
```
### Logs
```
GET /api/v1/datasource/:id/logs # List sync logs
GET /api/v1/datasource/logs/:log_id # Get specific log
```
### Metadata
```
GET /api/v1/datasource/types # Available connectors
```
## Implementation Guide
### Adding a New Connector
1. **Create the connector package** in `internal/datasource/connector/<type>/`
```go
// connector/<type>/client.go - API client
// connector/<type>/connector.go - Implements Connector interface
// connector/<type>/types.go - Type-specific models
```
2. **Implement the Connector interface**
```go
type YourConnector struct {
// fields
}
func (c *YourConnector) Type() string {
return types.ConnectorTypeYourType
}
func (c *YourConnector) Validate(ctx, config) error {
// Test connectivity and credentials
}
func (c *YourConnector) ListResources(ctx, config) ([]types.Resource, error) {
// List available resources (documents, spaces, etc)
}
func (c *YourConnector) FetchAll(ctx, config, resourceIDs) ([]types.FetchedItem, error) {
// Full sync - fetch all items from specified resources
}
func (c *YourConnector) FetchIncremental(ctx, config, cursor) ([]types.FetchedItem, *types.SyncCursor, error) {
// Incremental sync - fetch only changed items since cursor
}
```
3. **Register the connector** in the container (in container.go)
```go
container.Provide(func() datasource.Connector {
return yourconnector.NewConnector()
})
```
4. **Add metadata** in `connector.go`
```go
types.ConnectorTypeYourType: {
Type: types.ConnectorTypeYourType,
Name: "Your Platform",
AuthType: "oauth2",
// ...
}
```
## Database Schema
### data_sources table
- Stores data source configurations
- Indexes: tenant_id, knowledge_base_id, type, status
- Encrypted config field (AES-256-GCM)
### sync_logs table
- Stores sync operation history
- Indexes: data_source_id, tenant_id, status, started_at
- Foreign key to data_sources with cascade delete
## Key Design Decisions
1. **Adapter Pattern**: Each connector is a separate implementation, easy to add new ones
2. **Async Task Queue**: Syncs are non-blocking, using asynq (or sync mode in Lite)
3. **Incremental Sync**: Supports cursor-based incremental syncs to minimize API calls
4. **Encryption**: API keys/tokens stored encrypted with AES-256-GCM
5. **Multi-tenant**: All operations are tenant-isolated
6. **Channel Tracking**: Knowledge.Channel field tracks data source type for origin tracking
## Future Enhancements
1. **Webhook Support**: Real-time push from platforms that support webhooks
2. **Conflict Resolution**: Advanced merge/conflict strategies for overlapping content
3. **Rate Limiting**: Per-connector rate limiting and backoff strategies
4. **Scheduling**: Full cron scheduler with time zone support
5. **Monitoring**: Metrics, alerting, and sync health dashboards
6. **Filtering**: User-defined filters for selective syncing (by title, date, tags, etc)
7. **Transformation**: Content transformation pipelines (e.g., extract tables, summarize)
8. **Deduplication**: Smart deduplication across multiple sources
## Testing
The framework is designed to be testable:
- Mock connectors can be created for testing
- No external dependencies required for core logic
- Database operations isolated and mockable
## Error Handling
All operations return detailed error information:
- Connection failures are captured and status updated
- Partial failures are logged (some items succeed, some fail)
- Error messages stored in DataSource for user troubleshooting
+310
View File
@@ -0,0 +1,310 @@
package datasource
import (
"context"
"github.com/Tencent/WeKnora/internal/types"
)
// Connector is the interface that all external data source connectors must implement.
// Each connector (Feishu, Notion, Confluence, etc.) provides an implementation of this interface.
type Connector interface {
// Type returns the connector type identifier (e.g., "feishu", "notion")
Type() string
// Validate verifies that the provided configuration is valid by testing connectivity
// and checking credentials. Returns error if validation fails.
Validate(ctx context.Context, config *types.DataSourceConfig) error
// ListResources lists available resources that can be synced (documents, spaces, folders, etc.)
// Returns a list of Resource objects that the user can select for syncing.
//
// parentID controls lazy (on-demand) loading of hierarchical resources:
// - parentID == "" → return the top-level resources (e.g. Feishu wiki spaces).
// - parentID != "" → return only the direct children of that resource.
// Connectors whose listing is already flat or returns the full tree in a single
// call may ignore parentID for the root call and return an empty slice for any
// non-empty parentID.
ListResources(ctx context.Context, config *types.DataSourceConfig, parentID string) ([]types.Resource, error)
// ResolveResourceAncestors resolves, for each of the given resource IDs, the
// ExternalIDs of every ancestor whose direct children must be loaded so a
// lazily-loaded picker can reveal a pre-existing (possibly deeply nested)
// selection. The returned set is deduplicated and unordered.
//
// It exists so connectors that load their tree one level at a time (e.g. the
// Feishu wiki) can expose, in O(depth) per selection, the path back to the
// root without re-traversing the whole tree. Connectors that already return
// the full tree (Notion) or a flat list (Yuque) have nothing to reveal and
// return an empty slice.
ResolveResourceAncestors(
ctx context.Context, config *types.DataSourceConfig, resourceIDs []string,
) ([]string, error)
// FetchAll performs a full sync of the specified resources.
// Returns all items from the given resource IDs.
FetchAll(ctx context.Context, config *types.DataSourceConfig, resourceIDs []string) ([]types.FetchedItem, error)
// FetchIncremental performs an incremental sync based on the provided cursor.
// Returns items that have changed since the last sync, a new cursor for the next sync,
// and an error if the operation fails.
FetchIncremental(ctx context.Context, config *types.DataSourceConfig, cursor *types.SyncCursor) ([]types.FetchedItem, *types.SyncCursor, error)
}
// StreamHandler receives items and progress checkpoints emitted during a
// streaming fetch. The service implements it to ingest each item as it arrives
// (bounding memory to one item instead of the whole wiki) and to persist the
// connector cursor at page boundaries, so a sync that times out mid-traversal
// resumes from the last checkpoint instead of restarting from scratch
// (Tencent/WeKnora#2136).
type StreamHandler interface {
// Emit ingests a single fetched item. Returning an error aborts the
// stream: the connector stops fetching and propagates the error, since a
// failed ingest means the sync is failing and further API calls are wasted.
Emit(ctx context.Context, item types.FetchedItem) error
// Checkpoint persists the cursor reached so far. The cursor is only valid
// for the duration of the call (the connector may keep mutating its backing
// maps afterwards), so implementations must serialize it synchronously.
//
// The cursor MUST be a complete resumable snapshot, not a delta: resuming
// from it must reproduce all progress so far. This is what lets the service
// treat a checkpoint as a safe restart point and, for a full sync, drop the
// prior baseline without losing already-synced state.
Checkpoint(ctx context.Context, cursor *types.SyncCursor) error
}
// StreamingConnector is an optional interface. Connectors that implement it let
// the service interleave fetch→ingest→checkpoint so a large sync persists
// incrementally and resumes after a timeout, rather than holding every item in
// memory and losing all progress on retry. Connectors that do not implement it
// fall back to FetchAll / FetchIncremental unchanged.
type StreamingConnector interface {
Connector
// FetchStream walks the configured resources starting from cursor (nil =
// from the beginning / full sync), calling h.Emit for each changed item and
// h.Checkpoint at page boundaries. It returns the final cursor for the next
// sync. Nodes already recorded in cursor at their current edit time are
// skipped, which is what makes a resumed sync converge.
FetchStream(
ctx context.Context, config *types.DataSourceConfig,
cursor *types.SyncCursor, h StreamHandler,
) (*types.SyncCursor, error)
}
// ConnectorRegistry manages the registration and lookup of available connectors
type ConnectorRegistry struct {
connectors map[string]Connector
}
// NewConnectorRegistry creates a new connector registry
func NewConnectorRegistry() *ConnectorRegistry {
return &ConnectorRegistry{
connectors: make(map[string]Connector),
}
}
// Register registers a connector with the registry
func (r *ConnectorRegistry) Register(connector Connector) error {
if connector == nil {
return ErrConnectorNil
}
if connector.Type() == "" {
return ErrConnectorTypeEmpty
}
r.connectors[connector.Type()] = connector
return nil
}
// Get retrieves a connector by type
func (r *ConnectorRegistry) Get(connectorType string) (Connector, error) {
connector, exists := r.connectors[connectorType]
if !exists {
return nil, ErrConnectorNotFound
}
return connector, nil
}
// List returns all registered connector types
func (r *ConnectorRegistry) List() []string {
types := make([]string, 0, len(r.connectors))
for t := range r.connectors {
types = append(types, t)
}
return types
}
// ConnectorMetadata provides metadata about available connectors
type ConnectorMetadata struct {
Type string `json:"type"`
Name string `json:"name"`
Description string `json:"description"`
Icon string `json:"icon,omitempty"`
Priority int `json:"priority"` // Priority order for UI display (lower = higher priority)
AuthType string `json:"auth_type"` // "oauth2", "api_key", "token", etc.
Capabilities []string `json:"capabilities"` // "incremental", "webhook", "deletion_sync", etc.
}
// GetConnectorMetadata returns metadata for all available connectors
// This is used by the frontend to display connector options
var ConnectorMetadataRegistry = map[string]ConnectorMetadata{
types.ConnectorTypeFeishu: {
Type: types.ConnectorTypeFeishu,
Name: "Feishu (飞书)",
Description: "Sync documents, wikis, and content from Feishu",
Priority: 0,
AuthType: "oauth2",
Capabilities: []string{"incremental", "deletion_sync"},
},
types.ConnectorTypeLark: {
Type: types.ConnectorTypeLark,
Name: "Lark",
Description: "Sync documents, wikis, and content from Lark (Feishu international)",
Priority: 0,
AuthType: "oauth2",
Capabilities: []string{"incremental", "deletion_sync"},
},
types.ConnectorTypeFeishuDrive: {
Type: types.ConnectorTypeFeishuDrive,
Name: "Feishu Drive (飞书云盘)",
Description: "Sync documents and files from a Feishu Drive folder",
Priority: 0,
AuthType: "oauth2",
Capabilities: []string{"incremental", "deletion_sync"},
},
types.ConnectorTypeLarkDrive: {
Type: types.ConnectorTypeLarkDrive,
Name: "Lark Drive",
Description: "Sync documents and files from a Lark Drive folder",
Priority: 0,
AuthType: "oauth2",
Capabilities: []string{"incremental", "deletion_sync"},
},
types.ConnectorTypeNotion: {
Type: types.ConnectorTypeNotion,
Name: "Notion",
Description: "Sync pages and databases from Notion",
Priority: 1,
AuthType: "api_key",
Capabilities: []string{"incremental"},
},
types.ConnectorTypeConfluence: {
Type: types.ConnectorTypeConfluence,
Name: "Confluence",
Description: "Sync spaces and pages from Atlassian Confluence",
Priority: 2,
AuthType: "api_key",
Capabilities: []string{"incremental"},
},
types.ConnectorTypeYuque: {
Type: types.ConnectorTypeYuque,
Name: "Yuque (语雀)",
Description: "Sync knowledge bases and documents from Yuque",
Priority: 3,
AuthType: "api_key",
Capabilities: []string{"incremental"},
},
types.ConnectorTypeIMA: {
Type: types.ConnectorTypeIMA,
Name: "Tencent IMA (ima.qq.com)",
Description: "Sync knowledge bases and documents from Tencent IMA",
Priority: 3,
AuthType: "api_key",
Capabilities: []string{"incremental", "deletion_sync"},
},
types.ConnectorTypeGitHub: {
Type: types.ConnectorTypeGitHub,
Name: "GitHub",
Description: "Sync repositories, wikis, and issues from GitHub",
Priority: 4,
AuthType: "oauth2",
Capabilities: []string{"incremental"},
},
types.ConnectorTypeGoogleDrive: {
Type: types.ConnectorTypeGoogleDrive,
Name: "Google Drive",
Description: "Sync documents and files from Google Drive",
Priority: 5,
AuthType: "oauth2",
Capabilities: []string{"incremental"},
},
types.ConnectorTypeOneDrive: {
Type: types.ConnectorTypeOneDrive,
Name: "OneDrive / SharePoint",
Description: "Sync documents and files from Microsoft OneDrive",
Priority: 6,
AuthType: "oauth2",
Capabilities: []string{"incremental"},
},
types.ConnectorTypeDingTalk: {
Type: types.ConnectorTypeDingTalk,
Name: "DingTalk (钉钉)",
Description: "Sync documents and content from DingTalk",
Priority: 7,
AuthType: "api_key",
Capabilities: []string{"incremental"},
},
types.ConnectorTypeWebCrawler: {
Type: types.ConnectorTypeWebCrawler,
Name: "Web Crawler (Sitemap)",
Description: "Crawl websites via Sitemap.xml",
Priority: 9,
AuthType: "none",
Capabilities: []string{},
},
types.ConnectorTypeSlack: {
Type: types.ConnectorTypeSlack,
Name: "Slack",
Description: "Sync channel messages and files from Slack",
Priority: 10,
AuthType: "oauth2",
Capabilities: []string{"incremental"},
},
types.ConnectorTypeIMAP: {
Type: types.ConnectorTypeIMAP,
Name: "Email (IMAP)",
Description: "Sync email content from IMAP servers",
Priority: 11,
AuthType: "password",
Capabilities: []string{},
},
types.ConnectorTypeRSS: {
Type: types.ConnectorTypeRSS,
Name: "RSS / Atom Feed",
Description: "Sync articles from RSS/Atom feeds",
Priority: 12,
AuthType: "custom",
Capabilities: []string{"incremental"},
},
types.ConnectorTypeGitLab: {
Type: types.ConnectorTypeGitLab,
Name: "GitLab",
Description: "Sync files from GitLab projects",
Priority: 8,
AuthType: "token",
Capabilities: []string{"incremental", "hierarchical"},
},
}
// ListAvailableConnectors returns all available connector metadata
// sorted by priority
func ListAvailableConnectors() []ConnectorMetadata {
metadata := make([]ConnectorMetadata, 0, len(ConnectorMetadataRegistry))
for _, meta := range ConnectorMetadataRegistry {
metadata = append(metadata, meta)
}
// Sort by priority (insertion sort for simplicity)
for i := 1; i < len(metadata); i++ {
key := metadata[i]
j := i - 1
for j >= 0 && metadata[j].Priority > key.Priority {
metadata[j+1] = metadata[j]
j--
}
metadata[j+1] = key
}
return metadata
}
@@ -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",
},
}
}
@@ -0,0 +1,267 @@
package gitlab
import (
"context"
"encoding/base64"
"encoding/json"
"errors"
"fmt"
"io"
"net/http"
"net/url"
"strconv"
"strings"
"time"
"github.com/Tencent/WeKnora/internal/datasource"
)
type client struct {
baseURL, token string
http *http.Client
}
type apiError struct {
endpoint string
status int
}
func (e *apiError) Error() string {
return fmt.Sprintf("gitlab API %s: status %d", e.endpoint, e.status)
}
type project struct {
ID int64 `json:"id"`
PathWithNamespace string `json:"path_with_namespace"`
Name string `json:"name"`
WebURL string `json:"web_url"`
DefaultBranch string `json:"default_branch"`
Namespace struct {
ID int64 `json:"id"`
} `json:"namespace"`
}
type treeEntry struct {
ID string `json:"id"`
Name string `json:"name"`
Type string `json:"type"`
Path string `json:"path"`
}
type comparison struct {
Diffs []struct {
OldPath string `json:"old_path"`
NewPath string `json:"new_path"`
NewFile bool `json:"new_file"`
DeletedFile bool `json:"deleted_file"`
RenamedFile bool `json:"renamed_file"`
} `json:"diffs"`
CompareTimeout bool `json:"compare_timeout"`
CompareSameRef bool `json:"compare_same_ref"`
}
func newClient(baseURL, token string) (*client, error) {
baseURL = strings.TrimRight(strings.TrimSpace(baseURL), "/")
if baseURL == "" || strings.TrimSpace(token) == "" {
return nil, fmt.Errorf("GitLab platform configuration is missing")
}
if err := datasource.ValidateConnectorBaseURL(baseURL); err != nil {
return nil, err
}
if !strings.Contains(baseURL, "://") {
baseURL = "https://" + baseURL
}
if !strings.HasSuffix(baseURL, "/api/v4") {
baseURL += "/api/v4"
}
return &client{baseURL: baseURL, token: token, http: datasource.NewConnectorHTTPClient(30 * time.Second)}, nil
}
func (c *client) get(ctx context.Context, endpoint string, out interface{}) error {
req, err := http.NewRequestWithContext(ctx, http.MethodGet, c.baseURL+endpoint, nil)
if err != nil {
return err
}
req.Header.Set("PRIVATE-TOKEN", c.token)
resp, err := c.http.Do(req)
if err != nil {
return err
}
defer resp.Body.Close()
body, err := io.ReadAll(resp.Body)
if err != nil {
return err
}
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
return &apiError{endpoint: endpoint, status: resp.StatusCode}
}
return json.Unmarshal(body, out)
}
func (c *client) getRaw(ctx context.Context, endpoint string) ([]byte, error) {
req, err := http.NewRequestWithContext(ctx, http.MethodGet, c.baseURL+endpoint, nil)
if err != nil {
return nil, err
}
req.Header.Set("PRIVATE-TOKEN", c.token)
resp, err := c.http.Do(req)
if err != nil {
return nil, err
}
defer resp.Body.Close()
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
return nil, &apiError{endpoint: endpoint, status: resp.StatusCode}
}
return io.ReadAll(resp.Body)
}
// projectPath encodes a GitLab project identifier for URL path segments.
// Numeric IDs are used verbatim. Namespace paths accept either "group/project"
// or a once-encoded "group%2Fproject" without double-encoding percent signs.
func projectPath(id string) string {
id = strings.TrimSpace(id)
if id == "" {
return ""
}
if _, err := strconv.ParseInt(id, 10, 64); err == nil {
return id
}
decoded := id
if strings.Contains(id, "%") {
if unescaped, err := url.PathUnescape(id); err == nil {
decoded = unescaped
}
}
if strings.Contains(decoded, "/") {
parts := strings.Split(decoded, "/")
for i, part := range parts {
parts[i] = url.PathEscape(part)
}
return strings.Join(parts, "%2F")
}
return url.PathEscape(decoded)
}
func (c *client) project(ctx context.Context, id string) (*project, error) {
var p project
err := c.get(ctx, "/projects/"+projectPath(id), &p)
return &p, err
}
func (c *client) projects(ctx context.Context) ([]project, error) {
var p []project
err := c.get(ctx, "/projects?membership=true&per_page=100&order_by=path_with_namespace&sort=asc", &p)
return p, err
}
// ping verifies that the supplied private token is accepted by this GitLab
// instance. It avoids project-list ordering parameters that older GitLab
// deployments may reject even when the token is valid.
func (c *client) ping(ctx context.Context) error {
var user struct {
ID int64 `json:"id"`
}
return c.get(ctx, "/user", &user)
}
func (c *client) commitSHA(ctx context.Context, id, ref string) (string, error) {
var v struct {
ID string `json:"id"`
}
err := c.get(ctx, "/projects/"+projectPath(id)+"/repository/commits/"+url.PathEscape(ref), &v)
return v.ID, err
}
func (c *client) tree(ctx context.Context, id, ref, dir string) ([]treeEntry, error) {
q := url.Values{"ref": {ref}, "per_page": {"100"}, "page": {"1"}}
if dir != "" {
q.Set("path", dir)
}
endpoint := "/projects/" + projectPath(id) + "/repository/tree"
var all []treeEntry
for {
var page []treeEntry
nextPage, err := c.getTreePage(ctx, endpoint, q, &page)
if err != nil {
return nil, err
}
all = append(all, page...)
if nextPage == "" {
return all, nil
}
q.Set("page", nextPage)
}
}
func (c *client) getTreePage(ctx context.Context, endpoint string, query url.Values, out interface{}) (string, error) {
req, err := http.NewRequestWithContext(ctx, http.MethodGet, c.baseURL+endpoint+"?"+query.Encode(), nil)
if err != nil {
return "", err
}
req.Header.Set("PRIVATE-TOKEN", c.token)
resp, err := c.http.Do(req)
if err != nil {
return "", err
}
defer resp.Body.Close()
body, err := io.ReadAll(resp.Body)
if err != nil {
return "", err
}
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
return "", &apiError{endpoint: endpoint, status: resp.StatusCode}
}
if err := json.Unmarshal(body, out); err != nil {
return "", err
}
return resp.Header.Get("X-Next-Page"), nil
}
func (c *client) raw(ctx context.Context, id, ref, file string) ([]byte, error) {
q := url.Values{"ref": {ref}}
encodedFile := gitlabFilePathEscape(file)
rawEndpoint := "/projects/" + projectPath(id) + "/repository/files/" + encodedFile + "/raw?" + q.Encode()
content, err := c.getRaw(ctx, rawEndpoint)
if err == nil {
return content, nil
}
var apiErr *apiError
if !errors.As(err, &apiErr) || apiErr.status != http.StatusNotFound {
return nil, fmt.Errorf("gitlab raw file: %w", err)
}
// Some GitLab deployments expose the file detail endpoint but return 404
// for the otherwise standard /raw route. The detail response contains the
// same content as base64 and provides a compatible fallback.
var detail struct {
Encoding string `json:"encoding"`
Content string `json:"content"`
}
detailEndpoint := "/projects/" + projectPath(id) + "/repository/files/" + encodedFile + "?" + q.Encode()
if err := c.get(ctx, detailEndpoint, &detail); err != nil {
return nil, fmt.Errorf("gitlab file content: %w", err)
}
if detail.Encoding != "base64" {
return nil, fmt.Errorf("gitlab file content: unsupported encoding %q", detail.Encoding)
}
content, err = base64.StdEncoding.DecodeString(detail.Content)
if err != nil {
return nil, fmt.Errorf("gitlab file content: decode base64: %w", err)
}
return content, nil
}
// gitlabFilePathEscape mirrors the company GitLab raw-file route: only ASCII
// letters, digits, hyphen and underscore remain literal. In particular dots,
// path separators and UTF-8 bytes must be percent-encoded.
func gitlabFilePathEscape(file string) string {
const hex = "0123456789ABCDEF"
var b strings.Builder
b.Grow(len(file) * 3)
for _, c := range []byte(file) {
if (c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z') || (c >= '0' && c <= '9') || c == '-' || c == '_' {
b.WriteByte(c)
continue
}
b.WriteByte('%')
b.WriteByte(hex[c>>4])
b.WriteByte(hex[c&0x0f])
}
return b.String()
}
func (c *client) compare(ctx context.Context, id, from, to string) (*comparison, error) {
q := url.Values{"from": {from}, "to": {to}}
var v comparison
err := c.get(ctx, "/projects/"+projectPath(id)+"/repository/compare?"+q.Encode(), &v)
return &v, err
}
@@ -0,0 +1,408 @@
package gitlab
import (
"context"
"errors"
"fmt"
"net/http"
"net/http/httptest"
"strings"
"testing"
"github.com/Tencent/WeKnora/internal/types"
secutils "github.com/Tencent/WeKnora/internal/utils"
)
func allowLocalGitLabServer(t *testing.T) {
t.Helper()
t.Setenv("SSRF_WHITELIST", "127.0.0.1,::1,localhost")
secutils.ResetSSRFWhitelistForTest()
t.Cleanup(secutils.ResetSSRFWhitelistForTest)
}
func TestConnectorValidateUsesDataSourceCredentials(t *testing.T) {
allowLocalGitLabServer(t)
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.URL.Path != "/api/v4/user" {
http.NotFound(w, r)
return
}
if got := r.Header.Get("PRIVATE-TOKEN"); got != "per-source-token" {
t.Fatalf("PRIVATE-TOKEN = %q", got)
}
w.Header().Set("Content-Type", "application/json")
_, _ = w.Write([]byte(`{"id": 42}`))
}))
defer server.Close()
ds := &types.DataSourceConfig{
Credentials: map[string]interface{}{
"base_url": server.URL,
"access_token": "per-source-token",
},
}
if err := NewConnector().Validate(context.Background(), ds); err != nil {
t.Fatalf("Validate() error = %v", err)
}
}
func TestConnectorValidateReturnsGitLabAPIError(t *testing.T) {
allowLocalGitLabServer(t)
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.URL.Path != "/api/v4/user" {
http.NotFound(w, r)
return
}
http.Error(w, "invalid token", http.StatusUnauthorized)
}))
defer server.Close()
err := NewConnector().Validate(context.Background(), &types.DataSourceConfig{
Credentials: map[string]interface{}{
"base_url": server.URL,
"access_token": "invalid-token",
},
})
var apiErr *apiError
if !errors.As(err, &apiErr) {
t.Fatalf("Validate() error = %v, want GitLab API error", err)
}
if apiErr.endpoint != "/user" || apiErr.status != http.StatusUnauthorized {
t.Fatalf("apiError = %#v", apiErr)
}
}
func TestConnectorValidateRejectsMissingCredentials(t *testing.T) {
err := NewConnector().Validate(context.Background(), &types.DataSourceConfig{
Credentials: map[string]interface{}{"base_url": "https://gitlab.example.com"},
})
if err == nil || err.Error() != "GitLab platform configuration is missing" {
t.Fatalf("Validate() error = %v", err)
}
}
func TestConnectorValidateRejectsMissingProjectsOnSave(t *testing.T) {
allowLocalGitLabServer(t)
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.URL.Path != "/api/v4/user" {
http.NotFound(w, r)
return
}
w.Header().Set("Content-Type", "application/json")
_, _ = w.Write([]byte(`{"id": 42}`))
}))
defer server.Close()
err := NewConnector().Validate(context.Background(), &types.DataSourceConfig{
Credentials: map[string]interface{}{
"base_url": server.URL,
"access_token": "token",
},
Settings: map[string]interface{}{
"projects": []interface{}{},
},
})
if err == nil {
t.Fatal("expected missing projects validation error")
}
}
func TestConnectorConfiguredDoesNotReuseRegistryClient(t *testing.T) {
allowLocalGitLabServer(t)
var requestedTokens []string
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.URL.Path != "/api/v4/user" {
http.NotFound(w, r)
return
}
requestedTokens = append(requestedTokens, r.Header.Get("PRIVATE-TOKEN"))
w.Header().Set("Content-Type", "application/json")
_, _ = w.Write([]byte(`{"id": 42}`))
}))
defer server.Close()
connector := NewConnector()
first := &types.DataSourceConfig{
Credentials: map[string]interface{}{"base_url": server.URL, "access_token": "token-a"},
}
second := &types.DataSourceConfig{
Credentials: map[string]interface{}{"base_url": server.URL, "access_token": "token-b"},
}
if err := connector.Validate(context.Background(), first); err != nil {
t.Fatalf("first validate: %v", err)
}
if err := connector.Validate(context.Background(), second); err != nil {
t.Fatalf("second validate: %v", err)
}
if len(requestedTokens) != 2 || requestedTokens[0] != "token-a" || requestedTokens[1] != "token-b" {
t.Fatalf("requested tokens = %#v", requestedTokens)
}
}
func TestNewClientNormalizesAPIBaseURL(t *testing.T) {
allowLocalGitLabServer(t)
server := httptest.NewServer(http.NotFoundHandler())
defer server.Close()
c, err := newClient(server.URL+"/", "token")
if err != nil {
t.Fatal(err)
}
if got, want := c.baseURL, server.URL+"/api/v4"; got != want {
t.Fatalf("baseURL = %q, want %q", got, want)
}
}
func TestFetchIncrementalSyncsMultipleProjects(t *testing.T) {
allowLocalGitLabServer(t)
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "application/json")
switch r.URL.Path {
case "/api/v4/projects/1":
_, _ = w.Write([]byte(`{"id":1,"name":"project-one","path_with_namespace":"group/project-one","web_url":"https://gitlab.test/group/project-one","default_branch":"master"}`))
case "/api/v4/projects/2":
_, _ = w.Write([]byte(`{"id":2,"name":"project-two","path_with_namespace":"group/project-two","web_url":"https://gitlab.test/group/project-two","default_branch":"master"}`))
case "/api/v4/projects/1/repository/commits/master":
_, _ = w.Write([]byte(`{"id":"commit-1"}`))
case "/api/v4/projects/2/repository/commits/master":
_, _ = w.Write([]byte(`{"id":"commit-2"}`))
case "/api/v4/projects/1/repository/tree":
_, _ = w.Write([]byte(`[{"name":"one.md","type":"blob","path":"one.md"}]`))
case "/api/v4/projects/2/repository/tree":
_, _ = w.Write([]byte(`[{"name":"two.md","type":"blob","path":"two.md"}]`))
default:
if strings.HasPrefix(r.URL.EscapedPath(), "/api/v4/projects/1/repository/files/one%2Emd/raw") {
_, _ = w.Write([]byte("one"))
return
}
if strings.HasPrefix(r.URL.EscapedPath(), "/api/v4/projects/2/repository/files/two%2Emd/raw") {
_, _ = w.Write([]byte("two"))
return
}
http.NotFound(w, r)
}
}))
defer server.Close()
connector := NewConnector()
config := &types.DataSourceConfig{
Credentials: map[string]interface{}{
"base_url": server.URL,
"access_token": "token",
},
Settings: map[string]interface{}{
"projects": []interface{}{
map[string]interface{}{"project_id": "1", "ref": "master", "paths": []interface{}{}},
map[string]interface{}{"project_id": "2", "ref": "master", "paths": []interface{}{}},
},
}}
items, cursor, err := connector.FetchIncremental(context.Background(), config, nil)
if err != nil {
t.Fatalf("first sync: %v", err)
}
if len(items) != 2 {
t.Fatalf("first sync items = %d, want 2", len(items))
}
if cursor == nil || fmt.Sprint(cursor.ConnectorCursor["projects"]) == "" {
t.Fatal("first sync did not return per-project cursor state")
}
items, _, err = connector.FetchIncremental(context.Background(), config, cursor)
if err != nil {
t.Fatalf("second sync: %v", err)
}
if len(items) != 0 {
t.Fatalf("second sync items = %d, want 0", len(items))
}
}
type gitLabStreamRecorder struct {
items []types.FetchedItem
checkpoints []*types.SyncCursor
}
func (h *gitLabStreamRecorder) Emit(_ context.Context, item types.FetchedItem) error {
h.items = append(h.items, item)
return nil
}
func (h *gitLabStreamRecorder) Checkpoint(_ context.Context, cursor *types.SyncCursor) error {
h.checkpoints = append(h.checkpoints, cursor)
return nil
}
func TestFetchStreamFiltersUnsupportedFilesAndCheckpointsProjects(t *testing.T) {
allowLocalGitLabServer(t)
var rawRequests []string
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "application/json")
switch r.URL.Path {
case "/api/v4/projects/1":
_, _ = w.Write([]byte(`{"id":1,"name":"docs","path_with_namespace":"group/docs","web_url":"https://gitlab.test/group/docs","default_branch":"main"}`))
case "/api/v4/projects/1/repository/commits/main":
_, _ = w.Write([]byte(`{"id":"commit-1"}`))
case "/api/v4/projects/1/repository/tree":
_, _ = w.Write([]byte(`[
{"name":"README.md","type":"blob","path":"README.md"},
{"name":"server.go","type":"blob","path":"server.go"},
{"name":"payload.exe","type":"blob","path":"payload.exe"}
]`))
default:
if strings.Contains(r.URL.EscapedPath(), "/repository/files/") {
rawRequests = append(rawRequests, r.URL.EscapedPath())
_, _ = w.Write([]byte("# readme"))
return
}
http.NotFound(w, r)
}
}))
defer server.Close()
connector := NewConnector()
config := &types.DataSourceConfig{
Credentials: map[string]interface{}{
"base_url": server.URL,
"access_token": "token",
},
Settings: map[string]interface{}{
"projects": []interface{}{map[string]interface{}{"project_id": "1", "paths": []interface{}{}}},
},
}
handler := &gitLabStreamRecorder{}
next, err := connector.FetchStream(context.Background(), config, nil, handler)
if err != nil {
t.Fatalf("FetchStream() error = %v", err)
}
if len(handler.items) != 1 || handler.items[0].FileName != "docs-main/README.md" {
t.Fatalf("emitted items = %#v", handler.items)
}
if len(rawRequests) != 1 || !strings.Contains(rawRequests[0], "README%2Emd") {
t.Fatalf("raw requests = %v, want only README.md", rawRequests)
}
if len(handler.checkpoints) != 1 || next == nil {
t.Fatalf("checkpoints = %d, next = %#v", len(handler.checkpoints), next)
}
projects, _ := next.ConnectorCursor["projects"].(map[string]string)
if projects["1"] != "commit-1" {
t.Fatalf("cursor projects = %#v", next.ConnectorCursor["projects"])
}
}
func TestIsSupportedFile(t *testing.T) {
for _, tc := range []struct {
file string
want bool
}{
{file: "docs/guide.MD", want: true},
{file: "docs/guide.mdx", want: true},
{file: "report.pdf", want: true},
{file: "src/main.go", want: false},
{file: "archive.tar.gz", want: false},
{file: "LICENSE", want: false},
} {
if got := isSupportedFile(tc.file); got != tc.want {
t.Errorf("isSupportedFile(%q) = %v, want %v", tc.file, got, tc.want)
}
}
}
func TestTreeFollowsGitLabPagination(t *testing.T) {
allowLocalGitLabServer(t)
var requestedPages []string
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.URL.Path != "/api/v4/projects/1/repository/tree" {
http.NotFound(w, r)
return
}
requestedPages = append(requestedPages, r.URL.Query().Get("page"))
w.Header().Set("Content-Type", "application/json")
switch r.URL.Query().Get("page") {
case "1":
w.Header().Set("X-Next-Page", "2")
_, _ = w.Write([]byte(`[{"name":"one.md","type":"blob","path":"one.md"}]`))
case "2":
_, _ = w.Write([]byte(`[{"name":"two.md","type":"blob","path":"two.md"}]`))
default:
http.Error(w, "unexpected page", http.StatusBadRequest)
}
}))
defer server.Close()
c, err := newClient(server.URL, "token")
if err != nil {
t.Fatal(err)
}
entries, err := c.tree(context.Background(), "1", "main", "")
if err != nil {
t.Fatalf("tree() error = %v", err)
}
if len(entries) != 2 || entries[1].Path != "two.md" {
t.Fatalf("entries = %#v", entries)
}
if got, want := strings.Join(requestedPages, ","), "1,2"; got != want {
t.Fatalf("requested pages = %q, want %q", got, want)
}
}
func TestProjectPathEncodesNamespaceWithoutDoubleEscaping(t *testing.T) {
for _, tc := range []struct {
in, want string
}{
{in: "12345", want: "12345"},
{in: "group/project", want: "group%2Fproject"},
{in: "group%2Fproject", want: "group%2Fproject"},
{in: "my group/my project", want: "my%20group%2Fmy%20project"},
} {
if got := projectPath(tc.in); got != tc.want {
t.Errorf("projectPath(%q) = %q, want %q", tc.in, got, tc.want)
}
}
}
func TestRawFallsBackToBase64FileDetail(t *testing.T) {
t.Setenv("SSRF_WHITELIST", "127.0.0.1,::1,localhost")
secutils.ResetSSRFWhitelistForTest()
t.Cleanup(secutils.ResetSSRFWhitelistForTest)
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
switch r.URL.EscapedPath() {
case "/api/v4/projects/18724/repository/files/docs%2Finternal%2Freadme%2Emd/raw":
http.NotFound(w, r)
case "/api/v4/projects/18724/repository/files/docs%2Finternal%2Freadme%2Emd":
w.Header().Set("Content-Type", "application/json")
_, _ = w.Write([]byte(`{"encoding":"base64","content":"SGVsbG8sIEdpdExhYiE="}`))
default:
http.NotFound(w, r)
}
}))
defer server.Close()
c, err := newClient(server.URL, "token")
if err != nil {
t.Fatal(err)
}
content, err := c.raw(context.Background(), "18724", "master", "docs/internal/readme.md")
if err != nil {
t.Fatal(err)
}
if string(content) != "Hello, GitLab!" {
t.Fatalf("content = %q", content)
}
}
func TestGitlabFilePathEscape(t *testing.T) {
got := gitlabFilePathEscape("docs/internal/中文-file.md")
want := "docs%2Finternal%2F%E4%B8%AD%E6%96%87-file%2Emd"
if got != want {
t.Fatalf("gitlabFilePathEscape() = %q, want %q", got, want)
}
}
@@ -0,0 +1,436 @@
package gitlab
import (
"context"
"encoding/json"
"fmt"
"path"
"strings"
"time"
"github.com/Tencent/WeKnora/internal/datasource"
"github.com/Tencent/WeKnora/internal/types"
)
var _ datasource.StreamingConnector = (*Connector)(nil)
type Connector struct {
client *client
canonicalBase string
}
// NewConnector creates a stateless connector. Each data source provides its
// own GitLab URL and access token in its encrypted credentials.
func NewConnector() *Connector { return &Connector{} }
func (c *Connector) Type() string { return types.ConnectorTypeGitLab }
func (c *Connector) configured(ds *types.DataSourceConfig) (*Connector, error) {
if ds == nil {
return nil, datasource.ErrInvalidConfig
}
baseURL, _ := ds.Credentials["base_url"].(string)
token, _ := ds.Credentials["access_token"].(string)
client, err := newClient(baseURL, token)
if err != nil {
return nil, err
}
return &Connector{client: client, canonicalBase: client.baseURL}, nil
}
func (c *Connector) Validate(ctx context.Context, ds *types.DataSourceConfig) error {
configured, err := c.configured(ds)
if err != nil {
return err
}
if err := configured.client.ping(ctx); err != nil {
return err
}
if ds != nil {
if _, ok := ds.Settings["projects"]; ok {
if _, err := parseConfig(ds); err != nil {
return err
}
}
}
return nil
}
func (c *Connector) ListResources(ctx context.Context, ds *types.DataSourceConfig, parent string) ([]types.Resource, error) {
var err error
if c, err = c.configured(ds); err != nil {
return nil, err
}
if _, err := parseConfig(ds); err != nil {
return nil, err
}
if parent == "" {
ps, err := c.client.projects(ctx)
if err != nil {
return nil, err
}
out := make([]types.Resource, 0, len(ps))
for _, p := range ps {
out = append(out, types.Resource{ExternalID: fmt.Sprint(p.ID), Name: p.PathWithNamespace, Type: "project", URL: p.WebURL, HasChildren: true})
}
return out, nil
}
id, dir := splitResourceID(parent)
p, err := c.client.project(ctx, id)
if err != nil {
return nil, err
}
ref := p.DefaultBranch
entries, err := c.client.tree(ctx, id, ref, dir)
if err != nil {
return nil, err
}
out := make([]types.Resource, 0, len(entries))
for _, e := range entries {
if e.Type == "tree" {
out = append(out, types.Resource{ExternalID: id + ":" + e.Path, Name: e.Name, Type: "directory", ParentID: parent, HasChildren: true})
}
}
return out, nil
}
func (c *Connector) ResolveResourceAncestors(context.Context, *types.DataSourceConfig, []string) ([]string, error) {
return []string{}, nil
}
func (c *Connector) FetchAll(ctx context.Context, ds *types.DataSourceConfig, _ []string) ([]types.FetchedItem, error) {
var err error
if c, err = c.configured(ds); err != nil {
return nil, err
}
cfg, err := parseConfig(ds)
if err != nil {
return nil, err
}
var out []types.FetchedItem
for _, s := range cfg.Projects {
p, err := c.client.project(ctx, s.ProjectID)
if err != nil {
return nil, err
}
ref := s.Ref
if ref == "" {
ref = p.DefaultBranch
}
files, err := c.files(ctx, s.ProjectID, ref, s.Paths)
if err != nil {
return nil, err
}
for _, file := range files {
item, err := c.item(ctx, p, ref, file)
if err != nil {
return nil, err
}
out = append(out, item)
}
}
return out, nil
}
type cursor struct {
Projects map[string]string `json:"projects"`
}
func (c *Connector) FetchIncremental(ctx context.Context, ds *types.DataSourceConfig, old *types.SyncCursor) ([]types.FetchedItem, *types.SyncCursor, error) {
var err error
if c, err = c.configured(ds); err != nil {
return nil, nil, err
}
cfg, err := parseConfig(ds)
if err != nil {
return nil, nil, err
}
prev := cursor{Projects: map[string]string{}}
if old != nil {
b, _ := json.Marshal(old.ConnectorCursor)
_ = json.Unmarshal(b, &prev)
if prev.Projects == nil {
prev.Projects = map[string]string{}
}
}
var out []types.FetchedItem
next := cursor{Projects: map[string]string{}}
for _, s := range cfg.Projects {
p, err := c.client.project(ctx, s.ProjectID)
if err != nil {
return nil, nil, err
}
ref := s.Ref
if ref == "" {
ref = p.DefaultBranch
}
head, err := c.client.commitSHA(ctx, s.ProjectID, ref)
if err != nil {
return nil, nil, err
}
previous := prev.Projects[s.ProjectID]
if previous == "" {
files, err := c.files(ctx, s.ProjectID, ref, s.Paths)
if err != nil {
return nil, nil, err
}
for _, f := range files {
item, err := c.item(ctx, p, ref, f)
if err != nil {
return nil, nil, err
}
out = append(out, item)
}
} else if previous != head {
diff, err := c.client.compare(ctx, s.ProjectID, previous, head)
if err != nil || diff.CompareTimeout {
// A compare may be unavailable after history rewrites or truncated by
// the server. Re-enumerate the configured scope to preserve updates.
files, listErr := c.files(ctx, s.ProjectID, ref, s.Paths)
if listErr != nil {
return nil, nil, fmt.Errorf("gitlab list files %s: %w", s.ProjectID, listErr)
}
for _, f := range files {
item, itemErr := c.item(ctx, p, ref, f)
if itemErr != nil {
return nil, nil, itemErr
}
out = append(out, item)
}
} else {
for _, d := range diff.Diffs {
if d.DeletedFile {
if c.inScope(d.OldPath, s.Paths) && isSupportedFile(d.OldPath) {
out = append(out, c.deleted(p, ref, d.OldPath))
}
continue
}
if d.RenamedFile && c.inScope(d.OldPath, s.Paths) && isSupportedFile(d.OldPath) {
out = append(out, c.deleted(p, ref, d.OldPath))
}
if c.inScope(d.NewPath, s.Paths) && isSupportedFile(d.NewPath) {
item, err := c.item(ctx, p, ref, d.NewPath)
if err != nil {
return nil, nil, err
}
out = append(out, item)
}
}
}
}
next.Projects[s.ProjectID] = head
}
raw, _ := json.Marshal(next)
return out, &types.SyncCursor{LastSyncTime: time.Now().UTC(), ConnectorCursor: map[string]interface{}{"projects": next.Projects, "raw": string(raw)}}, nil
}
// FetchStream is the production sync path. It emits each supported repository
// file as soon as it is read, so large GitLab projects do not accumulate their
// contents in memory. The service prefers this method for StreamingConnector
// implementations; FetchAll and FetchIncremental remain compatibility methods
// required by the base Connector interface.
func (c *Connector) FetchStream(
ctx context.Context, ds *types.DataSourceConfig, old *types.SyncCursor, h datasource.StreamHandler,
) (*types.SyncCursor, error) {
var err error
if c, err = c.configured(ds); err != nil {
return nil, err
}
cfg, err := parseConfig(ds)
if err != nil {
return nil, err
}
prev := cursor{Projects: map[string]string{}}
if old != nil {
b, _ := json.Marshal(old.ConnectorCursor)
_ = json.Unmarshal(b, &prev)
if prev.Projects == nil {
prev.Projects = map[string]string{}
}
}
next := cursor{Projects: make(map[string]string, len(cfg.Projects))}
for projectID, commit := range prev.Projects {
next.Projects[projectID] = commit
}
for _, selection := range cfg.Projects {
project, err := c.client.project(ctx, selection.ProjectID)
if err != nil {
return nil, err
}
ref := selection.Ref
if ref == "" {
ref = project.DefaultBranch
}
head, err := c.client.commitSHA(ctx, selection.ProjectID, ref)
if err != nil {
return nil, err
}
previous := prev.Projects[selection.ProjectID]
switch {
case previous == "":
err = c.streamFiles(ctx, project, ref, selection.Paths, h)
case previous != head:
err = c.streamChanges(ctx, project, ref, previous, head, selection.Paths, h)
}
if err != nil {
return nil, err
}
next.Projects[selection.ProjectID] = head
checkpoint := gitLabCursor(next)
if err := h.Checkpoint(ctx, checkpoint); err != nil {
return nil, err
}
}
return gitLabCursor(next), nil
}
func gitLabCursor(value cursor) *types.SyncCursor {
raw, _ := json.Marshal(value)
return &types.SyncCursor{
LastSyncTime: time.Now().UTC(),
ConnectorCursor: map[string]interface{}{"projects": value.Projects, "raw": string(raw)},
}
}
func (c *Connector) streamChanges(
ctx context.Context, project *project, ref, from, to string, roots []string, h datasource.StreamHandler,
) error {
diff, err := c.client.compare(ctx, fmt.Sprint(project.ID), from, to)
if err != nil || diff.CompareTimeout {
// A compare can be unavailable after history rewrites or be truncated by
// GitLab. Re-enumerating the configured scope preserves file updates.
return c.streamFiles(ctx, project, ref, roots, h)
}
for _, change := range diff.Diffs {
if change.DeletedFile {
if c.inScope(change.OldPath, roots) && isSupportedFile(change.OldPath) {
if err := h.Emit(ctx, c.deleted(project, ref, change.OldPath)); err != nil {
return err
}
}
continue
}
if change.RenamedFile && c.inScope(change.OldPath, roots) && isSupportedFile(change.OldPath) {
if err := h.Emit(ctx, c.deleted(project, ref, change.OldPath)); err != nil {
return err
}
}
if c.inScope(change.NewPath, roots) && isSupportedFile(change.NewPath) {
item, err := c.item(ctx, project, ref, change.NewPath)
if err != nil {
return err
}
if err := h.Emit(ctx, item); err != nil {
return err
}
}
}
return nil
}
func (c *Connector) streamFiles(
ctx context.Context, project *project, ref string, roots []string, h datasource.StreamHandler,
) error {
return c.walkFiles(ctx, fmt.Sprint(project.ID), ref, roots, func(file string) error {
item, err := c.item(ctx, project, ref, file)
if err != nil {
return err
}
return h.Emit(ctx, item)
})
}
func (c *Connector) files(ctx context.Context, id, ref string, roots []string) ([]string, error) {
var out []string
err := c.walkFiles(ctx, id, ref, roots, func(file string) error {
out = append(out, file)
return nil
})
return out, err
}
// walkFiles visits supported blobs while traversing the selected directories.
// It deliberately does not collect paths, keeping streaming sync memory bounded
// by the traversal depth plus the current file body.
func (c *Connector) walkFiles(ctx context.Context, id, ref string, roots []string, visit func(string) error) error {
if len(roots) == 0 {
roots = []string{""}
}
var walk func(string) error
walk = func(dir string) error {
entries, err := c.client.tree(ctx, id, ref, dir)
if err != nil {
return err
}
for _, e := range entries {
if e.Type == "tree" {
if err := walk(e.Path); err != nil {
return err
}
} else if e.Type == "blob" && isSupportedFile(e.Path) {
if err := visit(e.Path); err != nil {
return err
}
}
}
return nil
}
for _, r := range roots {
if err := walk(r); err != nil {
return err
}
}
return nil
}
func isSupportedFile(file string) bool {
_, ok := gitLabSupportedFileExtensions[strings.ToLower(path.Ext(file))]
return ok
}
// gitLabSupportedFileExtensions limits repository sync to formats the current
// knowledge import pipeline can process. This is intentionally connector-local:
// GitLab exposes arbitrary repository blobs, unlike document-centric sources.
var gitLabSupportedFileExtensions = map[string]struct{}{
".pdf": {}, ".txt": {}, ".docx": {}, ".doc": {}, ".epub": {},
".html": {}, ".htm": {}, ".mhtml": {}, ".md": {}, ".markdown": {}, ".mdx": {},
".png": {}, ".jpg": {}, ".jpeg": {}, ".gif": {},
".csv": {}, ".xlsx": {}, ".xls": {}, ".pptx": {}, ".ppt": {}, ".json": {},
".mp3": {}, ".wav": {}, ".m4a": {}, ".flac": {}, ".ogg": {},
}
func (c *Connector) item(ctx context.Context, p *project, ref, file string) (types.FetchedItem, error) {
body, err := c.client.raw(ctx, fmt.Sprint(p.ID), ref, file)
if err != nil {
return types.FetchedItem{}, err
}
id := fmt.Sprintf("gitlab:%s:%d:%s:%s", c.canonicalBase, p.ID, ref, file)
return types.FetchedItem{ExternalID: id, Title: p.PathWithNamespace + "/" + file, FileName: knowledgeRelativePath(p.Name, ref, file), Content: body, ContentType: "text/plain", UpdatedAt: time.Now().UTC(), SourceResourceID: fmt.Sprint(p.ID), Metadata: map[string]string{"channel": types.ConnectorTypeGitLab, "source_type": "gitlab", "gitlab_project_id": fmt.Sprint(p.ID), "gitlab_ref": ref, "gitlab_path": file, "gitlab_url": p.WebURL + "/-/blob/" + ref + "/" + file}}, nil
}
// knowledgeRelativePath maps a repository file to the KB folder convention:
// <GitLab project name>-<branch>/<repository-relative file path>.
func knowledgeRelativePath(projectName, ref, file string) string {
root := strings.TrimSpace(projectName) + "-" + strings.ReplaceAll(strings.TrimSpace(ref), "/", "-")
return path.Join(root, file)
}
func (c *Connector) deleted(p *project, ref, file string) types.FetchedItem {
return types.FetchedItem{ExternalID: fmt.Sprintf("gitlab:%s:%d:%s:%s", c.canonicalBase, p.ID, ref, file), IsDeleted: true, Metadata: map[string]string{"channel": types.ConnectorTypeGitLab, "gitlab_path": file}}
}
func (c *Connector) inScope(file string, roots []string) bool {
if len(roots) == 0 {
return true
}
for _, r := range roots {
if file == r || strings.HasPrefix(file, r+"/") {
return true
}
}
return false
}
func splitResourceID(value string) (string, string) {
parts := strings.SplitN(value, ":", 2)
if len(parts) == 1 {
return value, ""
}
return parts[0], parts[1]
}
@@ -0,0 +1,108 @@
package gitlab
import (
"fmt"
"path"
"sort"
"strings"
"github.com/Tencent/WeKnora/internal/datasource"
"github.com/Tencent/WeKnora/internal/types"
)
type projectSelection struct {
ProjectID string `json:"project_id"`
Ref string `json:"ref"`
Paths []string `json:"paths"`
}
type config struct {
Projects []projectSelection `json:"projects"`
}
func parseConfig(ds *types.DataSourceConfig) (*config, error) {
if ds == nil {
return nil, datasource.ErrInvalidConfig
}
raw, ok := ds.Settings["projects"]
if !ok {
return nil, fmt.Errorf("%w: settings.projects is required", datasource.ErrInvalidConfig)
}
items, ok := raw.([]interface{})
if !ok {
return nil, fmt.Errorf("%w: settings.projects must be an array", datasource.ErrInvalidConfig)
}
out := &config{Projects: make([]projectSelection, 0, len(items))}
seen := map[string]bool{}
for _, rawProject := range items {
m, ok := rawProject.(map[string]interface{})
if !ok {
return nil, fmt.Errorf("%w: invalid project selection", datasource.ErrInvalidConfig)
}
id, _ := m["project_id"].(string)
id = strings.TrimSpace(id)
if id == "" || seen[id] {
return nil, fmt.Errorf("%w: project_id must be unique and non-empty", datasource.ErrInvalidConfig)
}
seen[id] = true
ref, _ := m["ref"].(string)
p := projectSelection{ProjectID: id, Ref: strings.TrimSpace(ref)}
if rawPaths, exists := m["paths"]; exists {
values, ok := rawPaths.([]interface{})
if !ok {
return nil, fmt.Errorf("%w: paths must be an array", datasource.ErrInvalidConfig)
}
for _, value := range values {
s, ok := value.(string)
if !ok {
return nil, fmt.Errorf("%w: path must be a string", datasource.ErrInvalidConfig)
}
normalized, err := normalizePath(s)
if err != nil {
return nil, err
}
p.Paths = append(p.Paths, normalized)
}
}
p.Paths = collapsePaths(p.Paths)
out.Projects = append(out.Projects, p)
}
if len(out.Projects) == 0 {
return nil, fmt.Errorf("%w: at least one project is required", datasource.ErrInvalidConfig)
}
return out, nil
}
func normalizePath(value string) (string, error) {
v := strings.Trim(strings.TrimSpace(value), "/")
if v == "" {
return "", nil
}
if strings.Contains(v, "\\") {
return "", fmt.Errorf("%w: path must use forward slashes", datasource.ErrInvalidConfig)
}
if path.Clean(v) != v || v == "." || strings.HasPrefix(v, "../") || strings.Contains(v, "/../") {
return "", fmt.Errorf("%w: invalid repository path", datasource.ErrInvalidConfig)
}
return v, nil
}
func collapsePaths(paths []string) []string {
if len(paths) == 0 {
return nil
}
sort.Strings(paths)
for _, p := range paths {
if p == "" {
return nil
}
}
out := make([]string, 0, len(paths))
for _, p := range paths {
if len(out) > 0 && (p == out[len(out)-1] || strings.HasPrefix(p, out[len(out)-1]+"/")) {
continue
}
out = append(out, p)
}
return out
}
@@ -0,0 +1,47 @@
package gitlab
import (
"testing"
"github.com/Tencent/WeKnora/internal/types"
)
func TestParseConfigCollapsesDirectories(t *testing.T) {
cfg, err := parseConfig(&types.DataSourceConfig{Settings: map[string]interface{}{
"projects": []interface{}{map[string]interface{}{
"project_id": "team/docs", "ref": "main",
"paths": []interface{}{"docs/guide", "docs", "docs"},
}},
}})
if err != nil {
t.Fatal(err)
}
if len(cfg.Projects) != 1 || len(cfg.Projects[0].Paths) != 1 || cfg.Projects[0].Paths[0] != "docs" {
t.Fatalf("unexpected config: %#v", cfg)
}
}
func TestParseConfigRootMeansWholeProject(t *testing.T) {
cfg, err := parseConfig(&types.DataSourceConfig{Settings: map[string]interface{}{
"projects": []interface{}{map[string]interface{}{"project_id": "42", "paths": []interface{}{"/"}}},
}})
if err != nil {
t.Fatal(err)
}
if len(cfg.Projects[0].Paths) != 0 {
t.Fatalf("root path must collapse to whole project: %#v", cfg.Projects[0].Paths)
}
}
func TestNormalizePathRejectsTraversal(t *testing.T) {
if _, err := normalizePath("docs/../secrets"); err == nil {
t.Fatal("expected traversal rejection")
}
}
func TestKnowledgeRelativePathPreservesRepositoryTreeBelowProjectAndBranch(t *testing.T) {
got := knowledgeRelativePath("knowledge", "feature/login", "docs/guide/install.md")
if got != "knowledge-feature-login/docs/guide/install.md" {
t.Fatalf("knowledge relative path = %q", got)
}
}
+403
View File
@@ -0,0 +1,403 @@
// Package ima implements the WeKnora data source connector for Tencent IMA
// (ima.qq.com), syncing documents and files out of IMA knowledge bases via the
// /openapi/wiki/v1 OpenAPI.
package ima
import (
"bytes"
"context"
"encoding/json"
"fmt"
"io"
"net/http"
"sync"
"time"
"github.com/Tencent/WeKnora/internal/datasource"
"github.com/Tencent/WeKnora/internal/logger"
"github.com/Tencent/WeKnora/internal/utils"
)
const (
defaultTimeout = 60 * time.Second
defaultPageSize = 50
userAgent = "WeKnora-IMA-Connector/1.0"
// IMA hard limit: get_knowledge_list max limit=50; search_knowledge_base max=20.
searchPageSize = 20
// Default per-item download timeout — IMA COS URLs can be slow for large PDFs.
downloadTimeout = 120 * time.Second
// Max body we accept from get_media_info-URLs. 200MB matches IMA's largest
maxDownloadBytes = 200 * 1024 * 1024
)
// client wraps the IMA OpenAPI. It is safe for concurrent use.
type client struct {
baseURL string
clientID string
apiKey string
httpClient *http.Client
// downloadClient is a separate client with a longer timeout for pulling
// media bodies from COS/CDN URLs returned by get_media_info.
downloadClient *http.Client
logCredentialOnce sync.Once
}
// newClient constructs a client with a normalized base URL.
func newClient(cfg *Config) *client {
return &client{
baseURL: cfg.GetBaseURL(),
clientID: cfg.ClientID,
apiKey: cfg.APIKey,
httpClient: datasource.NewConnectorHTTPClient(defaultTimeout),
downloadClient: datasource.NewConnectorHTTPClient(downloadTimeout),
}
}
// callAPI executes an authenticated POST to /openapi/wiki/v1/<action>.
func (c *client) callAPI(ctx context.Context, action string, req interface{}, result interface{}) error {
return c.callAPIAt(ctx, apiBasePath, action, req, result)
}
// callAPIAt executes an authenticated POST to <basePath>/<action>. It parses
// the standard {code, msg, data} envelope and unmarshals `data` into `result`.
// A non-zero business code is returned as a plain error so the caller can
// surface `msg` to the user.
//
// basePath is a parameter because notes live under a second namespace
// (/openapi/note/v1) that shares this authentication and envelope.
func (c *client) callAPIAt(
ctx context.Context, basePath, action string, req interface{}, result interface{},
) error {
const (
maxRetries = 3
max5xxRetries = 1
retry5xxDelay = 2 * time.Second
)
backoff := []time.Duration{2 * time.Second, 4 * time.Second, 8 * time.Second}
c.logCredentialOnce.Do(func() {
logger.Infof(ctx, "[IMA] client configured client_id=%s api_key=%s base=%s",
redact(c.clientID), redact(c.apiKey), c.baseURL)
})
var body []byte
if req != nil {
b, err := json.Marshal(req)
if err != nil {
return fmt.Errorf("marshal request: %w", err)
}
body = b
} else {
body = []byte("{}")
}
url := c.baseURL + basePath + "/" + action
// Both namespaces expose distinct actions, but log the full path so the two
// are never ambiguous.
label := basePath + "/" + action
var lastErr error
for attempt := 0; attempt <= maxRetries; attempt++ {
httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(body))
if err != nil {
return fmt.Errorf("create request: %w", err)
}
httpReq.Header.Set("ima-openapi-clientid", c.clientID)
httpReq.Header.Set("ima-openapi-apikey", c.apiKey)
httpReq.Header.Set("Content-Type", "application/json; charset=utf-8")
httpReq.Header.Set("User-Agent", userAgent)
if attempt == 0 {
logger.Infof(ctx, "[IMA] POST %s", label)
} else {
logger.Infof(ctx, "[IMA] POST %s (retry %d/%d)", label, attempt, maxRetries)
}
resp, err := c.httpClient.Do(httpReq)
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: %w", readErr)
if attempt < maxRetries {
if sErr := sleepCtx(ctx, backoff[attempt]); sErr != nil {
return sErr
}
continue
}
return lastErr
}
// The body is only ever logged on an error path: successful responses
// (get_media_info in particular) carry pre-signed storage URLs and
// per-request auth headers that must not land in the server log.
bodyPreview := truncate(string(respBody), 500)
logger.Infof(ctx, "[IMA] POST %s → status=%d bodyLen=%d", label, resp.StatusCode, len(respBody))
// HTTP-level auth failures: surface as ErrInvalidCredentials.
if resp.StatusCode == http.StatusUnauthorized || resp.StatusCode == http.StatusForbidden {
return fmt.Errorf("%w: status=%d body=%s",
datasource.ErrInvalidCredentials, resp.StatusCode, bodyPreview)
}
if resp.StatusCode == http.StatusTooManyRequests {
wait := backoff[minInt(attempt, len(backoff)-1)]
lastErr = fmt.Errorf("ima rate limited: status=429 body=%s", bodyPreview)
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("ima server error: status=%d body=%s", resp.StatusCode, bodyPreview)
if attempt < max5xxRetries {
if sErr := sleepCtx(ctx, retry5xxDelay); sErr != nil {
return sErr
}
continue
}
return lastErr
}
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
return fmt.Errorf("ima api http error: status=%d body=%s", resp.StatusCode, bodyPreview)
}
var env apiEnvelope
if err := json.Unmarshal(respBody, &env); err != nil {
return fmt.Errorf("decode envelope: %w", err)
}
// Business-level errors.
// 110030 (无权限) is mapped to ErrInvalidCredentials so the service marks
// the source in `error` state and stops scheduling until re-authenticated.
// 110021 (限频) is retried. Everything else — including 110001 (参数非法),
// which is a bug on our side rather than a credential problem — surfaces
// verbatim so the user sees IMA's own message.
if code := env.statusCode(); code != 0 {
if code == 110030 {
return fmt.Errorf("%w: ima code=%d msg=%s",
datasource.ErrInvalidCredentials, code, env.message())
}
if code == 110021 && attempt < maxRetries {
lastErr = fmt.Errorf("ima rate limited: code=%d msg=%s", code, env.message())
if sErr := sleepCtx(ctx, backoff[minInt(attempt, len(backoff)-1)]); sErr != nil {
return sErr
}
continue
}
return fmt.Errorf("ima api error: code=%d msg=%s", code, env.message())
}
if result != nil && len(env.Data) > 0 && string(env.Data) != "null" {
if err := json.Unmarshal(env.Data, result); err != nil {
return fmt.Errorf("decode data: %w", err)
}
}
return nil
}
return lastErr
}
// SearchKnowledgeBase — POST /search_knowledge_base. Empty query returns the
// list of every knowledge base visible to the current token (see api.md §7).
func (c *client) SearchKnowledgeBase(
ctx context.Context, query, cursor string, limit int,
) (*searchKnowledgeBaseResp, error) {
if limit <= 0 || limit > searchPageSize {
limit = searchPageSize
}
req := map[string]interface{}{
"query": query,
"cursor": cursor,
"limit": uint64(limit),
}
var resp searchKnowledgeBaseResp
if err := c.callAPI(ctx, "search_knowledge_base", req, &resp); err != nil {
return nil, err
}
return &resp, nil
}
// GetKnowledgeBase — POST /get_knowledge_base for a batch of ids (1-20).
func (c *client) GetKnowledgeBase(ctx context.Context, ids []string) (map[string]knowledgeBaseInfo, error) {
if len(ids) == 0 {
return map[string]knowledgeBaseInfo{}, nil
}
req := map[string]interface{}{"ids": ids}
var resp getKnowledgeBaseResp
if err := c.callAPI(ctx, "get_knowledge_base", req, &resp); err != nil {
return nil, err
}
if resp.Infos == nil {
resp.Infos = map[string]knowledgeBaseInfo{}
}
return resp.Infos, nil
}
// GetAddableKnowledgeBaseList — POST /get_addable_knowledge_base_list.
// Lists knowledge bases the current token has permission to add content to.
// This is the authoritative "what KBs can this credential see?" endpoint
// (see api.md §7). Used as the primary enumeration source in ListResources —
// search_knowledge_base with empty query has been observed to return an empty
// list even when the token owns KBs, so we prefer this endpoint.
func (c *client) GetAddableKnowledgeBaseList(
ctx context.Context, cursor string, limit int,
) (*getAddableKnowledgeBaseListResp, error) {
if limit <= 0 || limit > defaultPageSize {
limit = defaultPageSize
}
req := map[string]interface{}{
"cursor": cursor,
"limit": uint64(limit),
}
var resp getAddableKnowledgeBaseListResp
if err := c.callAPI(ctx, "get_addable_knowledge_base_list", req, &resp); err != nil {
return nil, err
}
return &resp, nil
}
// GetKnowledgeList — POST /get_knowledge_list. folderID may be empty for root.
func (c *client) GetKnowledgeList(
ctx context.Context, kbID, folderID, cursor string, limit int,
) (*getKnowledgeListResp, error) {
if limit <= 0 || limit > defaultPageSize {
limit = defaultPageSize
}
req := map[string]interface{}{
"cursor": cursor,
"limit": uint64(limit),
"knowledge_base_id": kbID,
}
if folderID != "" {
req["folder_id"] = folderID
}
var resp getKnowledgeListResp
if err := c.callAPI(ctx, "get_knowledge_list", req, &resp); err != nil {
return nil, err
}
return &resp, nil
}
// GetMediaInfo — POST /get_media_info. Returns URL access info and (for notes)
// the notebook_id under notebook_ext_info.
func (c *client) GetMediaInfo(ctx context.Context, mediaID string) (*getMediaInfoResp, error) {
req := map[string]interface{}{"media_id": mediaID}
var resp getMediaInfoResp
if err := c.callAPI(ctx, "get_media_info", req, &resp); err != nil {
return nil, err
}
return &resp, nil
}
// GetNoteContent — POST /openapi/note/v1/get_doc_content. Returns an IMA note's
// body as plain text (target_content_format=0).
//
// noteID is the notebook_id that get_media_info reports for a note: the wiki
// namespace never exposes note bodies, it only points at them.
func (c *client) GetNoteContent(ctx context.Context, noteID string) (string, error) {
req := map[string]interface{}{
"note_id": noteID,
"target_content_format": 0,
}
var resp getDocContentResp
if err := c.callAPIAt(ctx, noteBasePath, "get_doc_content", req, &resp); err != nil {
return "", err
}
return resp.Content, nil
}
// DownloadURL fetches a URL returned by GetMediaInfo, respecting the auth
// headers IMA may return alongside it. Enforces maxDownloadBytes to avoid a
// runaway body blowing up sync memory.
func (c *client) DownloadURL(ctx context.Context, u urlInfo) ([]byte, string, error) {
if u.URL == "" {
return nil, "", fmt.Errorf("empty url")
}
// The URL comes from an API response, so it is attacker-influenced as far
// as this process is concerned. downloadClient already guards the dial and
// every redirect hop; this rejects an unsafe target before the first
// connection is even attempted.
if err := utils.ValidateURLForSSRF(u.URL); err != nil {
return nil, "", fmt.Errorf("media URL rejected: %w", err)
}
req, err := http.NewRequestWithContext(ctx, http.MethodGet, u.URL, nil)
if err != nil {
return nil, "", fmt.Errorf("create download request: %w", err)
}
for k, v := range u.Headers {
req.Header.Set(k, v)
}
req.Header.Set("User-Agent", userAgent)
resp, err := c.downloadClient.Do(req)
if err != nil {
return nil, "", fmt.Errorf("download: %w", err)
}
defer func() { _ = resp.Body.Close() }()
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
return nil, "", fmt.Errorf("download http error: status=%d", resp.StatusCode)
}
body, err := io.ReadAll(io.LimitReader(resp.Body, maxDownloadBytes+1))
if err != nil {
return nil, "", fmt.Errorf("read download body: %w", err)
}
if int64(len(body)) > maxDownloadBytes {
return nil, "", fmt.Errorf("download body exceeds %d bytes", maxDownloadBytes)
}
return body, resp.Header.Get("Content-Type"), nil
}
// sleepCtx pauses for d, returning early if ctx is cancelled.
func sleepCtx(ctx context.Context, d time.Duration) error {
t := time.NewTimer(d)
defer t.Stop()
select {
case <-t.C:
return nil
case <-ctx.Done():
return ctx.Err()
}
}
func truncate(s string, maxLen int) string {
if len(s) <= maxLen {
return s
}
return s[:maxLen] + "..."
}
// redact returns a masked form of a credential for logging (never log full).
func redact(t string) string {
if len(t) < 12 {
return "***"
}
return t[:6] + "..." + t[len(t)-4:]
}
func minInt(a, b int) int {
if a < b {
return a
}
return b
}
@@ -0,0 +1,612 @@
package ima
import (
"context"
"encoding/json"
"fmt"
"sort"
"strings"
"time"
"unicode/utf8"
"github.com/Tencent/WeKnora/internal/datasource"
"github.com/Tencent/WeKnora/internal/logger"
"github.com/Tencent/WeKnora/internal/types"
)
// Compile-time proof that *Connector satisfies the datasource.Connector interface.
var _ datasource.Connector = (*Connector)(nil)
// Connector implements datasource.Connector for Tencent IMA (ima.qq.com).
type Connector struct{}
// NewConnector creates a new IMA connector.
func NewConnector() *Connector { return &Connector{} }
// Type returns the connector type identifier.
func (c *Connector) Type() string { return types.ConnectorTypeIMA }
// Validate verifies the credentials by calling get_addable_knowledge_base_list
// — the endpoint most likely to succeed even when the token has zero KBs, and
// the same one ListResources uses. It returns 110030 (无权限) when the token
// itself is invalid, which client.callAPI already maps to ErrInvalidCredentials.
func (c *Connector) Validate(ctx context.Context, config *types.DataSourceConfig) error {
cfg, err := parseIMAConfig(config)
if err != nil {
return err
}
cli := newClient(cfg)
if _, err := cli.GetAddableKnowledgeBaseList(ctx, "", defaultPageSize); err != nil {
return fmt.Errorf("ima connection failed: %w", err)
}
return nil
}
// ResolveResourceAncestors returns nothing: IMA knowledge bases are exposed as
// a flat list of top-level resources, so a lazy-loaded picker has no ancestors
// to reveal.
func (c *Connector) ResolveResourceAncestors(
_ context.Context, _ *types.DataSourceConfig, _ []string,
) ([]string, error) {
return []string{}, nil
}
// ListResources returns the flat list of knowledge bases the token can read.
// parentID is honoured only for the "no children" contract — see the note on
// Connector.ListResources in internal/datasource/connector.go.
//
// Primary source is get_addable_knowledge_base_list, which returns the KBs the
// current OpenAPI credential has permission to operate on. When that endpoint
// returns an empty list (e.g. a legacy tenant only exposes read scopes) we
// fall back to search_knowledge_base with an empty query so users still see
// something to pick.
func (c *Connector) ListResources(
ctx context.Context, config *types.DataSourceConfig, parentID string,
) ([]types.Resource, error) {
if parentID != "" {
return []types.Resource{}, nil
}
cfg, err := parseIMAConfig(config)
if err != nil {
return nil, err
}
cli := newClient(cfg)
type kbLite struct {
ID string
Name string
CoverURL string
}
var bases []kbLite
cursor := ""
for {
resp, err := cli.GetAddableKnowledgeBaseList(ctx, cursor, defaultPageSize)
if err != nil {
return nil, fmt.Errorf("get_addable_knowledge_base_list: %w", err)
}
for _, b := range resp.AddableKnowledgeBaseList {
bases = append(bases, kbLite{ID: b.ID, Name: b.Name})
}
if resp.IsEnd || resp.NextCursor == "" {
break
}
cursor = resp.NextCursor
}
logger.Infof(ctx, "[IMA] get_addable_knowledge_base_list returned %d knowledge bases", len(bases))
if len(bases) == 0 {
cursor = ""
for {
resp, err := cli.SearchKnowledgeBase(ctx, "", cursor, searchPageSize)
if err != nil {
return nil, fmt.Errorf("search_knowledge_base fallback: %w", err)
}
for _, b := range resp.InfoList {
bases = append(bases, kbLite(b))
}
if resp.IsEnd || resp.NextCursor == "" {
break
}
cursor = resp.NextCursor
}
logger.Infof(ctx, "[IMA] search_knowledge_base fallback returned %d knowledge bases", len(bases))
}
ids := make([]string, 0, len(bases))
for _, b := range bases {
ids = append(ids, b.ID)
}
details := map[string]knowledgeBaseInfo{}
for i := 0; i < len(ids); i += 20 {
end := i + 20
if end > len(ids) {
end = len(ids)
}
batch, err := cli.GetKnowledgeBase(ctx, ids[i:end])
if err != nil {
logger.Warnf(ctx, "[IMA] get_knowledge_base batch failed (skipping enrichment): %v", err)
continue
}
for k, v := range batch {
details[k] = v
}
}
out := make([]types.Resource, 0, len(bases))
for _, b := range bases {
desc := ""
coverURL := b.CoverURL
if d, ok := details[b.ID]; ok {
desc = d.Description
if coverURL == "" {
coverURL = d.CoverURL
}
}
out = append(out, types.Resource{
ExternalID: b.ID,
Name: b.Name,
Type: "knowledge_base",
Description: desc,
URL: cfg.GetBaseURL(),
Metadata: map[string]interface{}{
"cover_url": coverURL,
},
})
}
sort.Slice(out, func(i, j int) bool { return out[i].ExternalID < out[j].ExternalID })
logger.Infof(ctx, "[IMA] ListResources returning %d knowledge bases to UI", len(out))
return out, nil
}
// FetchAll performs a full sync of every knowledge base in resourceIDs.
func (c *Connector) FetchAll(
ctx context.Context, config *types.DataSourceConfig, resourceIDs []string,
) ([]types.FetchedItem, error) {
items, _, err := c.walk(ctx, config, resourceIDs, nil, false)
return items, err
}
// FetchIncremental syncs new / replaced / removed items against the cursor.
//
// Each IMA item is identified by a stable logical_key (see logicalKey), NOT by
// media_id — because IMA reassigns media_id whenever a same-named file is
// replaced in place, and we want that replacement to surface as an UPDATE to
// the same knowledge item rather than a delete-plus-insert pair.
//
// Emitted this cycle:
//
// - logical_key new since last sync → add (new content)
// - logical_key present, media_id changed → update (re-fetched content;
// ingest layer's existing
// external_id match triggers
// delete-and-recreate)
// - logical_key unchanged, media_id unchanged → skip
// - logical_key present last sync, absent this sync → IsDeleted tombstone
//
// LIMITATION: IMA still exposes no per-item updated_at, so an in-place edit
// that keeps the SAME media_id (same file bytes replaced under the hood by
// IMA, if that ever happens) is invisible to us. Users needing a full content
// refresh should periodically run a Full sync.
func (c *Connector) FetchIncremental(
ctx context.Context, config *types.DataSourceConfig, cursor *types.SyncCursor,
) ([]types.FetchedItem, *types.SyncCursor, error) {
resourceIDs := config.ResourceIDs
if len(resourceIDs) == 0 {
return nil, nil, fmt.Errorf("no resource IDs (knowledge base IDs) configured")
}
var prev *imaCursor
if cursor != nil && cursor.ConnectorCursor != nil {
var p imaCursor
b, _ := json.Marshal(cursor.ConnectorCursor)
_ = json.Unmarshal(b, &p)
prev = &p
}
items, newCursor, err := c.walk(ctx, config, resourceIDs, prev, true)
if err != nil {
return nil, nil, err
}
cursorMap := make(map[string]interface{})
b, _ := json.Marshal(newCursor)
_ = json.Unmarshal(b, &cursorMap)
return items, &types.SyncCursor{
LastSyncTime: newCursor.LastSyncTime,
ConnectorCursor: cursorMap,
}, nil
}
// walk is the shared implementation for FetchAll / FetchIncremental.
// When incremental is false, prev is ignored and returned cursor is nil.
func (c *Connector) walk(
ctx context.Context,
config *types.DataSourceConfig,
resourceIDs []string,
prev *imaCursor,
incremental bool,
) ([]types.FetchedItem, *imaCursor, error) {
cfg, err := parseIMAConfig(config)
if err != nil {
return nil, nil, err
}
cli := newClient(cfg)
newCursor := &imaCursor{
LastSyncTime: time.Now(),
KBLogical: make(map[string]map[string]string),
}
var out []types.FetchedItem
for _, kbID := range resourceIDs {
files, folderPath, err := listAllKBFiles(ctx, cli, kbID)
if err != nil {
return nil, nil, fmt.Errorf("list KB %s: %w", kbID, err)
}
// present holds every logical_key IMA listed this cycle; it drives
// delete detection. syncedLogical holds only the keys we actually
// resolved this cycle and is what gets persisted as the cursor, so a
// transient failure is retried next run rather than being remembered
// as "unchanged" forever.
present := make(map[string]struct{}, len(files))
syncedLogical := make(map[string]string, len(files))
keys := make([]string, len(files))
for i, f := range files {
key := logicalKey(kbID, f.ParentFolderID, f.Title)
keys[i] = key
present[key] = struct{}{}
}
newCursor.KBLogical[kbID] = syncedLogical
var replaced, failed int
for i, f := range files {
key := keys[i]
// Incremental skip / replacement detection.
if incremental && prev != nil && prev.KBLogical != nil {
if prevSet, ok := prev.KBLogical[kbID]; ok {
if prevMediaID, seen := prevSet[key]; seen {
if prevMediaID == f.MediaID {
// Unchanged: carry the entry forward so it is not
// re-fetched next cycle either.
syncedLogical[key] = f.MediaID
continue
}
replaced++ // media_id changed → replacement, fall through to re-fetch
}
}
}
item, outcome := fetchOneMedia(ctx, cli, kbID, key, f, folderPath)
switch outcome {
case fetchOK:
syncedLogical[key] = f.MediaID
out = append(out, item)
case fetchSkipped:
// Deterministic skip (unsupported media type, no URL): record it
// so we do not pay for get_media_info on every future sync.
syncedLogical[key] = f.MediaID
case fetchFailed:
// Transient: leave the key out of the cursor so the next run
// retries. `present` still contains it, so it is not mistaken
// for a deletion.
failed++
}
}
// Deletion detection (incremental only) — based on logical_key vanishing
// from the listing, so neither a same-name replacement (media_id churns
// but logical_key stays) nor a failed download fires a spurious delete.
deleted := 0
if incremental && prev != nil && prev.KBLogical != nil {
if prevSet, ok := prev.KBLogical[kbID]; ok {
for prevKey := range prevSet {
if _, still := present[prevKey]; !still {
out = append(out, types.FetchedItem{
ExternalID: prevKey,
IsDeleted: true,
SourceResourceID: kbID,
})
deleted++
}
}
}
}
logger.Infof(ctx, "[IMA] KB %s: total=%d replaced=%d failed=%d deleted=%d",
kbID, len(files), replaced, failed, deleted)
}
if !incremental {
return out, nil, nil
}
return out, newCursor, nil
}
// walkedFile is an internal representation of a knowledge item after we've
// resolved its folder location within the KB tree.
type walkedFile struct {
knowledgeInfo
// FolderPath is the "/A/B" style path (folder names, no leading slash)
// of the containing folder — empty for root-level items.
FolderPath string
}
// listAllKBFiles recursively enumerates a knowledge base into a flat list of
// files, resolving each item's folder path from the on-the-fly folder tree.
// folderPath maps folder_id → readable path (for metadata).
func listAllKBFiles(
ctx context.Context, cli *client, kbID string,
) ([]walkedFile, map[string]string, error) {
folderPath := map[string]string{}
var out []walkedFile
// BFS/DFS over folders. Start from root (empty folder_id).
type todo struct {
folderID string
path string
}
stack := []todo{{folderID: "", path: ""}}
for len(stack) > 0 {
cur := stack[len(stack)-1]
stack = stack[:len(stack)-1]
cursor := ""
for {
resp, err := cli.GetKnowledgeList(ctx, kbID, cur.folderID, cursor, defaultPageSize)
if err != nil {
return nil, nil, err
}
for _, raw := range resp.KnowledgeList {
// Probe each entry: an entry with a non-empty folder_id is a
// folder; otherwise it's a knowledge item (file / note / etc.).
var probe struct {
FolderID string `json:"folder_id"`
MediaID string `json:"media_id"`
}
_ = json.Unmarshal(raw, &probe)
if probe.FolderID != "" && probe.MediaID == "" {
var fi folderInfo
if err := json.Unmarshal(raw, &fi); err != nil {
continue
}
child := cur.path
if child == "" {
child = fi.Name
} else {
child = cur.path + "/" + fi.Name
}
folderPath[fi.FolderID] = child
stack = append(stack, todo{folderID: fi.FolderID, path: child})
continue
}
if probe.MediaID == "" {
continue // unrecognized shape, skip defensively
}
var ki knowledgeInfo
if err := json.Unmarshal(raw, &ki); err != nil {
continue
}
out = append(out, walkedFile{
knowledgeInfo: ki,
FolderPath: cur.path,
})
}
if resp.IsEnd || resp.NextCursor == "" {
break
}
cursor = resp.NextCursor
}
}
return out, folderPath, nil
}
// fetchNote resolves an IMA note (media_type=11). Notes carry no downloadable
// body: get_media_info only reports the notebook_id, and the text has to be
// read from the separate note namespace.
//
// The body arrives as plain text but is ingested as Markdown, because IMA notes
// are authored as rich text and the export keeps its heading and list markers —
// parsing it as Markdown preserves that structure for chunking, and text with
// no markers is unaffected.
func fetchNote(
ctx context.Context, cli *client,
kbID string, externalID string, f walkedFile, folderPath map[string]string,
info *getMediaInfoResp,
) (types.FetchedItem, fetchOutcome) {
noteID := info.NotebookExtInfo.NotebookID
if noteID == "" {
logger.Warnf(ctx, "[IMA] note %s (title=%q) has no notebook_id, skipping", f.MediaID, f.Title)
return types.FetchedItem{}, fetchSkipped
}
content, err := cli.GetNoteContent(ctx, noteID)
if err != nil {
logger.Warnf(ctx, "[IMA] get_doc_content(note %s, title=%q) failed, will retry next sync: %v",
noteID, f.Title, err)
return types.FetchedItem{}, fetchFailed
}
if strings.TrimSpace(content) == "" {
logger.Infof(ctx, "[IMA] note %s (title=%q) is empty, skipping", noteID, f.Title)
return types.FetchedItem{}, fetchSkipped
}
fileName := sanitizeFileName(f.Title)
if !strings.HasSuffix(strings.ToLower(fileName), ".md") {
fileName += ".md"
}
return types.FetchedItem{
ExternalID: externalID,
Title: f.Title,
Content: []byte(content),
ContentType: "text/markdown",
FileName: fileName,
SourceResourceID: kbID,
UpdatedAt: time.Now(),
Metadata: baseMetadata(externalID, f, folderPath, info, kbID),
}, fetchOK
}
// fetchOutcome distinguishes the three ways fetchOneMedia can end, because the
// caller has to treat them differently when building the sync cursor: only a
// transient failure must stay out of the cursor so it is retried next run.
type fetchOutcome int
const (
// fetchOK: the item was resolved and should be ingested.
fetchOK fetchOutcome = iota
// fetchSkipped: IMA cannot give us this item's content and never will
// (unsupported media type, no URL). Deterministic, so it is safe to
// remember in the cursor and stop re-probing it.
fetchSkipped
// fetchFailed: a transient error (API call or download). Must be retried.
fetchFailed
)
// fetchOneMedia calls get_media_info for a single item and, when possible,
// downloads its body.
//
// externalID is the stable logical_key computed by the caller — same value
// across syncs even after an IMA same-name replacement mints a new media_id.
// The raw media_id is still preserved in metadata for observability.
func fetchOneMedia(
ctx context.Context, cli *client,
kbID string, externalID string, f walkedFile, folderPath map[string]string,
) (types.FetchedItem, fetchOutcome) {
info, err := cli.GetMediaInfo(ctx, f.MediaID)
if err != nil {
logger.Warnf(ctx, "[IMA] get_media_info(%s) failed, will retry next sync: %v", f.MediaID, err)
return types.FetchedItem{}, fetchFailed
}
if info.MediaType == mediaTypeNote {
return fetchNote(ctx, cli, kbID, externalID, f, folderPath, info)
}
if isSkippableMediaType(info.MediaType) {
logger.Infof(ctx, "[IMA] skip media %s (title=%q media_type=%d): unsupported by the IMA OpenAPI",
f.MediaID, f.Title, info.MediaType)
return types.FetchedItem{}, fetchSkipped
}
if info.URLInfo.URL == "" {
logger.Warnf(ctx, "[IMA] media %s (title=%q) has no url_info.url, skipping", f.MediaID, f.Title)
return types.FetchedItem{}, fetchSkipped
}
ext := extensionForMediaType(info.MediaType)
// A media type with no fixed extension (web page, WeChat article, ...) is
// normally handed to the ingest layer as a bare URL so WeKnora fetches it
// itself. That only works for publicly reachable links: when IMA attaches
// auth headers the URL points at IMA-hosted storage and WeKnora's own
// fetch — which cannot carry those headers — would be rejected. In that
// case download here, where the headers are available, and infer the
// extension from the response instead.
if ext == "" {
if len(info.URLInfo.Headers) == 0 {
return types.FetchedItem{
ExternalID: externalID,
Title: f.Title,
URL: info.URLInfo.URL,
SourceResourceID: kbID,
UpdatedAt: time.Now(),
Metadata: baseMetadata(externalID, f, folderPath, info, kbID),
}, fetchOK
}
ext = "html"
}
body, ct, err := cli.DownloadURL(ctx, info.URLInfo)
if err != nil {
logger.Warnf(ctx, "[IMA] download media %s (title=%q) failed, will retry next sync: %v",
f.MediaID, f.Title, err)
return types.FetchedItem{}, fetchFailed
}
// IMA labels every image media_type=9, so trust the response's own
// Content-Type over the media_type default when it names a real format.
if detected := extensionForContentType(ct); detected != "" {
ext = detected
}
if ct == "" || strings.HasPrefix(ct, "application/octet-stream") {
ct = mimeForExtension(ext)
}
fileName := sanitizeFileName(f.Title)
if !strings.HasSuffix(strings.ToLower(fileName), "."+ext) {
fileName = fileName + "." + ext
}
return types.FetchedItem{
ExternalID: externalID,
Title: f.Title,
Content: body,
ContentType: ct,
FileName: fileName,
URL: info.URLInfo.URL,
UpdatedAt: time.Now(),
SourceResourceID: kbID,
Metadata: baseMetadata(externalID, f, folderPath, info, kbID),
}, fetchOK
}
// baseMetadata builds the metadata map preserved on every ingested item.
// externalID (the caller's logical_key) is stored so operators can join the
// stable identity back to the raw media_id — useful when debugging why a
// same-name replacement showed up as an update instead of an add.
func baseMetadata(
externalID string, f walkedFile, folderPath map[string]string,
info *getMediaInfoResp, kbID string,
) map[string]string {
m := map[string]string{
"channel": types.ChannelIMA,
"media_id": f.MediaID,
"ima_logical_key": externalID,
"knowledge_base_id": kbID,
"folder_path": f.FolderPath,
"media_type": fmt.Sprintf("%d", info.MediaType),
}
if f.ParentFolderID != "" {
m["parent_folder_id"] = f.ParentFolderID
}
if fp, ok := folderPath[f.ParentFolderID]; ok && fp != "" {
m["folder_path"] = fp
}
if info.NotebookExtInfo.NotebookID != "" {
m["notebook_id"] = info.NotebookExtInfo.NotebookID
}
return m
}
// sanitizeFileName removes filesystem-hostile characters and truncates to a
// safe UTF-8 boundary.
func sanitizeFileName(name string) string {
if name == "" {
return "untitled"
}
replacer := strings.NewReplacer(
"/", "_", "\\", "_", ":", "_", "*", "_",
"?", "_", "\"", "_", "<", "_", ">", "_", "|", "_",
)
result := replacer.Replace(name)
const maxBytes = 200
if len(result) > maxBytes {
result = result[:maxBytes]
for len(result) > 0 {
r, size := utf8.DecodeLastRuneInString(result)
if r != utf8.RuneError || size != 1 {
break
}
result = result[:len(result)-1]
}
}
return result
}
@@ -0,0 +1,511 @@
package ima
import (
"context"
"strings"
"testing"
"github.com/Tencent/WeKnora/internal/datasource"
"github.com/Tencent/WeKnora/internal/types"
)
func TestConnector_Type(t *testing.T) {
if got := NewConnector().Type(); got != types.ConnectorTypeIMA {
t.Errorf("Type() = %q, want %q", got, types.ConnectorTypeIMA)
}
}
// TestFetchIncremental_TransientFailureIsRetried is the regression test for the
// cursor bug: the cursor used to be built from the raw listing before any
// content was fetched, so an item whose download failed was recorded as
// "seen with this media_id" and every later incremental sync skipped it as
// unchanged. The document was then silently missing from the KB forever.
func TestFetchIncremental_TransientFailureIsRetried(t *testing.T) {
f := newFakeIMA(t)
good := fakeFile{MediaID: "m-good", Title: "Good", MediaType: mediaTypeMarkdown, Body: "# ok"}
flaky := fakeFile{
MediaID: "m-flaky", Title: "Flaky", MediaType: mediaTypeMarkdown,
Body: "# later", DownloadFails: true,
}
f.setKB("kb1", []fakeFile{good, flaky})
c := NewConnector()
cfg := f.config("kb1")
goodKey := logicalKey("kb1", "", "Good")
flakyKey := logicalKey("kb1", "", "Flaky")
items, cursor, err := c.FetchIncremental(context.Background(), cfg, nil)
if err != nil {
t.Fatalf("first FetchIncremental: %v", err)
}
if _, ok := findItem(items, flakyKey); ok {
t.Fatalf("failed download must not be emitted, got %s", describeItems(items))
}
mustFindItem(t, items, goodKey)
state := decodeCursor(t, cursor)
if _, recorded := state.KBLogical["kb1"][flakyKey]; recorded {
t.Fatal("a failed item must not be recorded in the cursor, otherwise it is never retried")
}
if state.KBLogical["kb1"][goodKey] != "m-good" {
t.Fatalf("successful item missing from cursor: %#v", state.KBLogical["kb1"])
}
// Second pass: the download now succeeds and the item must be re-attempted.
flaky.DownloadFails = false
f.setKB("kb1", []fakeFile{good, flaky})
items2, cursor2, err := c.FetchIncremental(context.Background(), cfg, cursor)
if err != nil {
t.Fatalf("second FetchIncremental: %v", err)
}
retried := mustFindItem(t, items2, flakyKey)
if string(retried.Content) != "# later" {
t.Errorf("retried content = %q, want %q", retried.Content, "# later")
}
if _, ok := findItem(items2, goodKey); ok {
t.Errorf("unchanged item must be skipped on the second pass, got %s", describeItems(items2))
}
if decodeCursor(t, cursor2).KBLogical["kb1"][flakyKey] != "m-flaky" {
t.Error("the retried item should now be recorded in the cursor")
}
}
// TestFetchIncremental_FailureIsNotADeletion guards the other half of the fix:
// keeping a failed item out of the cursor must not make the next sync mistake
// it for a document that vanished from IMA.
func TestFetchIncremental_FailureIsNotADeletion(t *testing.T) {
f := newFakeIMA(t)
good := fakeFile{MediaID: "m-good", Title: "Good", MediaType: mediaTypeMarkdown, Body: "# ok"}
flaky := fakeFile{MediaID: "m-flaky", Title: "Flaky", MediaType: mediaTypeMarkdown, DownloadFails: true}
f.setKB("kb1", []fakeFile{good, flaky})
c := NewConnector()
cfg := f.config("kb1")
flakyKey := logicalKey("kb1", "", "Flaky")
_, cursor, err := c.FetchIncremental(context.Background(), cfg, nil)
if err != nil {
t.Fatalf("first FetchIncremental: %v", err)
}
// Still failing, still listed by IMA: no tombstone may be emitted.
items, _, err := c.FetchIncremental(context.Background(), cfg, cursor)
if err != nil {
t.Fatalf("second FetchIncremental: %v", err)
}
for _, it := range items {
if it.IsDeleted && it.ExternalID == flakyKey {
t.Fatalf("a still-listed item that merely failed to download must not be tombstoned: %s",
describeItems(items))
}
}
}
func TestFetchIncremental_RemovedItemIsTombstoned(t *testing.T) {
f := newFakeIMA(t)
keep := fakeFile{MediaID: "m-keep", Title: "Keep", MediaType: mediaTypeMarkdown, Body: "# keep"}
drop := fakeFile{MediaID: "m-drop", Title: "Drop", MediaType: mediaTypeMarkdown, Body: "# drop"}
f.setKB("kb1", []fakeFile{keep, drop})
c := NewConnector()
cfg := f.config("kb1")
dropKey := logicalKey("kb1", "", "Drop")
_, cursor, err := c.FetchIncremental(context.Background(), cfg, nil)
if err != nil {
t.Fatalf("first FetchIncremental: %v", err)
}
f.setKB("kb1", []fakeFile{keep})
items, _, err := c.FetchIncremental(context.Background(), cfg, cursor)
if err != nil {
t.Fatalf("second FetchIncremental: %v", err)
}
tombstone, ok := findItem(items, dropKey)
if !ok || !tombstone.IsDeleted {
t.Fatalf("removed item must be tombstoned, got %s", describeItems(items))
}
if tombstone.SourceResourceID != "kb1" {
t.Errorf("tombstone SourceResourceID = %q, want kb1", tombstone.SourceResourceID)
}
}
// TestFetchIncremental_SameNameReplacementKeepsExternalID covers IMA minting a
// new media_id when a file is replaced in place: the item must come back as an
// update under the same external id rather than a delete plus an insert.
func TestFetchIncremental_SameNameReplacementKeepsExternalID(t *testing.T) {
f := newFakeIMA(t)
f.setKB("kb1", []fakeFile{{MediaID: "m-v1", Title: "Doc", MediaType: mediaTypeMarkdown, Body: "v1"}})
c := NewConnector()
cfg := f.config("kb1")
key := logicalKey("kb1", "", "Doc")
_, cursor, err := c.FetchIncremental(context.Background(), cfg, nil)
if err != nil {
t.Fatalf("first FetchIncremental: %v", err)
}
f.setKB("kb1", []fakeFile{{MediaID: "m-v2", Title: "Doc", MediaType: mediaTypeMarkdown, Body: "v2"}})
items, cursor2, err := c.FetchIncremental(context.Background(), cfg, cursor)
if err != nil {
t.Fatalf("second FetchIncremental: %v", err)
}
updated := mustFindItem(t, items, key)
if updated.IsDeleted {
t.Fatal("a same-name replacement must be an update, not a delete")
}
if string(updated.Content) != "v2" {
t.Errorf("content = %q, want v2", updated.Content)
}
if updated.Metadata["media_id"] != "m-v2" {
t.Errorf("metadata media_id = %q, want m-v2", updated.Metadata["media_id"])
}
for _, it := range items {
if it.IsDeleted {
t.Errorf("no tombstone expected, got %s", describeItems(items))
}
}
if decodeCursor(t, cursor2).KBLogical["kb1"][key] != "m-v2" {
t.Error("cursor should track the new media_id")
}
}
// TestFetchIncremental_UnsupportedTypeIsProbedOnce checks that a deterministic
// skip is remembered, so an IMA knowledge base full of notes does not re-pay
// for get_media_info on every scheduled sync.
func TestFetchIncremental_UnsupportedTypeIsProbedOnce(t *testing.T) {
f := newFakeIMA(t)
f.setKB("kb1", []fakeFile{{MediaID: "m-note", Title: "Note", MediaType: mediaTypeNote}})
c := NewConnector()
cfg := f.config("kb1")
items, cursor, err := c.FetchIncremental(context.Background(), cfg, nil)
if err != nil {
t.Fatalf("first FetchIncremental: %v", err)
}
if len(items) != 0 {
t.Fatalf("note must be skipped, got %s", describeItems(items))
}
if f.callCount("get_media_info:m-note") != 1 {
t.Fatalf("get_media_info calls = %d, want 1", f.callCount("get_media_info:m-note"))
}
if _, _, err := c.FetchIncremental(context.Background(), cfg, cursor); err != nil {
t.Fatalf("second FetchIncremental: %v", err)
}
if got := f.callCount("get_media_info:m-note"); got != 1 {
t.Errorf("get_media_info calls after second sync = %d, want 1 (skip should be cached)", got)
}
}
// TestFetchAll_NoteBodyIsReadFromNoteNamespace covers IMA notes. They expose no
// url_info at all: get_media_info only reports a notebook_id, and the body has
// to be read from /openapi/note/v1/get_doc_content.
func TestFetchAll_NoteBodyIsReadFromNoteNamespace(t *testing.T) {
f := newFakeIMA(t)
f.setKB("kb1", []fakeFile{{
MediaID: "note_1",
Title: "Meeting notes",
MediaType: mediaTypeNote,
NotebookID: "987654321",
NoteBody: "# Standup\n- shipped the connector",
}})
items, err := NewConnector().FetchAll(context.Background(), f.config("kb1"), []string{"kb1"})
if err != nil {
t.Fatalf("FetchAll: %v", err)
}
note := mustFindItem(t, items, logicalKey("kb1", "", "Meeting notes"))
if string(note.Content) != "# Standup\n- shipped the connector" {
t.Errorf("content = %q, want the note body", note.Content)
}
if !strings.HasSuffix(note.FileName, ".md") {
t.Errorf("FileName = %q, want an .md suffix", note.FileName)
}
if note.ContentType != "text/markdown" {
t.Errorf("ContentType = %q, want text/markdown", note.ContentType)
}
if note.Metadata["notebook_id"] != "987654321" {
t.Errorf("metadata notebook_id = %q, want 987654321", note.Metadata["notebook_id"])
}
if f.callCount("get_doc_content:987654321") != 1 {
t.Errorf("get_doc_content calls = %d, want 1", f.callCount("get_doc_content:987654321"))
}
}
// TestFetchIncremental_NoteReadFailureIsRetried checks that a note whose body
// could not be read follows the same retry path as a failed download rather
// than being remembered as synced.
func TestFetchIncremental_NoteReadFailureIsRetried(t *testing.T) {
f := newFakeIMA(t)
note := fakeFile{
MediaID: "note_1", Title: "Broken", MediaType: mediaTypeNote,
NotebookID: "111", NoteBody: "recovered", NoteFails: true,
}
f.setKB("kb1", []fakeFile{note})
c := NewConnector()
cfg := f.config("kb1")
key := logicalKey("kb1", "", "Broken")
items, cursor, err := c.FetchIncremental(context.Background(), cfg, nil)
if err != nil {
t.Fatalf("first FetchIncremental: %v", err)
}
if len(items) != 0 {
t.Fatalf("a note that could not be read must not be emitted, got %s", describeItems(items))
}
if _, recorded := decodeCursor(t, cursor).KBLogical["kb1"][key]; recorded {
t.Fatal("a failed note read must not be recorded in the cursor")
}
note.NoteFails = false
f.setKB("kb1", []fakeFile{note})
items2, _, err := c.FetchIncremental(context.Background(), cfg, cursor)
if err != nil {
t.Fatalf("second FetchIncremental: %v", err)
}
if got := mustFindItem(t, items2, key); string(got.Content) != "recovered" {
t.Errorf("content = %q, want recovered", got.Content)
}
}
func TestFetchAll_EmptyNoteIsSkipped(t *testing.T) {
f := newFakeIMA(t)
f.setKB("kb1", []fakeFile{{
MediaID: "note_1", Title: "Blank", MediaType: mediaTypeNote,
NotebookID: "222", NoteBody: " \n ",
}})
items, err := NewConnector().FetchAll(context.Background(), f.config("kb1"), []string{"kb1"})
if err != nil {
t.Fatalf("FetchAll: %v", err)
}
if len(items) != 0 {
t.Fatalf("an empty note must be skipped, got %s", describeItems(items))
}
}
// TestFetchAll_AISessionIsStillSkipped pins the types IMA genuinely cannot
// export, so widening note support does not quietly widen these too.
func TestFetchAll_AISessionIsStillSkipped(t *testing.T) {
f := newFakeIMA(t)
f.setKB("kb1", []fakeFile{
{MediaID: "m-ai", Title: "Chat", MediaType: mediaTypeAISession},
{MediaID: "m-video", Title: "Clip", MediaType: mediaTypeVideo},
})
items, err := NewConnector().FetchAll(context.Background(), f.config("kb1"), []string{"kb1"})
if err != nil {
t.Fatalf("FetchAll: %v", err)
}
if len(items) != 0 {
t.Fatalf("AI sessions and video parses must be skipped, got %s", describeItems(items))
}
}
// TestFetchAll_AuthenticatedURLIsDownloaded covers media types with no fixed
// extension. When IMA attaches auth headers the URL is IMA-hosted and WeKnora's
// own fetch could not authenticate, so the connector must download it here and
// forward the headers.
func TestFetchAll_AuthenticatedURLIsDownloaded(t *testing.T) {
f := newFakeIMA(t)
f.setKB("kb1", []fakeFile{{
MediaID: "m-web",
Title: "Article",
MediaType: mediaTypeWeb,
Body: "<html>body</html>",
ContentType: "text/html",
URLHeaders: map[string]string{"X-Ima-Token": "secret"},
}})
items, err := NewConnector().FetchAll(context.Background(), f.config("kb1"), []string{"kb1"})
if err != nil {
t.Fatalf("FetchAll: %v", err)
}
item := mustFindItem(t, items, logicalKey("kb1", "", "Article"))
if string(item.Content) != "<html>body</html>" {
t.Errorf("content = %q, want the downloaded body", item.Content)
}
if !strings.HasSuffix(item.FileName, ".html") {
t.Errorf("FileName = %q, want an .html suffix", item.FileName)
}
if got := f.downloadHeaders["m-web"].Get("X-Ima-Token"); got != "secret" {
t.Errorf("auth header forwarded = %q, want %q", got, "secret")
}
}
// TestFetchAll_PublicURLStaysURLOnly is the counterpart: without auth headers
// the link is publicly reachable, so it is handed to the ingest layer as a URL
// and WeKnora renders the live page instead of a snapshot.
func TestFetchAll_PublicURLStaysURLOnly(t *testing.T) {
f := newFakeIMA(t)
f.setKB("kb1", []fakeFile{{MediaID: "m-web", Title: "Article", MediaType: mediaTypeWeb, Body: "ignored"}})
items, err := NewConnector().FetchAll(context.Background(), f.config("kb1"), []string{"kb1"})
if err != nil {
t.Fatalf("FetchAll: %v", err)
}
item := mustFindItem(t, items, logicalKey("kb1", "", "Article"))
if len(item.Content) != 0 {
t.Errorf("public URL must not be downloaded, got %d bytes", len(item.Content))
}
if item.URL == "" {
t.Error("URL-only item must carry a URL")
}
if f.callCount("download:m-web") != 0 {
t.Error("no download should have been attempted")
}
}
// TestFetchAll_ImageExtensionFollowsContentType guards against naming a JPEG
// ".png": IMA reports every image as media_type=9, so only the download's
// Content-Type reveals the real format, and the extension drives the ingest
// pipeline's file-type checks.
func TestFetchAll_ImageExtensionFollowsContentType(t *testing.T) {
f := newFakeIMA(t)
f.setKB("kb1", []fakeFile{{
MediaID: "m-img",
Title: "Photo",
MediaType: mediaTypeImage,
Body: "\xff\xd8\xff\xe0jpeg",
ContentType: "image/jpeg",
}})
items, err := NewConnector().FetchAll(context.Background(), f.config("kb1"), []string{"kb1"})
if err != nil {
t.Fatalf("FetchAll: %v", err)
}
item := mustFindItem(t, items, logicalKey("kb1", "", "Photo"))
if !strings.HasSuffix(item.FileName, ".jpg") {
t.Errorf("FileName = %q, want a .jpg suffix for an image/jpeg body", item.FileName)
}
}
func TestFetchAll_SkipsItemWithoutURL(t *testing.T) {
f := newFakeIMA(t)
f.setKB("kb1", []fakeFile{{MediaID: "m-empty", Title: "Empty", MediaType: mediaTypePDF, NoURL: true}})
items, err := NewConnector().FetchAll(context.Background(), f.config("kb1"), []string{"kb1"})
if err != nil {
t.Fatalf("FetchAll: %v", err)
}
if len(items) != 0 {
t.Fatalf("item without a URL must be skipped, got %s", describeItems(items))
}
}
// TestFetchAll_ResolvesFolderPath checks the mixed file/folder listing walk.
func TestFetchAll_ResolvesFolderPath(t *testing.T) {
f := newFakeIMA(t)
f.setKB("kb1",
[]fakeFile{
{MediaID: "m-root", Title: "Root", MediaType: mediaTypeMarkdown, Body: "root"},
{MediaID: "m-deep", Title: "Deep", ParentFolderID: "f2", MediaType: mediaTypeMarkdown, Body: "deep"},
},
fakeFolder{FolderID: "f1", Name: "Outer"},
fakeFolder{FolderID: "f2", Name: "Inner", ParentFolderID: "f1"},
)
items, err := NewConnector().FetchAll(context.Background(), f.config("kb1"), []string{"kb1"})
if err != nil {
t.Fatalf("FetchAll: %v", err)
}
deep := mustFindItem(t, items, logicalKey("kb1", "f2", "Deep"))
if deep.Metadata["folder_path"] != "Outer/Inner" {
t.Errorf("folder_path = %q, want Outer/Inner", deep.Metadata["folder_path"])
}
root := mustFindItem(t, items, logicalKey("kb1", "", "Root"))
if root.Metadata["folder_path"] != "" {
t.Errorf("root folder_path = %q, want empty", root.Metadata["folder_path"])
}
}
func TestListResources_ReturnsAddableBases(t *testing.T) {
f := newFakeIMA(t)
f.setKB("kb-b", nil)
f.setKB("kb-a", nil)
resources, err := NewConnector().ListResources(context.Background(), f.config(), "")
if err != nil {
t.Fatalf("ListResources: %v", err)
}
if len(resources) != 2 {
t.Fatalf("got %d resources, want 2", len(resources))
}
if resources[0].ExternalID != "kb-a" || resources[1].ExternalID != "kb-b" {
t.Errorf("resources are not sorted by id: %+v", resources)
}
if resources[0].Description != "desc kb-a" {
t.Errorf("description not enriched from get_knowledge_base: %q", resources[0].Description)
}
}
// TestListResources_FallsBackToSearch covers tenants whose credential exposes
// read-only scopes, where get_addable_knowledge_base_list comes back empty.
func TestListResources_FallsBackToSearch(t *testing.T) {
f := newFakeIMA(t)
f.searchBases = []searchedKnowledgeBaseInfo{{ID: "kb-x", Name: "X", CoverURL: "https://cover"}}
resources, err := NewConnector().ListResources(context.Background(), f.config(), "")
if err != nil {
t.Fatalf("ListResources: %v", err)
}
if len(resources) != 1 || resources[0].ExternalID != "kb-x" {
t.Fatalf("expected the search fallback to surface kb-x, got %+v", resources)
}
if resources[0].Metadata["cover_url"] != "https://cover" {
t.Errorf("cover_url = %v, want https://cover", resources[0].Metadata["cover_url"])
}
}
func TestListResources_ChildrenAreEmpty(t *testing.T) {
f := newFakeIMA(t)
f.setKB("kb1", nil)
resources, err := NewConnector().ListResources(context.Background(), f.config(), "kb1")
if err != nil {
t.Fatalf("ListResources: %v", err)
}
if len(resources) != 0 {
t.Errorf("knowledge bases have no children, got %+v", resources)
}
}
func TestValidate_InvalidCredentials(t *testing.T) {
f := newFakeIMA(t)
cfg := f.config()
cfg.Credentials["client_id"] = ""
err := NewConnector().Validate(context.Background(), cfg)
if err == nil {
t.Fatal("expected an error for a missing client_id")
}
if !strings.Contains(err.Error(), "client_id") {
t.Errorf("error = %v, want it to name the missing field", err)
}
}
func TestFetchIncremental_RequiresResourceIDs(t *testing.T) {
f := newFakeIMA(t)
if _, _, err := NewConnector().FetchIncremental(context.Background(), f.config(), nil); err == nil {
t.Fatal("expected an error when no knowledge base is selected")
}
}
func TestConnectorIsRegisteredInMetadata(t *testing.T) {
meta, ok := datasource.ConnectorMetadataRegistry[types.ConnectorTypeIMA]
if !ok {
t.Fatal("IMA connector missing from the metadata registry")
}
if meta.Type != types.ConnectorTypeIMA {
t.Errorf("metadata type = %q, want %q", meta.Type, types.ConnectorTypeIMA)
}
}
@@ -0,0 +1,359 @@
package ima
import (
"encoding/json"
"fmt"
"net/http"
"net/http/httptest"
"os"
"strings"
"sync"
"testing"
"github.com/Tencent/WeKnora/internal/types"
secutils "github.com/Tencent/WeKnora/internal/utils"
)
// TestMain whitelists loopback for SSRF so the httptest servers (127.0.0.1)
// are reachable. Production keeps the default strict SSRF policy.
func TestMain(m *testing.M) {
_ = os.Setenv("SSRF_WHITELIST", "127.0.0.1,localhost")
secutils.ResetSSRFWhitelistForTest()
os.Exit(m.Run())
}
// fakeFile is one entry in a fake knowledge base listing.
type fakeFile struct {
MediaID string
Title string
ParentFolderID string
// MediaType is what get_media_info reports for this file.
MediaType int32
// Body is what the download URL serves.
Body string
// ContentType is the download response's Content-Type ("" omits the header).
ContentType string
// URLHeaders are the auth headers IMA attaches to url_info.
URLHeaders map[string]string
// NoURL makes get_media_info return an empty url_info.url.
NoURL bool
// InfoFails makes get_media_info return a 500 for this media_id.
InfoFails bool
// DownloadFails makes the download URL return a 500.
DownloadFails bool
// NotebookID is reported under notebook_ext_info for notes.
NotebookID string
// NoteBody is what get_doc_content serves for NotebookID.
NoteBody string
// NoteFails makes get_doc_content return a business error.
NoteFails bool
}
// fakeFolder is a folder entry in a fake knowledge base listing.
type fakeFolder struct {
FolderID string
Name string
ParentFolderID string
}
// fakeIMA is an in-process stand-in for the IMA OpenAPI. Only the endpoints
// the connector calls are implemented.
type fakeIMA struct {
server *httptest.Server
mu sync.Mutex
// files and folders are keyed by knowledge base id, then by parent folder
// id ("" for the KB root).
files map[string]map[string][]fakeFile
folders map[string]map[string][]fakeFolder
// bases is what get_addable_knowledge_base_list returns.
bases []addableKnowledgeBaseInfo
// searchBases is what search_knowledge_base returns.
searchBases []searchedKnowledgeBaseInfo
// calls counts requests per action, plus "download:<media_id>".
calls map[string]int
// downloadHeaders records the headers seen by the last download of a media.
downloadHeaders map[string]http.Header
}
func newFakeIMA(t *testing.T) *fakeIMA {
t.Helper()
f := &fakeIMA{
files: map[string]map[string][]fakeFile{},
folders: map[string]map[string][]fakeFolder{},
calls: map[string]int{},
downloadHeaders: map[string]http.Header{},
}
mux := http.NewServeMux()
mux.HandleFunc(apiBasePath+"/", f.handleAPI)
mux.HandleFunc(noteBasePath+"/", f.handleNoteAPI)
mux.HandleFunc("/dl/", f.handleDownload)
f.server = httptest.NewServer(mux)
t.Cleanup(f.server.Close)
return f
}
// setKB replaces the listing of a knowledge base and registers it as addable.
func (f *fakeIMA) setKB(kbID string, files []fakeFile, folders ...fakeFolder) {
f.mu.Lock()
defer f.mu.Unlock()
byParent := map[string][]fakeFile{}
for _, file := range files {
byParent[file.ParentFolderID] = append(byParent[file.ParentFolderID], file)
}
f.files[kbID] = byParent
foldersByParent := map[string][]fakeFolder{}
for _, folder := range folders {
foldersByParent[folder.ParentFolderID] = append(foldersByParent[folder.ParentFolderID], folder)
}
f.folders[kbID] = foldersByParent
for _, b := range f.bases {
if b.ID == kbID {
return
}
}
f.bases = append(f.bases, addableKnowledgeBaseInfo{ID: kbID, Name: "KB " + kbID})
}
func (f *fakeIMA) callCount(key string) int {
f.mu.Lock()
defer f.mu.Unlock()
return f.calls[key]
}
func (f *fakeIMA) record(key string) {
f.mu.Lock()
defer f.mu.Unlock()
f.calls[key]++
}
// findFile locates a file by media_id across every knowledge base.
func (f *fakeIMA) findFile(mediaID string) (fakeFile, bool) {
f.mu.Lock()
defer f.mu.Unlock()
for _, byParent := range f.files {
for _, files := range byParent {
for _, file := range files {
if file.MediaID == mediaID {
return file, true
}
}
}
}
return fakeFile{}, false
}
func (f *fakeIMA) handleAPI(w http.ResponseWriter, r *http.Request) {
action := strings.TrimPrefix(r.URL.Path, apiBasePath+"/")
f.record(action)
var req map[string]interface{}
_ = json.NewDecoder(r.Body).Decode(&req)
str := func(k string) string {
if v, ok := req[k].(string); ok {
return v
}
return ""
}
switch action {
case "get_addable_knowledge_base_list":
f.mu.Lock()
resp := getAddableKnowledgeBaseListResp{AddableKnowledgeBaseList: f.bases, IsEnd: true}
f.mu.Unlock()
writeEnvelope(w, 0, "", resp)
case "search_knowledge_base":
f.mu.Lock()
resp := searchKnowledgeBaseResp{InfoList: f.searchBases, IsEnd: true}
f.mu.Unlock()
writeEnvelope(w, 0, "", resp)
case "get_knowledge_base":
infos := map[string]knowledgeBaseInfo{}
if ids, ok := req["ids"].([]interface{}); ok {
for _, raw := range ids {
id, _ := raw.(string)
infos[id] = knowledgeBaseInfo{ID: id, Name: "KB " + id, Description: "desc " + id}
}
}
writeEnvelope(w, 0, "", getKnowledgeBaseResp{Infos: infos})
case "get_knowledge_list":
kbID, folderID := str("knowledge_base_id"), str("folder_id")
f.mu.Lock()
files := f.files[kbID][folderID]
folders := f.folders[kbID][folderID]
f.mu.Unlock()
var list []json.RawMessage
for _, folder := range folders {
b, _ := json.Marshal(folderInfo{
FolderID: folder.FolderID,
Name: folder.Name,
ParentFolderID: folder.ParentFolderID,
})
list = append(list, b)
}
for _, file := range files {
// IMA omits media_type in list responses.
b, _ := json.Marshal(map[string]interface{}{
"media_id": file.MediaID,
"title": file.Title,
"parent_folder_id": file.ParentFolderID,
})
list = append(list, b)
}
writeEnvelope(w, 0, "", getKnowledgeListResp{KnowledgeList: list, IsEnd: true})
case "get_media_info":
mediaID := str("media_id")
f.record("get_media_info:" + mediaID)
file, ok := f.findFile(mediaID)
if !ok {
writeEnvelope(w, 110001, "unknown media", nil)
return
}
if file.InfoFails {
w.WriteHeader(http.StatusInternalServerError)
return
}
resp := getMediaInfoResp{MediaType: file.MediaType}
resp.NotebookExtInfo = notebookExtInfo{NotebookID: file.NotebookID}
// Notes carry no url_info; the body lives in the note namespace.
if !file.NoURL && file.MediaType != mediaTypeNote {
resp.URLInfo = urlInfo{
URL: f.server.URL + "/dl/" + mediaID,
Headers: file.URLHeaders,
}
}
writeEnvelope(w, 0, "", resp)
default:
writeEnvelope(w, 110001, "unsupported action "+action, nil)
}
}
// handleNoteAPI serves the /openapi/note/v1 namespace.
func (f *fakeIMA) handleNoteAPI(w http.ResponseWriter, r *http.Request) {
action := strings.TrimPrefix(r.URL.Path, noteBasePath+"/")
f.record("note/" + action)
var req map[string]interface{}
_ = json.NewDecoder(r.Body).Decode(&req)
if action != "get_doc_content" {
writeEnvelope(w, 110012, "unsupported note action "+action, nil)
return
}
noteID, _ := req["note_id"].(string)
f.record("get_doc_content:" + noteID)
f.mu.Lock()
defer f.mu.Unlock()
for _, byParent := range f.files {
for _, files := range byParent {
for _, file := range files {
if file.NotebookID != noteID || noteID == "" {
continue
}
if file.NoteFails {
writeEnvelope(w, 110011, "note read failed", nil)
return
}
writeEnvelope(w, 0, "", getDocContentResp{Content: file.NoteBody})
return
}
}
}
writeEnvelope(w, 110001, "unknown note", nil)
}
func (f *fakeIMA) handleDownload(w http.ResponseWriter, r *http.Request) {
mediaID := strings.TrimPrefix(r.URL.Path, "/dl/")
f.record("download:" + mediaID)
f.mu.Lock()
f.downloadHeaders[mediaID] = r.Header.Clone()
f.mu.Unlock()
file, ok := f.findFile(mediaID)
if !ok || file.DownloadFails {
w.WriteHeader(http.StatusInternalServerError)
return
}
if file.ContentType != "" {
w.Header().Set("Content-Type", file.ContentType)
}
_, _ = w.Write([]byte(file.Body))
}
func writeEnvelope(w http.ResponseWriter, code int, msg string, data interface{}) {
raw, _ := json.Marshal(data)
w.Header().Set("Content-Type", "application/json")
_ = json.NewEncoder(w).Encode(apiEnvelope{Code: code, Msg: msg, Data: raw})
}
func (f *fakeIMA) config(resourceIDs ...string) *types.DataSourceConfig {
return &types.DataSourceConfig{
Type: types.ConnectorTypeIMA,
Credentials: map[string]interface{}{
"client_id": "cid",
"api_key": "key",
"base_url": f.server.URL,
},
ResourceIDs: resourceIDs,
}
}
// decodeCursor converts the connector cursor back into its typed form.
func decodeCursor(t *testing.T, cur *types.SyncCursor) *imaCursor {
t.Helper()
if cur == nil {
t.Fatal("cursor is nil")
}
b, err := json.Marshal(cur.ConnectorCursor)
if err != nil {
t.Fatalf("marshal cursor: %v", err)
}
var out imaCursor
if err := json.Unmarshal(b, &out); err != nil {
t.Fatalf("unmarshal cursor: %v", err)
}
return &out
}
// findItem returns the fetched item with the given external id.
func findItem(items []types.FetchedItem, externalID string) (types.FetchedItem, bool) {
for _, it := range items {
if it.ExternalID == externalID {
return it, true
}
}
return types.FetchedItem{}, false
}
func mustFindItem(t *testing.T, items []types.FetchedItem, externalID string) types.FetchedItem {
t.Helper()
it, ok := findItem(items, externalID)
if !ok {
t.Fatalf("item %s not found in %s", externalID, describeItems(items))
}
return it
}
func describeItems(items []types.FetchedItem) string {
var parts []string
for _, it := range items {
parts = append(parts, fmt.Sprintf("{id=%s title=%q deleted=%t}", it.ExternalID, it.Title, it.IsDeleted))
}
return "[" + strings.Join(parts, " ") + "]"
}
+433
View File
@@ -0,0 +1,433 @@
package ima
import (
"crypto/sha256"
"encoding/hex"
"encoding/json"
"fmt"
"strings"
"time"
"github.com/Tencent/WeKnora/internal/datasource"
"github.com/Tencent/WeKnora/internal/types"
)
// DefaultBaseURL is the IMA OpenAPI base URL. The Path prefix is baked into
// the client (see client.go); users only configure the host.
const DefaultBaseURL = "https://ima.qq.com"
// Base path for knowledge-base OpenAPI requests.
const apiBasePath = "/openapi/wiki/v1"
// Base path for the note OpenAPI. Notes live in their own namespace: the wiki
// endpoints only ever hand back a notebook_id for them, and the body has to be
// read from here.
const noteBasePath = "/openapi/note/v1"
// Config holds IMA-specific configuration decoded from
// DataSourceConfig.Credentials. Both credentials are opaque strings and
// stored encrypted at rest by DataSourceConfig.ToJSON (see internal/types/datasource.go).
type Config struct {
// ClientID: value of the `ima-openapi-clientid` request header,
// obtained from https://ima.qq.com/agent-interface.
ClientID string `json:"client_id"`
// APIKey: value of the `ima-openapi-apikey` request header.
APIKey string `json:"api_key"`
// BaseURL lets users point at an on-prem / test IMA endpoint.
// Defaults to https://ima.qq.com when empty.
BaseURL string `json:"base_url,omitempty"`
}
// GetBaseURL returns the normalized base URL (empty → default, no trailing slash).
func (c *Config) GetBaseURL() string {
url := strings.TrimSpace(c.BaseURL)
if url == "" {
return DefaultBaseURL
}
if !strings.Contains(url, "://") {
url = "https://" + url
}
return strings.TrimRight(url, "/")
}
// parseIMAConfig extracts and validates IMA-specific configuration.
// Uses JSON marshal/unmarshal roundtrip so extra fields are ignored gracefully.
func parseIMAConfig(config *types.DataSourceConfig) (*Config, error) {
if config == nil {
return nil, fmt.Errorf("%w: config is nil", datasource.ErrInvalidConfig)
}
credBytes, err := json.Marshal(config.Credentials)
if err != nil {
return nil, fmt.Errorf("marshal credentials: %w", err)
}
var cfg Config
if err := json.Unmarshal(credBytes, &cfg); err != nil {
return nil, fmt.Errorf("parse ima credentials: %w", err)
}
if strings.TrimSpace(cfg.ClientID) == "" {
return nil, fmt.Errorf("%w: client_id is required", datasource.ErrInvalidCredentials)
}
if strings.TrimSpace(cfg.APIKey) == "" {
return nil, fmt.Errorf("%w: api_key is required", datasource.ErrInvalidCredentials)
}
if err := datasource.ValidateConnectorBaseURL(cfg.GetBaseURL()); err != nil {
return nil, err
}
return &cfg, nil
}
// apiEnvelope is the uniform IMA response wrapper: `{ code, msg, data }`.
// A non-zero code is a business-level error and MUST be surfaced to the user.
//
// Some published IMA references spell the same envelope `{ retcode, errmsg }`.
// Accepting both costs nothing and avoids the failure mode where a mismatch
// leaves Code at its zero value and every API error is silently read as
// success.
type apiEnvelope struct {
Code int `json:"code"`
Msg string `json:"msg"`
RetErr *int `json:"retcode"`
ErrMsg string `json:"errmsg"`
Data json.RawMessage `json:"data"`
}
// statusCode returns the business status code under either spelling.
func (e apiEnvelope) statusCode() int {
if e.Code != 0 {
return e.Code
}
if e.RetErr != nil {
return *e.RetErr
}
return 0
}
// message returns the human-readable error under either spelling.
func (e apiEnvelope) message() string {
if e.Msg != "" {
return e.Msg
}
return e.ErrMsg
}
// KnowledgeBaseInfo mirrors IMA's KnowledgeBaseInfo (get_knowledge_base).
type knowledgeBaseInfo struct {
ID string `json:"id"`
Name string `json:"name"`
CoverURL string `json:"cover_url"`
Description string `json:"description"`
RecommendedQuestions []string `json:"recommended_questions"`
}
// searchedKnowledgeBaseInfo — search_knowledge_base returns fewer fields.
type searchedKnowledgeBaseInfo struct {
ID string `json:"id"`
Name string `json:"name"`
CoverURL string `json:"cover_url"`
}
// searchKnowledgeBaseResp — search_knowledge_base response.
type searchKnowledgeBaseResp struct {
InfoList []searchedKnowledgeBaseInfo `json:"info_list"`
IsEnd bool `json:"is_end"`
NextCursor string `json:"next_cursor"`
}
// addableKnowledgeBaseInfo — get_addable_knowledge_base_list item.
// Contains only id + name (no cover_url), see api.md §7.
type addableKnowledgeBaseInfo struct {
ID string `json:"id"`
Name string `json:"name"`
}
// getAddableKnowledgeBaseListResp — get_addable_knowledge_base_list response.
// NOTE: the JSON field is `addable_knowledge_base_list`, not `info_list`.
type getAddableKnowledgeBaseListResp struct {
AddableKnowledgeBaseList []addableKnowledgeBaseInfo `json:"addable_knowledge_base_list"`
IsEnd bool `json:"is_end"`
NextCursor string `json:"next_cursor"`
}
// getKnowledgeBaseResp — get_knowledge_base response (map keyed by id).
type getKnowledgeBaseResp struct {
Infos map[string]knowledgeBaseInfo `json:"infos"`
}
// knowledgeInfo mirrors IMA's KnowledgeInfo (files, notes, etc.).
// IMPORTANT: IMA does not expose media_type in the list response, we must call
// get_media_info to discover it. However, folders are distinguished by presence
// of a non-empty folder_id — see folderInfo below.
type knowledgeInfo struct {
MediaID string `json:"media_id"`
Title string `json:"title"`
ParentFolderID string `json:"parent_folder_id"`
MediaType int32 `json:"media_type,omitempty"`
}
// folderInfo mirrors IMA's FolderInfo.
type folderInfo struct {
FolderID string `json:"folder_id"`
Name string `json:"name"`
FileNumber json.Number `json:"file_number"`
FolderNumber json.Number `json:"folder_number"`
ParentFolderID string `json:"parent_folder_id"`
IsTop bool `json:"is_top"`
}
// FileCount safely converts FileNumber to int64 (0 on empty/invalid).
func (f folderInfo) FileCount() int64 {
if f.FileNumber == "" {
return 0
}
n, err := f.FileNumber.Int64()
if err != nil {
return 0
}
return n
}
// FolderCount safely converts FolderNumber to int64 (0 on empty/invalid).
func (f folderInfo) FolderCount() int64 {
if f.FolderNumber == "" {
return 0
}
n, err := f.FolderNumber.Int64()
if err != nil {
return 0
}
return n
}
// getKnowledgeListResp — get_knowledge_list response.
// According to `references/api.md`, the response returns knowledge_list mixed
// with folders; we decode a loose JSON view so entries without a folder_id are
// treated as knowledge and those with folder_id are treated as folders.
type getKnowledgeListResp struct {
// KnowledgeList is a loose slice of raw JSON objects because IMA returns
// files and folders in the same array. We iterate and inspect fields.
KnowledgeList []json.RawMessage `json:"knowledge_list"`
IsEnd bool `json:"is_end"`
NextCursor string `json:"next_cursor"`
CurrentPath []folderInfo `json:"current_path"`
}
// urlInfo mirrors IMA's URLInfo (present on get_media_info for URL-backed media).
type urlInfo struct {
URL string `json:"url"`
Headers map[string]string `json:"headers"`
}
// notebookExtInfo — set on notes (media_type=11).
type notebookExtInfo struct {
NotebookID string `json:"notebook_id"`
}
// getMediaInfoResp — get_media_info response.
type getMediaInfoResp struct {
MediaType int32 `json:"media_type"`
URLInfo urlInfo `json:"url_info"`
NotebookExtInfo notebookExtInfo `json:"notebook_ext_info"`
}
// getDocContentResp — note/v1 get_doc_content response. Content is the note
// body; with target_content_format=0 it is plain text.
type getDocContentResp struct {
Content string `json:"content"`
}
// imaCursor tracks per-knowledge-base state seen during the last sync.
//
// KBLogical: { logical_key: media_id } — the authoritative source for both
// delete detection AND "replacement" detection. IMA reassigns media_id
// whenever a same-named file is replaced in place, so a logical key (see
// logicalKey) is what we treat as a document's stable identity.
// A logical_key present last sync but absent this sync ⇒ delete.
// A logical_key whose media_id changed ⇒ replacement (re-fetch content; the
// emitted FetchedItem carries the same ExternalID so the ingest layer's
// "existing external_id → delete-and-recreate" path (see
// datasource_service.go applyFetchedItem) treats it as an update.
//
// Only items that were successfully synced (or deterministically skipped)
// this cycle are recorded, so a transient download failure leaves the key
// absent and the next run retries it instead of treating it as unchanged.
//
// KBMedia is a legacy field retained only so cursors persisted by earlier
// builds still deserialize; it is never written or consulted.
type imaCursor struct {
LastSyncTime time.Time `json:"last_sync_time"`
KBLogical map[string]map[string]string `json:"kb_logical,omitempty"`
KBMedia map[string]map[string]string `json:"kb_media,omitempty"`
}
// logicalKey derives a stable identity for an IMA item that survives
// same-name replacement (which mints a new media_id server-side).
//
// The key is a short SHA-256 hex over (kb_id, parent_folder_id, title).
// Rationale:
//
// - kb_id keeps identities isolated across knowledge bases (same file
// copied into two KBs is two distinct docs from WeKnora's perspective).
// - parent_folder_id + title anchor the doc to its folder location and
// display name — the tuple IMA's own check_repeated_names uses to gate
// uploads (see api.md), so IMA's own uniqueness constraint enforces
// this key's uniqueness inside a KB.
//
// media_type is deliberately NOT part of the key: get_knowledge_list omits it
// (it only arrives later from get_media_info), so folding it in would mean
// every external_id shifts the day IMA starts populating the field in list
// responses — which reads as "every document deleted and re-added".
//
// Emitted as "ima_" + first 32 hex chars (128 bits) — small enough to fit
// well within Postgres varchar(64) external_id columns, still >> collision
// resistant for any realistic KB size.
func logicalKey(kbID, parentFolderID, title string) string {
h := sha256.New()
// Delimiter can't appear in any component because it's neither a valid
// IMA id character nor allowed in titles once we've folded to bytes.
_, _ = fmt.Fprintf(h, "%s\x1f%s\x1f%s", kbID, parentFolderID, title)
return "ima_" + hex.EncodeToString(h.Sum(nil))[:32]
}
// IMA MediaType enum (see api.md §MediaType).
const (
mediaTypePDF int32 = 1
mediaTypeWeb int32 = 2
mediaTypeWord int32 = 3
mediaTypePPT int32 = 4
mediaTypeExcel int32 = 5
mediaTypeMPArticle int32 = 6
mediaTypeMarkdown int32 = 7
mediaTypeImage int32 = 9
mediaTypeNote int32 = 11
mediaTypeAISession int32 = 12
mediaTypeTXT int32 = 13
mediaTypeXmind int32 = 14
mediaTypeAudio int32 = 15
mediaTypeVideo int32 = 16
mediaTypeHTML int32 = 20
mediaTypeEPUB int32 = 21
)
// extensionForMediaType maps IMA media types to a file extension.
// Returns "" when the media type has no downloadable body (web page, note, AI
// session, video-parse) — callers should skip those or treat them as URL-only.
func extensionForMediaType(t int32) string {
switch t {
case mediaTypePDF:
return "pdf"
case mediaTypeWord:
return "docx"
case mediaTypePPT:
return "pptx"
case mediaTypeExcel:
return "xlsx"
case mediaTypeMarkdown:
return "md"
case mediaTypeImage:
return "png"
case mediaTypeTXT:
return "txt"
case mediaTypeXmind:
return "xmind"
case mediaTypeAudio:
return "mp3"
case mediaTypeHTML:
return "html"
case mediaTypeEPUB:
return "epub"
default:
return ""
}
}
// extensionForContentType maps a response Content-Type back to a file
// extension, ignoring any "; charset=..." suffix. Returns "" for types we
// have no better guess for than the media_type-derived default.
//
// IMA reports every image as media_type=9 regardless of the real encoding, so
// the download's Content-Type is the only signal that distinguishes a JPEG
// from a PNG. Naming a JPEG ".png" breaks the extension-based file-type checks
// downstream in the ingest pipeline.
func extensionForContentType(contentType string) string {
mediaType := strings.ToLower(strings.TrimSpace(contentType))
if i := strings.Index(mediaType, ";"); i >= 0 {
mediaType = strings.TrimSpace(mediaType[:i])
}
switch mediaType {
case "image/png":
return "png"
case "image/jpeg", "image/jpg":
return "jpg"
case "image/gif":
return "gif"
case "image/webp":
return "webp"
case "image/bmp":
return "bmp"
default:
return ""
}
}
// mimeForExtension maps a file extension to a canonical MIME type.
// Used when the downloaded URL response omits (or misreports) Content-Type.
func mimeForExtension(ext string) string {
switch strings.ToLower(strings.TrimPrefix(ext, ".")) {
case "pdf":
return "application/pdf"
case "doc":
return "application/msword"
case "docx":
return "application/vnd.openxmlformats-officedocument.wordprocessingml.document"
case "ppt":
return "application/vnd.ms-powerpoint"
case "pptx":
return "application/vnd.openxmlformats-officedocument.presentationml.presentation"
case "xls":
return "application/vnd.ms-excel"
case "xlsx":
return "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"
case "md":
return "text/markdown"
case "txt":
return "text/plain"
case "html", "htm":
return "text/html"
case "png":
return "image/png"
case "jpg", "jpeg":
return "image/jpeg"
case "webp":
return "image/webp"
case "gif":
return "image/gif"
case "bmp":
return "image/bmp"
case "epub":
return "application/epub+zip"
case "xmind":
return "application/x-xmind"
case "mp3":
return "audio/mpeg"
case "m4a":
return "audio/x-m4a"
case "wav":
return "audio/wav"
case "aac":
return "audio/aac"
default:
return "application/octet-stream"
}
}
// isSkippableMediaType reports whether IMA offers no way to read this type's
// content. AI sessions are only ever referenced by session_id with no read
// endpoint, and video parses cannot even be added outside the IMA desktop app.
// Notes are NOT listed here: they are read through the note namespace, see
// fetchNote.
func isSkippableMediaType(t int32) bool {
return t == mediaTypeAISession || t == mediaTypeVideo
}
@@ -0,0 +1,196 @@
package ima
import (
"encoding/json"
"errors"
"strings"
"testing"
"github.com/Tencent/WeKnora/internal/datasource"
"github.com/Tencent/WeKnora/internal/types"
)
func TestLogicalKey_IsStableAndScoped(t *testing.T) {
base := logicalKey("kb1", "f1", "Doc")
if base != logicalKey("kb1", "f1", "Doc") {
t.Error("logicalKey must be deterministic")
}
if !strings.HasPrefix(base, "ima_") || len(base) != 36 {
t.Errorf("key = %q, want an ima_ prefix and 36 chars so it fits varchar(64)", base)
}
for _, other := range []string{
logicalKey("kb2", "f1", "Doc"),
logicalKey("kb1", "f2", "Doc"),
logicalKey("kb1", "f1", "Other"),
} {
if other == base {
t.Errorf("keys must differ across kb/folder/title, both were %q", base)
}
}
}
// TestLogicalKey_DelimiterIsUnambiguous guards against a key collision from
// naively concatenating the components.
func TestLogicalKey_DelimiterIsUnambiguous(t *testing.T) {
if logicalKey("kb", "a", "bc") == logicalKey("kb", "ab", "c") {
t.Error("component boundaries must not be ambiguous")
}
}
func TestParseIMAConfig(t *testing.T) {
tests := []struct {
name string
credentials map[string]interface{}
wantErr error
wantErrText string
}{
{
name: "valid",
credentials: map[string]interface{}{"client_id": "cid", "api_key": "key"},
},
{
name: "missing client_id",
credentials: map[string]interface{}{"api_key": "key"},
wantErr: datasource.ErrInvalidCredentials,
},
{
name: "blank api_key",
credentials: map[string]interface{}{"client_id": "cid", "api_key": " "},
wantErr: datasource.ErrInvalidCredentials,
},
{
name: "loopback base_url rejected by SSRF policy",
credentials: map[string]interface{}{"client_id": "cid", "api_key": "key", "base_url": "http://169.254.169.254"},
wantErrText: "SSRF",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
_, err := parseIMAConfig(&types.DataSourceConfig{Credentials: tt.credentials})
switch {
case tt.wantErr != nil:
if !errors.Is(err, tt.wantErr) {
t.Fatalf("err = %v, want %v", err, tt.wantErr)
}
case tt.wantErrText != "":
if err == nil || !strings.Contains(err.Error(), tt.wantErrText) {
t.Fatalf("err = %v, want it to mention %q", err, tt.wantErrText)
}
default:
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
}
})
}
}
func TestParseIMAConfig_NilConfig(t *testing.T) {
if _, err := parseIMAConfig(nil); !errors.Is(err, datasource.ErrInvalidConfig) {
t.Fatalf("err = %v, want ErrInvalidConfig", err)
}
}
func TestConfig_GetBaseURL(t *testing.T) {
tests := []struct {
in string
want string
}{
{"", DefaultBaseURL},
{" ", DefaultBaseURL},
{"https://ima.example.com/", "https://ima.example.com"},
{"ima.example.com", "https://ima.example.com"},
{"http://ima.example.com//", "http://ima.example.com"},
}
for _, tt := range tests {
if got := (&Config{BaseURL: tt.in}).GetBaseURL(); got != tt.want {
t.Errorf("GetBaseURL(%q) = %q, want %q", tt.in, got, tt.want)
}
}
}
func TestExtensionForContentType(t *testing.T) {
tests := []struct {
in string
want string
}{
{"image/jpeg", "jpg"},
{"image/JPEG", "jpg"},
{"image/png; charset=binary", "png"},
{"image/webp", "webp"},
{"application/pdf", ""},
{"", ""},
}
for _, tt := range tests {
if got := extensionForContentType(tt.in); got != tt.want {
t.Errorf("extensionForContentType(%q) = %q, want %q", tt.in, got, tt.want)
}
}
}
func TestSanitizeFileName(t *testing.T) {
if got := sanitizeFileName(""); got != "untitled" {
t.Errorf("empty name = %q, want untitled", got)
}
if got := sanitizeFileName(`a/b\c:d*e?f"g<h>i|j`); strings.ContainsAny(got, `/\:*?"<>|`) {
t.Errorf("filesystem-hostile characters survived: %q", got)
}
// A long CJK title must be truncated on a rune boundary, not mid-sequence.
long := strings.Repeat("知识", 200)
got := sanitizeFileName(long)
if len(got) > 200 {
t.Errorf("len = %d bytes, want <= 200", len(got))
}
if !strings.HasPrefix(long, got) {
t.Errorf("truncation corrupted the prefix: %q", got)
}
}
func TestIsSkippableMediaType(t *testing.T) {
for _, mt := range []int32{mediaTypeAISession, mediaTypeVideo} {
if !isSkippableMediaType(mt) {
t.Errorf("media type %d should be skippable", mt)
}
}
// Notes are read through the note namespace rather than skipped.
for _, mt := range []int32{mediaTypePDF, mediaTypeMarkdown, mediaTypeWeb, mediaTypeNote} {
if isSkippableMediaType(mt) {
t.Errorf("media type %d should not be skippable", mt)
}
}
}
// TestAPIEnvelope_AcceptsBothSpellings covers the retcode/errmsg variant some
// IMA references describe: if only `code` were read, a non-zero retcode would
// decode as 0 and every API error would be silently treated as success.
func TestAPIEnvelope_AcceptsBothSpellings(t *testing.T) {
tests := []struct {
name string
body string
wantCode int
wantMsg string
}{
{"code/msg", `{"code":110030,"msg":"无权限"}`, 110030, "无权限"},
{"retcode/errmsg", `{"retcode":110030,"errmsg":"无权限"}`, 110030, "无权限"},
{"success", `{"code":0,"msg":"ok"}`, 0, "ok"},
{"retcode success", `{"retcode":0,"errmsg":"成功"}`, 0, "成功"},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
var env apiEnvelope
if err := json.Unmarshal([]byte(tt.body), &env); err != nil {
t.Fatalf("unmarshal: %v", err)
}
if got := env.statusCode(); got != tt.wantCode {
t.Errorf("statusCode() = %d, want %d", got, tt.wantCode)
}
if got := env.message(); got != tt.wantMsg {
t.Errorf("message() = %q, want %q", got, tt.wantMsg)
}
})
}
}
@@ -0,0 +1,530 @@
package notion
import (
"bytes"
"context"
"encoding/json"
"fmt"
"io"
"net/http"
"strconv"
"strings"
"time"
"golang.org/x/time/rate"
"github.com/Tencent/WeKnora/internal/datasource"
"github.com/Tencent/WeKnora/internal/logger"
"github.com/Tencent/WeKnora/internal/utils"
)
// notionClient wraps the Notion API with rate limiting and retry logic.
type notionClient struct {
token string
httpClient *http.Client
limiter *rate.Limiter
baseURL string
}
// newClient creates a new Notion API client.
func newClient(token, baseURL string) (*notionClient, error) {
if baseURL == "" {
baseURL = DefaultBaseURL
}
if err := datasource.ValidateConnectorBaseURL(baseURL); err != nil {
return nil, err
}
return &notionClient{
token: token,
httpClient: datasource.NewConnectorHTTPClient(30 * time.Second),
limiter: rate.NewLimiter(rate.Limit(3), 3),
baseURL: baseURL,
}, nil
}
const maxRetries = 3
// sleepWithContext pauses for the given duration, returning early if ctx is cancelled.
func sleepWithContext(ctx context.Context, d time.Duration) error {
t := time.NewTimer(d)
defer t.Stop()
select {
case <-t.C:
return nil
case <-ctx.Done():
return ctx.Err()
}
}
// doRequest performs an authenticated, rate-limited HTTP request to the Notion API.
func (c *notionClient) doRequest(ctx context.Context, method, path string, body interface{}) ([]byte, error) {
if err := c.limiter.Wait(ctx); err != nil {
return nil, fmt.Errorf("rate limiter: %w", err)
}
var bodyReader io.Reader
if body != nil {
bodyBytes, err := json.Marshal(body)
if err != nil {
return nil, fmt.Errorf("marshal request body: %w", err)
}
bodyReader = bytes.NewReader(bodyBytes)
}
req, err := http.NewRequestWithContext(ctx, method, c.baseURL+path, bodyReader)
if err != nil {
return nil, fmt.Errorf("create request: %w", err)
}
req.Header.Set("Authorization", "Bearer "+c.token)
req.Header.Set("Notion-Version", NotionAPIVersion)
req.Header.Set("Content-Type", "application/json")
var lastErr error
for attempt := 0; attempt <= maxRetries; attempt++ {
if attempt > 0 {
if body != nil {
bodyBytes, _ := json.Marshal(body)
req.Body = io.NopCloser(bytes.NewReader(bodyBytes))
}
}
resp, err := c.httpClient.Do(req)
if err != nil {
lastErr = err
if attempt < maxRetries {
if sErr := sleepWithContext(ctx, time.Duration(1<<attempt)*time.Second); sErr != nil {
return nil, sErr
}
continue
}
break
}
respBody, err := io.ReadAll(resp.Body)
resp.Body.Close()
if err != nil {
return nil, fmt.Errorf("read response: %w", err)
}
switch {
case resp.StatusCode >= 200 && resp.StatusCode < 300:
return respBody, nil
case resp.StatusCode == 401 || resp.StatusCode == 403:
return nil, fmt.Errorf("%w: %s", datasource.ErrInvalidCredentials, string(respBody))
case resp.StatusCode == 404:
return nil, fmt.Errorf("%w: %s", datasource.ErrResourceNotFound, path)
case resp.StatusCode == 429:
retryAfter := resp.Header.Get("Retry-After")
wait := 1 * time.Second
if secs, err := strconv.ParseFloat(retryAfter, 64); err == nil && secs > 0 {
wait = time.Duration(secs * float64(time.Second))
}
logger.Warnf(ctx, "[Notion] rate limited, retry after %v (attempt %d/%d)", wait, attempt+1, maxRetries)
lastErr = fmt.Errorf("rate limited: %s", string(respBody))
if attempt < maxRetries {
if sErr := sleepWithContext(ctx, wait); sErr != nil {
return nil, sErr
}
continue
}
case resp.StatusCode >= 500:
lastErr = fmt.Errorf("server error %d: %s", resp.StatusCode, string(respBody))
if attempt < maxRetries {
if sErr := sleepWithContext(ctx, time.Duration(1<<attempt)*time.Second); sErr != nil {
return nil, sErr
}
continue
}
default:
return nil, fmt.Errorf("unexpected status %d: %s", resp.StatusCode, string(respBody))
}
}
if lastErr != nil {
return nil, fmt.Errorf("%w: %v", datasource.ErrFetchFailed, lastErr)
}
return nil, datasource.ErrFetchFailed
}
// Ping verifies the API token is valid by calling GET /v1/users/me.
func (c *notionClient) Ping(ctx context.Context) error {
_, err := c.doRequest(ctx, http.MethodGet, "/v1/users/me", nil)
return err
}
// SearchPages returns all pages and databases accessible to the integration.
func (c *notionClient) SearchPages(ctx context.Context) ([]notionPage, error) {
return c.paginatePages(ctx, http.MethodPost, "/v1/search")
}
// GetPage retrieves a single page by ID.
func (c *notionClient) GetPage(ctx context.Context, pageID string) (*notionPage, error) {
respBody, err := c.doRequest(ctx, http.MethodGet, "/v1/pages/"+pageID, nil)
if err != nil {
return nil, err
}
var page notionPage
if err := json.Unmarshal(respBody, &page); err != nil {
return nil, fmt.Errorf("unmarshal page: %w", err)
}
page.Title = extractTitle(&page)
return &page, nil
}
// databaseInfo holds both the metadata and the data source ID from a single API call.
type databaseInfo struct {
Page notionPage
DataSourceID string
}
// GetDatabaseInfo retrieves a database container by ID, returning both
// the page metadata and the primary data source ID in a single API call.
func (c *notionClient) GetDatabaseInfo(ctx context.Context, dbID string) (*databaseInfo, error) {
respBody, err := c.doRequest(ctx, http.MethodGet, "/v1/databases/"+dbID, nil)
if err != nil {
return nil, err
}
var db notionPage
if err := json.Unmarshal(respBody, &db); err != nil {
return nil, fmt.Errorf("unmarshal database: %w", err)
}
db.Title = extractTitle(&db)
var dsResult struct {
DataSources []struct {
ID string `json:"id"`
} `json:"data_sources"`
}
if err := json.Unmarshal(respBody, &dsResult); err != nil {
return nil, fmt.Errorf("unmarshal data_sources: %w", err)
}
dsID := ""
if len(dsResult.DataSources) > 0 {
dsID = dsResult.DataSources[0].ID
}
return &databaseInfo{Page: db, DataSourceID: dsID}, nil
}
// GetDataSourceInfo retrieves a data source by ID, returning its metadata.
// In API 2025-09-03+, data_source objects hold the schema/properties and are
// the target for record queries. The response includes database_parent to
// locate the database in the workspace hierarchy.
func (c *notionClient) GetDataSourceInfo(ctx context.Context, dsID string) (*notionPage, error) {
respBody, err := c.doRequest(ctx, http.MethodGet, "/v1/data_sources/"+dsID, nil)
if err != nil {
return nil, err
}
var ds notionPage
if err := json.Unmarshal(respBody, &ds); err != nil {
return nil, fmt.Errorf("unmarshal data_source: %w", err)
}
ds.Title = extractTitle(&ds)
return &ds, nil
}
// GetBlockChildrenFlat fetches only the direct children of a block (no recursion).
// Used by discoverPages to quickly scan for child_page/child_database without
// fetching the full block tree content.
func (c *notionClient) GetBlockChildrenFlat(ctx context.Context, blockID string) ([]notionBlock, error) {
var allBlocks []notionBlock
var startCursor string
for {
path := fmt.Sprintf("/v1/blocks/%s/children", blockID)
if startCursor != "" {
path += "?start_cursor=" + startCursor
}
respBody, err := c.doRequest(ctx, http.MethodGet, path, nil)
if err != nil {
return nil, fmt.Errorf("get block children for %s: %w", blockID, err)
}
var resp paginatedResponse
if err := json.Unmarshal(respBody, &resp); err != nil {
return nil, fmt.Errorf("unmarshal block children response: %w", err)
}
var blocks []notionBlock
if err := json.Unmarshal(resp.Results, &blocks); err != nil {
return nil, fmt.Errorf("unmarshal blocks: %w", err)
}
allBlocks = append(allBlocks, blocks...)
if !resp.HasMore || resp.NextCursor == "" {
break
}
startCursor = resp.NextCursor
}
return allBlocks, nil
}
const maxBlockDepth = 5 // Limit recursion depth — deeper content has diminishing value for knowledge bases
const maxBlocksPerPage = 1000 // Limit total blocks fetched per page to prevent runaway API calls
// GetBlockChildrenAll recursively fetches all blocks under a given block ID,
// building a tree structure with Children populated for blocks with has_children=true.
// child_page and child_database blocks are NOT recursed into (handled by connector layer).
// Recursion is limited to maxBlockDepth to prevent excessive API calls on complex pages.
func (c *notionClient) GetBlockChildrenAll(ctx context.Context, blockID string) ([]notionBlock, error) {
return c.getBlockChildrenRecursive(ctx, blockID, 0)
}
func (c *notionClient) getBlockChildrenRecursive(ctx context.Context, blockID string, depth int) ([]notionBlock, error) {
var allBlocks []notionBlock
var startCursor string
for {
path := fmt.Sprintf("/v1/blocks/%s/children", blockID)
if startCursor != "" {
path += "?start_cursor=" + startCursor
}
respBody, err := c.doRequest(ctx, http.MethodGet, path, nil)
if err != nil {
return nil, fmt.Errorf("get block children for %s: %w", blockID, err)
}
var resp paginatedResponse
if err := json.Unmarshal(respBody, &resp); err != nil {
return nil, fmt.Errorf("unmarshal block children response: %w", err)
}
var blocks []notionBlock
if err := json.Unmarshal(resp.Results, &blocks); err != nil {
return nil, fmt.Errorf("unmarshal blocks: %w", err)
}
allBlocks = append(allBlocks, blocks...)
if len(allBlocks) >= maxBlocksPerPage || !resp.HasMore || resp.NextCursor == "" {
break
}
startCursor = resp.NextCursor
}
if depth >= maxBlockDepth {
return allBlocks, nil
}
for i := range allBlocks {
if !allBlocks[i].HasChildren {
continue
}
// Skip block types that don't contribute useful content
switch allBlocks[i].Type {
case "child_page", "child_database", "unsupported", "template", "breadcrumb", "table_of_contents":
continue
}
children, err := c.getBlockChildrenRecursive(ctx, allBlocks[i].ID, depth+1)
if err != nil {
logger.Warnf(ctx, "[Notion] failed to get children for block %s (depth %d): %v", allBlocks[i].ID, depth, err)
continue
}
allBlocks[i].Children = children
}
return allBlocks, nil
}
// QueryDatabaseAll retrieves all records from a database via POST /v1/data_sources/{id}/query.
// Accepts either a data_source_id (from search) or a database_id (from child_database blocks).
// For database_ids, resolves to data_source_id via GET /v1/databases/{id}.
func (c *notionClient) QueryDatabaseAll(ctx context.Context, id string) ([]notionPage, error) {
// Try as data_source_id directly
records, err := c.paginatePages(ctx, http.MethodPost, fmt.Sprintf("/v1/data_sources/%s/query", id))
if err == nil {
return records, nil
}
// If 404, id might be a database container ID — resolve to data_source_id
info, dbErr := c.GetDatabaseInfo(ctx, id)
if dbErr != nil {
return nil, fmt.Errorf("query database %s: not a data_source (%v) and not a database (%v)", id, err, dbErr)
}
if info.DataSourceID == "" {
return nil, fmt.Errorf("database %s has no data sources", id)
}
return c.paginatePages(ctx, http.MethodPost, fmt.Sprintf("/v1/data_sources/%s/query", info.DataSourceID))
}
// ResolveBlock re-fetches a single block to resolve file_upload URLs.
// When a block contains a file_upload type, re-fetching it returns the resolved
// download URL (temporary S3 signed URL, 1-hour expiry).
func (c *notionClient) ResolveBlock(ctx context.Context, blockID string) (*notionBlock, error) {
respBody, err := c.doRequest(ctx, http.MethodGet, "/v1/blocks/"+blockID, nil)
if err != nil {
return nil, err
}
var block notionBlock
if err := json.Unmarshal(respBody, &block); err != nil {
return nil, fmt.Errorf("unmarshal block: %w", err)
}
return &block, nil
}
const maxDownloadSize = 100 * 1024 * 1024 // 100MB — prevent OOM from oversized files
// DownloadFile downloads a file from the given URL (typically an S3 signed URL).
// Does not go through the rate limiter since it's not a Notion API call.
func (c *notionClient) DownloadFile(ctx context.Context, fileURL string) ([]byte, error) {
if err := utils.ValidateURLForSSRF(fileURL); err != nil {
return nil, fmt.Errorf("attachment URL rejected: %w", err)
}
req, err := http.NewRequestWithContext(ctx, http.MethodGet, fileURL, nil)
if err != nil {
return nil, fmt.Errorf("create download request: %w", err)
}
var lastErr error
for attempt := 0; attempt <= maxRetries; attempt++ {
resp, err := c.httpClient.Do(req)
if err != nil {
lastErr = err
if attempt < maxRetries {
if sErr := sleepWithContext(ctx, time.Duration(1<<attempt)*time.Second); sErr != nil {
return nil, sErr
}
continue
}
break
}
if resp.StatusCode != http.StatusOK {
resp.Body.Close()
lastErr = fmt.Errorf("download failed with status %d", resp.StatusCode)
if resp.StatusCode >= 500 && attempt < maxRetries {
if sErr := sleepWithContext(ctx, time.Duration(1<<attempt)*time.Second); sErr != nil {
return nil, sErr
}
continue
}
break
}
data, err := io.ReadAll(io.LimitReader(resp.Body, maxDownloadSize+1))
resp.Body.Close()
if err != nil {
return nil, err
}
if int64(len(data)) > maxDownloadSize {
return nil, fmt.Errorf("file exceeds maximum download size (%d MB)", maxDownloadSize/(1024*1024))
}
return data, nil
}
return nil, fmt.Errorf("download file: %w", lastErr)
}
// --- Shared pagination helper ---
// paginatePages fetches all pages from a paginated Notion API endpoint.
func (c *notionClient) paginatePages(ctx context.Context, method, path string) ([]notionPage, error) {
var allPages []notionPage
var startCursor string
for {
body := map[string]interface{}{
"page_size": 100,
}
if startCursor != "" {
body["start_cursor"] = startCursor
}
var respBody []byte
var err error
if method == http.MethodPost {
respBody, err = c.doRequest(ctx, method, path, body)
} else {
p := path
if startCursor != "" {
p += "?start_cursor=" + startCursor + "&page_size=100"
}
respBody, err = c.doRequest(ctx, method, p, nil)
}
if err != nil {
return nil, fmt.Errorf("paginate %s: %w", path, err)
}
var resp paginatedResponse
if err := json.Unmarshal(respBody, &resp); err != nil {
return nil, fmt.Errorf("unmarshal paginated response: %w", err)
}
var pages []notionPage
if err := json.Unmarshal(resp.Results, &pages); err != nil {
return nil, fmt.Errorf("unmarshal page results: %w", err)
}
for i := range pages {
pages[i].Title = extractTitle(&pages[i])
}
allPages = append(allPages, pages...)
if !resp.HasMore || resp.NextCursor == "" {
break
}
startCursor = resp.NextCursor
}
return allPages, nil
}
// --- Title extraction helpers ---
func extractTitle(page *notionPage) string {
// Try extracting from properties (works for pages)
if page.RawProperties != nil {
var props map[string]json.RawMessage
if err := json.Unmarshal(page.RawProperties, &props); err == nil {
for _, propRaw := range props {
var prop struct {
Type string `json:"type"`
Title []struct {
PlainText string `json:"plain_text"`
} `json:"title"`
}
if err := json.Unmarshal(propRaw, &prop); err != nil {
continue
}
if prop.Type == "title" && len(prop.Title) > 0 {
return joinPlainText(prop.Title)
}
}
}
}
// Fallback: top-level title array (works for databases)
if page.RawTitle != nil {
var titleSegments []struct {
PlainText string `json:"plain_text"`
}
if err := json.Unmarshal(page.RawTitle, &titleSegments); err == nil && len(titleSegments) > 0 {
return joinPlainText(titleSegments)
}
}
return ""
}
func joinPlainText(segments []struct {
PlainText string `json:"plain_text"`
}) string {
var sb strings.Builder
for _, s := range segments {
sb.WriteString(s.PlainText)
}
return sb.String()
}
@@ -0,0 +1,363 @@
package notion
import (
"context"
"encoding/json"
"net/http"
"net/http/httptest"
"testing"
"time"
secutils "github.com/Tencent/WeKnora/internal/utils"
)
func mustTestClient(t *testing.T, token, baseURL string) *notionClient {
t.Helper()
t.Setenv("SSRF_WHITELIST", "127.0.0.1,localhost")
secutils.ResetSSRFWhitelistForTest()
client, err := newClient(token, baseURL)
if err != nil {
t.Fatalf("newClient: %v", err)
}
return client
}
// fakeNotion builds an httptest.Server that emulates relevant Notion API endpoints.
func fakeNotion() (*httptest.Server, *Config) {
mux := http.NewServeMux()
// GET /v1/users/me — auth check
mux.HandleFunc("/v1/users/me", func(w http.ResponseWriter, r *http.Request) {
if r.Header.Get("Authorization") != "Bearer test-token" {
w.WriteHeader(401)
json.NewEncoder(w).Encode(map[string]interface{}{
"status": 401,
"code": "unauthorized",
"message": "API token is invalid.",
})
return
}
json.NewEncoder(w).Encode(map[string]interface{}{
"object": "user",
"id": "bot-user-id",
"type": "bot",
"name": "Test Integration",
})
})
// POST /v1/search — return pages and databases
mux.HandleFunc("/v1/search", func(w http.ResponseWriter, r *http.Request) {
json.NewEncoder(w).Encode(map[string]interface{}{
"object": "list",
"results": []interface{}{
map[string]interface{}{
"id": "page-1",
"object": "page",
"url": "https://notion.so/Page-1",
"last_edited_time": "2026-01-15T10:00:00.000Z",
"in_trash": false,
"parent": map[string]interface{}{"type": "workspace", "workspace": true},
"properties": map[string]interface{}{
"title": map[string]interface{}{
"type": "title",
"title": []interface{}{
map[string]interface{}{"plain_text": "Test Page"},
},
},
},
},
map[string]interface{}{
"id": "db-1",
"object": "data_source",
"url": "https://notion.so/DB-1",
"last_edited_time": "2026-01-16T10:00:00.000Z",
"in_trash": false,
"parent": map[string]interface{}{"type": "workspace", "workspace": true},
"title": []interface{}{
map[string]interface{}{"plain_text": "Test Database"},
},
},
},
"has_more": false,
"next_cursor": nil,
})
})
// GET /v1/blocks/{id}/children — return blocks
mux.HandleFunc("/v1/blocks/page-1/children", func(w http.ResponseWriter, r *http.Request) {
json.NewEncoder(w).Encode(map[string]interface{}{
"object": "list",
"results": []interface{}{
map[string]interface{}{
"id": "blk-1",
"type": "paragraph",
"has_children": false,
"paragraph": map[string]interface{}{
"rich_text": []interface{}{
map[string]interface{}{
"type": "text",
"plain_text": "Hello world",
"text": map[string]interface{}{"content": "Hello world"},
"annotations": map[string]interface{}{
"bold": false, "italic": false, "strikethrough": false,
"underline": false, "code": false, "color": "default",
},
},
},
},
},
map[string]interface{}{
"id": "blk-2",
"type": "child_page",
"has_children": true,
"child_page": map[string]interface{}{"title": "Sub Page"},
},
},
"has_more": false,
"next_cursor": nil,
})
})
// GET /v1/pages/{id}
mux.HandleFunc("/v1/pages/page-1", func(w http.ResponseWriter, r *http.Request) {
json.NewEncoder(w).Encode(map[string]interface{}{
"id": "page-1",
"object": "page",
"url": "https://notion.so/Page-1",
"last_edited_time": "2026-01-15T10:00:00.000Z",
"in_trash": false,
"parent": map[string]interface{}{"type": "workspace", "workspace": true},
"properties": map[string]interface{}{
"title": map[string]interface{}{
"type": "title",
"title": []interface{}{
map[string]interface{}{"plain_text": "Test Page"},
},
},
},
})
})
// GET /v1/databases/{id} — returns container with data_sources array (2025-09-03+)
mux.HandleFunc("/v1/databases/db-1", func(w http.ResponseWriter, r *http.Request) {
json.NewEncoder(w).Encode(map[string]interface{}{
"id": "db-1",
"object": "database",
"url": "https://notion.so/DB-1",
"last_edited_time": "2026-01-16T10:00:00.000Z",
"in_trash": false,
"parent": map[string]interface{}{"type": "workspace", "workspace": true},
"title": []interface{}{map[string]interface{}{"plain_text": "Test Database"}},
"data_sources": []interface{}{
map[string]interface{}{"id": "ds-1", "name": "Default"},
},
})
})
// GET /v1/data_sources/{id} — returns schema/properties
mux.HandleFunc("/v1/data_sources/ds-1", func(w http.ResponseWriter, r *http.Request) {
if r.Method == http.MethodGet {
json.NewEncoder(w).Encode(map[string]interface{}{
"id": "ds-1",
"object": "data_source",
"properties": map[string]interface{}{
"Name": map[string]interface{}{"type": "title", "title": map[string]interface{}{}},
"Status": map[string]interface{}{"type": "select", "select": map[string]interface{}{}},
},
})
return
}
http.NotFound(w, r)
})
// POST /v1/data_sources/{id}/query — return database records (2025-09-03+)
mux.HandleFunc("/v1/data_sources/ds-1/query", func(w http.ResponseWriter, r *http.Request) {
json.NewEncoder(w).Encode(map[string]interface{}{
"object": "list",
"results": []interface{}{
map[string]interface{}{
"id": "record-1",
"object": "page",
"url": "https://notion.so/Record-1",
"last_edited_time": "2026-01-17T10:00:00.000Z",
"in_trash": false,
"parent": map[string]interface{}{"type": "data_source_id", "data_source_id": "ds-1"},
"properties": map[string]interface{}{
"Name": map[string]interface{}{
"type": "title",
"title": []interface{}{
map[string]interface{}{"plain_text": "Record One"},
},
},
"Status": map[string]interface{}{
"type": "select",
"select": map[string]interface{}{"name": "Done"},
},
},
},
},
"has_more": false,
"next_cursor": nil,
})
})
// GET /v1/pages/record-1 — direct fetch of a database row (single-select case).
// fetchPage must recognize both database_id (older responses) and
// data_source_id (2025-09-03+) as a record and route to buildRecordItem.
mux.HandleFunc("/v1/pages/record-1", func(w http.ResponseWriter, r *http.Request) {
json.NewEncoder(w).Encode(map[string]interface{}{
"id": "record-1",
"object": "page",
"url": "https://notion.so/Record-1",
"last_edited_time": "2026-01-17T10:00:00.000Z",
"in_trash": false,
"parent": map[string]interface{}{"type": "data_source_id", "data_source_id": "ds-1"},
"properties": map[string]interface{}{
"Name": map[string]interface{}{
"type": "title",
"title": []interface{}{
map[string]interface{}{"plain_text": "Record One"},
},
},
"Status": map[string]interface{}{
"type": "select",
"select": map[string]interface{}{"name": "Done"},
},
},
})
})
// GET /v1/blocks/record-1/children — empty blocks for database record
mux.HandleFunc("/v1/blocks/record-1/children", func(w http.ResponseWriter, r *http.Request) {
json.NewEncoder(w).Encode(map[string]interface{}{
"object": "list",
"results": []interface{}{},
"has_more": false,
"next_cursor": nil,
})
})
ts := httptest.NewServer(mux)
cfg := &Config{APIKey: "test-token"}
return ts, cfg
}
func TestClientPing(t *testing.T) {
ts, cfg := fakeNotion()
defer ts.Close()
client := mustTestClient(t, cfg.APIKey, ts.URL)
if err := client.Ping(context.Background()); err != nil {
t.Fatalf("Ping() error: %v", err)
}
}
func TestClientPing_InvalidToken(t *testing.T) {
ts, _ := fakeNotion()
defer ts.Close()
client := mustTestClient(t, "wrong-token", ts.URL)
err := client.Ping(context.Background())
if err == nil {
t.Fatal("expected error for invalid token")
}
}
func TestClientSearchPages(t *testing.T) {
ts, cfg := fakeNotion()
defer ts.Close()
client := mustTestClient(t, cfg.APIKey, ts.URL)
pages, err := client.SearchPages(context.Background())
if err != nil {
t.Fatalf("SearchPages() error: %v", err)
}
if len(pages) != 2 {
t.Fatalf("expected 2 pages, got %d", len(pages))
}
// First result is a page
if pages[0].Object != "page" || pages[0].ID != "page-1" {
t.Errorf("page[0] = %+v", pages[0])
}
// Second result is a data_source (databases are returned as object="data_source"
// in API 2025-09-03+ which is what the mock simulates)
if pages[1].Object != "data_source" || pages[1].ID != "db-1" {
t.Errorf("page[1] = %+v", pages[1])
}
}
func TestClientGetBlockChildrenAll(t *testing.T) {
ts, cfg := fakeNotion()
defer ts.Close()
client := mustTestClient(t, cfg.APIKey, ts.URL)
blocks, err := client.GetBlockChildrenAll(context.Background(), "page-1")
if err != nil {
t.Fatalf("GetBlockChildrenAll() error: %v", err)
}
if len(blocks) != 2 {
t.Fatalf("expected 2 blocks, got %d", len(blocks))
}
if blocks[0].Type != "paragraph" {
t.Errorf("block[0].Type = %q", blocks[0].Type)
}
if blocks[1].Type != "child_page" {
t.Errorf("block[1].Type = %q", blocks[1].Type)
}
// child_page should NOT have its children recursed
if blocks[1].Children != nil {
t.Error("child_page children should not be recursed by client")
}
}
func TestClientGetPage(t *testing.T) {
ts, cfg := fakeNotion()
defer ts.Close()
client := mustTestClient(t, cfg.APIKey, ts.URL)
page, err := client.GetPage(context.Background(), "page-1")
if err != nil {
t.Fatalf("GetPage() error: %v", err)
}
if page.ID != "page-1" {
t.Errorf("ID = %q", page.ID)
}
if page.Title != "Test Page" {
t.Errorf("Title = %q, want %q", page.Title, "Test Page")
}
}
func TestClientQueryDatabaseAll(t *testing.T) {
ts, cfg := fakeNotion()
defer ts.Close()
client := mustTestClient(t, cfg.APIKey, ts.URL)
records, err := client.QueryDatabaseAll(context.Background(), "db-1")
if err != nil {
t.Fatalf("QueryDatabaseAll() error: %v", err)
}
if len(records) != 1 {
t.Fatalf("expected 1 record, got %d", len(records))
}
if records[0].ID != "record-1" {
t.Errorf("record ID = %q", records[0].ID)
}
}
// Ensure time import is used (referenced by fakeNotion indirectly via time.Time in types).
var _ = time.Now
func TestDownloadFile_RejectsLoopbackURL(t *testing.T) {
secutils.ResetSSRFWhitelistForTest()
t.Cleanup(secutils.ResetSSRFWhitelistForTest)
client, err := newClient("test-token", "https://api.notion.com")
if err != nil {
t.Fatalf("newClient: %v", err)
}
_, err = client.DownloadFile(context.Background(), "http://127.0.0.1/secret")
if err == nil {
t.Fatal("expected loopback attachment URL to be rejected")
}
}
@@ -0,0 +1,974 @@
package notion
import (
"context"
"encoding/json"
"fmt"
"sort"
"strings"
"time"
"github.com/Tencent/WeKnora/internal/logger"
"github.com/Tencent/WeKnora/internal/types"
)
const (
contentTypeMarkdown = "text/markdown"
objectTypePage = "page"
objectTypeDatabase = "database"
objectTypeAttachment = "attachment"
defaultUntitledName = "Untitled"
)
// Connector implements the datasource.Connector interface for Notion.
type Connector struct{}
// NewConnector creates a new Notion connector.
func NewConnector() *Connector {
return &Connector{}
}
// Type returns the connector type identifier.
func (c *Connector) Type() string {
return types.ConnectorTypeNotion
}
// Validate verifies that the Notion configuration is valid by testing connectivity.
func (c *Connector) Validate(ctx context.Context, config *types.DataSourceConfig) error {
notionCfg, err := parseNotionConfig(config)
if err != nil {
return err
}
client, err := newClient(notionCfg.APIKey, extractBaseURL(config))
if err != nil {
return err
}
return client.Ping(ctx)
}
// ResolveResourceAncestors has nothing to do for Notion: ListResources already
// returns the full tree with parent links, so any pre-existing selection is
// already present and revealed by the picker without on-demand loading.
func (c *Connector) ResolveResourceAncestors(
ctx context.Context, config *types.DataSourceConfig, resourceIDs []string,
) ([]string, error) {
return []string{}, nil
}
// ListResources lists all accessible Notion pages and databases as selectable resources.
// Returns all objects with parent-child relationships populated, allowing the frontend
// to render a tree view. Root objects have empty ParentID.
func (c *Connector) ListResources(
ctx context.Context, config *types.DataSourceConfig, parentID string,
) ([]types.Resource, error) {
// Notion returns the full hierarchy (with parent_id populated) in a single
// call, so children are already delivered with the root listing. Lazy-load
// requests for a specific parent therefore have nothing extra to return.
if parentID != "" {
return []types.Resource{}, nil
}
notionCfg, err := parseNotionConfig(config)
if err != nil {
return nil, err
}
client, err := newClient(notionCfg.APIKey, extractBaseURL(config))
if err != nil {
return nil, err
}
pages, err := client.SearchPages(ctx)
if err != nil {
return nil, fmt.Errorf("search notion pages: %w", err)
}
allIDs := make(map[string]bool, len(pages))
for _, p := range pages {
allIDs[p.ID] = true
}
// Precompute effective parent for each page.
// data_source objects need database_parent (their `parent` points at the
// database container, not the workspace location); see resolveParentID.
parentOf := make(map[string]string, len(pages))
for _, p := range pages {
if p.InTrash {
continue
}
parentOf[p.ID] = resolveParentID(p, allIDs)
}
childrenCount := make(map[string]int)
for _, pid := range parentOf {
if pid != "" {
childrenCount[pid]++
}
}
var resources []types.Resource
for _, p := range pages {
if p.InTrash {
continue
}
resourceType := objectTypePage
if p.isDatabase() {
resourceType = objectTypeDatabase
}
resources = append(resources, types.Resource{
ExternalID: p.ID,
Name: p.Title,
Type: resourceType,
URL: p.URL,
ParentID: parentOf[p.ID],
HasChildren: childrenCount[p.ID] > 0,
})
}
return resources, nil
}
// FetchAll performs a full sync of all documents from the specified resources.
func (c *Connector) FetchAll(ctx context.Context, config *types.DataSourceConfig, resourceIDs []string) ([]types.FetchedItem, error) {
notionCfg, err := parseNotionConfig(config)
if err != nil {
return nil, err
}
client, err := newClient(notionCfg.APIKey, extractBaseURL(config))
if err != nil {
return nil, err
}
visited := c.excludedSetFromListResources(ctx, config, resourceIDs)
var allItems []types.FetchedItem
for _, resourceID := range resourceIDs {
page, err := client.GetPage(ctx, resourceID)
if err == nil {
allItems = append(allItems, c.fetchPage(ctx, client, page, visited)...)
continue
}
// Not a page — treat as database/data_source.
// fetchDatabase handles both data_source IDs (from search) and
// database container IDs (from child_database blocks).
items := c.fetchDatabase(ctx, client, resourceID, visited)
if len(items) == 0 {
logger.Warnf(ctx, "[Notion] failed to fetch resource %s as page or database", resourceID)
}
allItems = append(allItems, items...)
}
return allItems, nil
}
// FetchIncremental performs an incremental sync based on cursor state.
// On first sync (nil cursor), delegates to FetchAll for efficiency (no discovery needed).
func (c *Connector) FetchIncremental(ctx context.Context, config *types.DataSourceConfig, cursor *types.SyncCursor) ([]types.FetchedItem, *types.SyncCursor, error) {
notionCfg, err := parseNotionConfig(config)
if err != nil {
return nil, nil, err
}
resourceIDs := config.ResourceIDs
if len(resourceIDs) == 0 {
return nil, nil, fmt.Errorf("no resource IDs configured")
}
client, err := newClient(notionCfg.APIKey, extractBaseURL(config))
if err != nil {
return nil, nil, err
}
// Parse previous cursor
var prevCursor notionCursor
if cursor != nil && cursor.ConnectorCursor != nil {
cursorBytes, _ := json.Marshal(cursor.ConnectorCursor)
_ = json.Unmarshal(cursorBytes, &prevCursor)
}
// First sync: use FetchAll then build cursor from Search API (one call, no block traversal)
isFirstSync := len(prevCursor.PageEditTimes) == 0
if isFirstSync {
logger.Infof(ctx, "[Notion] first sync, using FetchAll for %d resources", len(resourceIDs))
items, err := c.FetchAll(ctx, config, resourceIDs)
if err != nil {
return nil, nil, err
}
// Build cursor from fetched items' UpdatedAt timestamps.
// Record-level edit times are tracked individually (object_type == "page").
// Database container IDs are also added so incremental sync can detect
// whether a database actually changed, avoiding full record queries every cycle.
newEditTimes := make(map[string]time.Time)
for _, item := range items {
if item.Metadata["object_type"] == objectTypePage {
newEditTimes[item.ExternalID] = item.UpdatedAt
}
}
// Ensure all selected resourceIDs appear in the cursor.
// Pages already appear via their items; databases need explicit entries.
for _, rid := range resourceIDs {
if _, ok := newEditTimes[rid]; !ok {
newEditTimes[rid] = time.Now()
}
}
return items, buildCursor(newEditTimes), nil
}
// Subsequent syncs: discover all pages, diff against cursor, fetch only changes
logger.Infof(ctx, "[Notion] incremental sync, discovering pages")
pages, fetchVisited := c.discoverAllResources(ctx, client, resourceIDs)
logger.Infof(ctx, "[Notion] discovered %d pages", len(pages))
newEditTimes := make(map[string]time.Time)
pageByID := make(map[string]*notionPage, len(pages))
for i := range pages {
newEditTimes[pages[i].ID] = pages[i].LastEditedTime
pageByID[pages[i].ID] = &pages[i]
}
var changedItems []types.FetchedItem
changedCount := 0
for pageID, newTime := range newEditTimes {
prevTime, existed := prevCursor.PageEditTimes[pageID]
if existed && newTime.Equal(prevTime) {
continue
}
if fetchVisited[pageID] {
continue
}
changedCount++
pg := pageByID[pageID]
if pg == nil {
continue
}
logger.Debugf(ctx, "[Notion] changed: %s (%s, %s)", pg.Title, pg.ID, pg.Object)
if pg.isDatabase() {
// Incremental database sync: query records, diff against cursor,
// only fetch blocks for records whose edit time actually changed.
items, recordEditTimes := c.fetchDatabaseIncremental(ctx, client, pg.ID, prevCursor.PageEditTimes, fetchVisited)
changedItems = append(changedItems, items...)
// Merge record-level edit times into the cursor
for rid, rt := range recordEditTimes {
newEditTimes[rid] = rt
}
} else {
changedItems = append(changedItems, c.fetchPage(ctx, client, pg, fetchVisited)...)
}
}
// Detect deletions. Only emit IsDeleted for pages that previously belonged
// to a selected resource and have actually disappeared at source. Pages that
// are still visible but no longer reachable from any selected root (i.e. the
// user deselected their ancestor) live in fetchVisited as the "excluded" set
// returned by discoverAllResources, and must NOT be reported as deletions.
for pageID := range prevCursor.PageEditTimes {
if _, exists := newEditTimes[pageID]; exists {
continue
}
if fetchVisited[pageID] {
continue
}
changedItems = append(changedItems, types.FetchedItem{
ExternalID: pageID,
IsDeleted: true,
Metadata: map[string]string{"channel": types.ChannelNotion},
})
}
logger.Infof(ctx, "[Notion] incremental: %d changed, %d total items", changedCount, len(changedItems))
return changedItems, buildCursor(newEditTimes), nil
}
func buildCursor(editTimes map[string]time.Time) *types.SyncCursor {
now := time.Now()
cursorData := notionCursor{
PageEditTimes: editTimes,
}
cursorBytes, _ := json.Marshal(cursorData)
var cursorMap map[string]interface{}
_ = json.Unmarshal(cursorBytes, &cursorMap)
return &types.SyncCursor{
LastSyncTime: now,
ConnectorCursor: cursorMap,
}
}
// --- Internal methods ---
// fetchPage fetches a single page's content and attachments.
// If page is nil, it will be fetched from the API.
func (c *Connector) fetchPage(ctx context.Context, client *notionClient, page *notionPage, visited map[string]bool) []types.FetchedItem {
if page == nil {
return nil
}
if visited[page.ID] {
return nil
}
visited[page.ID] = true
if page.InTrash {
return nil
}
// Database records store content in properties, not blocks — delegate to
// buildRecordItem so they aren't dropped as empty pages. Record parent.type
// can be either "database_id" (older GetPage responses) or "data_source_id"
// (search API and 2025-09-03+ workspaces); accept both.
if page.Parent.Type == parentTypeDatabaseID || page.Parent.Type == parentTypeDataSourceID {
dbTitle := ""
if dbInfo, err := getDatabaseOrDataSourceInfo(ctx, client, page.Parent.GetParentID()); err != nil {
logger.Debugf(ctx, "[Notion] failed to resolve parent db title for record %s: %v", page.ID, err)
} else {
dbTitle = dbInfo.Page.Title
}
propNames := extractPropertySchema(*page)
if item := c.buildRecordItem(ctx, client, *page, propNames, dbTitle); item != nil {
return []types.FetchedItem{*item}
}
return nil
}
blocks, err := client.GetBlockChildrenAll(ctx, page.ID)
if err != nil {
logger.Warnf(ctx, "[Notion] failed to get blocks for page %s: %v", page.ID, err)
return nil
}
resolveFileUploads(ctx, client, blocks)
markdown, attachmentList := BlocksToMarkdown(blocks)
var items []types.FetchedItem
// Only skip truly empty pages (no content at all)
if strings.TrimSpace(markdown) != "" {
fileName := page.Title + ".md"
if page.Title == "" {
fileName = defaultUntitledName + ".md"
}
items = append(items, types.FetchedItem{
ExternalID: page.ID,
Title: page.Title,
Content: []byte(markdown),
ContentType: contentTypeMarkdown,
FileName: fileName,
URL: page.URL,
UpdatedAt: page.LastEditedTime,
Metadata: map[string]string{
"channel": types.ChannelNotion,
"object_type": objectTypePage,
},
})
}
// Download file attachments (PDF, documents, etc.)
// Skip images — they're already referenced in Markdown as ![](url),
// and downloading them as separate items requires VLM to process.
for _, att := range attachmentList {
if att.URL == "" || att.Type == "image" {
continue
}
data, err := client.DownloadFile(ctx, att.URL)
if err != nil {
logger.Warnf(ctx, "[Notion] failed to download attachment %s: %v", att.FileName, err)
continue
}
items = append(items, types.FetchedItem{
ExternalID: fmt.Sprintf("%s:%s", page.ID, att.FileName),
Title: att.FileName,
Content: data,
ContentType: mimeTypeForAttachment(att.Type),
FileName: att.FileName,
SourceResourceID: page.ID,
Metadata: map[string]string{
"channel": types.ChannelNotion,
"object_type": objectTypeAttachment,
},
})
}
for _, block := range blocks {
switch block.Type {
case "child_page":
childPage, err := client.GetPage(ctx, block.ID)
if err != nil {
logger.Warnf(ctx, "[Notion] failed to get child page %s: %v", block.ID, err)
continue
}
items = append(items, c.fetchPage(ctx, client, childPage, visited)...)
case "child_database":
items = append(items, c.fetchDatabase(ctx, client, block.ID, visited)...)
}
}
return items
}
// fetchDatabase syncs each database record as an individual knowledge item (full sync).
// Accepts either a data_source_id (from search) or database_id (from child_database blocks).
func (c *Connector) fetchDatabase(ctx context.Context, client *notionClient, id string, visited map[string]bool) []types.FetchedItem {
if visited[id] {
return nil
}
visited[id] = true
records, dbTitle, queryID, err := c.queryDatabaseRecords(ctx, client, id)
if err != nil || len(records) == 0 {
return nil
}
if queryID != "" && queryID != id {
if visited[queryID] {
return nil
}
visited[queryID] = true
}
for _, record := range records {
visited[record.ID] = true // Mark records as visited to avoid duplicate fetchPage calls
}
item := c.buildDatabaseItem(ctx, client, id, dbTitle, records)
if item != nil {
return []types.FetchedItem{*item}
}
return nil
}
// fetchDatabaseIncremental syncs only changed records by comparing edit times against cursor.
// Returns fetched items and a map of record_id → edit_time for cursor update.
func (c *Connector) fetchDatabaseIncremental(ctx context.Context, client *notionClient, id string, prevEditTimes map[string]time.Time, visited map[string]bool) ([]types.FetchedItem, map[string]time.Time) {
if visited[id] {
return nil, nil
}
visited[id] = true
records, dbTitle, queryID, err := c.queryDatabaseRecords(ctx, client, id)
if err != nil {
return nil, nil
}
if queryID != "" && queryID != id {
if visited[queryID] {
return nil, nil
}
visited[queryID] = true
}
recordEditTimes := make(map[string]time.Time, len(records))
changedCount := 0
for _, record := range records {
visited[record.ID] = true // Mark records as visited to avoid duplicate fetchPage calls
if record.InTrash {
continue
}
recordEditTimes[record.ID] = record.LastEditedTime
prevTime, existed := prevEditTimes[record.ID]
if !existed || !record.LastEditedTime.Equal(prevTime) {
changedCount++
}
}
logger.Infof(ctx, "[Notion] database %s incremental: %d changed out of %d records", id, changedCount, len(records))
// Rebuild the entire database table if any record changed, or if we don't have previous times (first sync)
if changedCount > 0 || len(prevEditTimes) == 0 {
item := c.buildDatabaseItem(ctx, client, id, dbTitle, records)
if item != nil {
return []types.FetchedItem{*item}, recordEditTimes
}
}
return nil, recordEditTimes
}
// queryDatabaseRecords resolves the database ID and queries all records.
// Returns the records, the database title, and the canonical data_source_id used for the query.
func (c *Connector) queryDatabaseRecords(ctx context.Context, client *notionClient, id string) ([]notionPage, string, string, error) {
dbInfo, err := getDatabaseOrDataSourceInfo(ctx, client, id)
if err != nil {
logger.Warnf(ctx, "[Notion] failed to get database/data_source info %s: %v", id, err)
return nil, "", "", err
}
queryID := dbInfo.DataSourceID
if queryID == "" {
queryID = id
}
records, err := client.QueryDatabaseAll(ctx, queryID)
if err != nil {
logger.Warnf(ctx, "[Notion] failed to query database %s: %v", id, err)
return nil, "", "", err
}
logger.Infof(ctx, "[Notion] database %s (%s): %d records", id, dbInfo.Page.Title, len(records))
return records, dbInfo.Page.Title, queryID, nil
}
// buildRecordItem converts a single database record into a FetchedItem
// with properties as header and block content as body.
func (c *Connector) buildRecordItem(ctx context.Context, client *notionClient, record notionPage, propNames []string, dbTitle string) *types.FetchedItem {
var content strings.Builder
title := record.Title
if title == "" {
title = defaultUntitledName
}
content.WriteString("# " + title + "\n\n")
if record.RawProperties != nil {
var props map[string]interface{}
json.Unmarshal(record.RawProperties, &props)
for _, name := range propNames {
if propMap, ok := props[name].(map[string]interface{}); ok {
val := propertyToString(propMap)
if val != "" {
content.WriteString("- **" + name + "**: " + val + "\n")
}
}
}
}
// Database records may contain page-style block content beyond their properties
blocks, err := client.GetBlockChildrenAll(ctx, record.ID)
if err != nil {
logger.Warnf(ctx, "[Notion] failed to get blocks for record %s: %v", record.ID, err)
} else if len(blocks) > 0 {
resolveFileUploads(ctx, client, blocks)
markdown, _ := BlocksToMarkdown(blocks)
if strings.TrimSpace(markdown) != "" {
content.WriteString("\n" + markdown)
}
}
bodyStr := strings.TrimSpace(content.String())
if bodyStr == "" {
return nil
}
return &types.FetchedItem{
ExternalID: record.ID,
Title: title,
Content: []byte(bodyStr),
ContentType: contentTypeMarkdown,
FileName: title + ".md",
URL: record.URL,
UpdatedAt: record.LastEditedTime,
Metadata: map[string]string{
"channel": types.ChannelNotion,
"object_type": objectTypePage,
"database": dbTitle,
},
}
}
// buildDatabaseItem converts a database into a single Markdown table document.
func (c *Connector) buildDatabaseItem(ctx context.Context, client *notionClient, id string, dbTitle string, records []notionPage) *types.FetchedItem {
if len(records) == 0 {
return nil
}
propNames := extractPropertySchema(records[0])
var content strings.Builder
title := dbTitle
if title == "" {
title = defaultUntitledName
}
content.WriteString("# " + title + "\n\n")
// Table header
content.WriteString("| Title ")
for _, name := range propNames {
content.WriteString("| " + strings.ReplaceAll(name, "|", "\\|") + " ")
}
content.WriteString("|\n|")
content.WriteString("---|")
for range propNames {
content.WriteString("---|")
}
content.WriteString("\n")
var extraContent strings.Builder
// Table rows
for _, record := range records {
if record.InTrash {
continue
}
recordTitle := record.Title
if recordTitle == "" {
recordTitle = defaultUntitledName
}
content.WriteString("| " + strings.ReplaceAll(recordTitle, "|", "\\|") + " ")
if record.RawProperties != nil {
var props map[string]interface{}
json.Unmarshal(record.RawProperties, &props)
for _, name := range propNames {
val := ""
if propMap, ok := props[name].(map[string]interface{}); ok {
val = propertyToString(propMap)
}
val = strings.ReplaceAll(val, "\n", "<br>")
val = strings.ReplaceAll(val, "|", "\\|")
content.WriteString("| " + val + " ")
}
}
content.WriteString("|\n")
// Fetch blocks for the record, in case it has page content
blocks, err := client.GetBlockChildrenAll(ctx, record.ID)
if err == nil && len(blocks) > 0 {
resolveFileUploads(ctx, client, blocks)
markdown, _ := BlocksToMarkdown(blocks)
if strings.TrimSpace(markdown) != "" {
extraContent.WriteString("\n## " + recordTitle + " 内容\n\n" + markdown + "\n")
}
}
}
if extraContent.Len() > 0 {
content.WriteString("\n" + extraContent.String())
}
bodyStr := strings.TrimSpace(content.String())
if bodyStr == "" {
return nil
}
updatedAt := time.Now()
if len(records) > 0 {
updatedAt = records[0].LastEditedTime
}
return &types.FetchedItem{
ExternalID: id,
Title: title,
Content: []byte(bodyStr),
ContentType: contentTypeMarkdown,
FileName: title + ".md",
URL: fmt.Sprintf("https://notion.so/%s", strings.ReplaceAll(id, "-", "")),
UpdatedAt: updatedAt,
Metadata: map[string]string{
"channel": types.ChannelNotion,
"object_type": objectTypeDatabase,
},
}
}
// discoverAllResources uses the Search API to find all pages and data_sources,
// then filters by parent chain via BFS from the selected resource IDs.
// Returns the included pages plus the excluded set (visible pages NOT under any
// selected root) so callers can seed `visited` for child-block recursion without
// a second SearchPages round-trip.
func (c *Connector) discoverAllResources(ctx context.Context, client *notionClient, resourceIDs []string) (included []notionPage, excluded map[string]bool) {
allPages, err := client.SearchPages(ctx)
if err != nil {
logger.Warnf(ctx, "[Notion] failed to search pages for discovery: %v", err)
return nil, nil
}
// data_source objects use database_parent for workspace hierarchy.
allIDs := make(map[string]bool, len(allPages))
pageByID := make(map[string]notionPage, len(allPages))
for _, p := range allPages {
allIDs[p.ID] = true
pageByID[p.ID] = p
}
childrenOf := make(map[string][]string)
for _, p := range allPages {
if p.InTrash {
continue
}
parentID := resolveParentID(p, allIDs)
if parentID != "" {
childrenOf[parentID] = append(childrenOf[parentID], p.ID)
}
}
// BFS from each resource root to collect all descendants
includedSet := make(map[string]bool)
queue := make([]string, 0, len(resourceIDs))
for _, id := range resourceIDs {
if _, ok := pageByID[id]; ok {
includedSet[id] = true
queue = append(queue, id)
}
}
for len(queue) > 0 {
current := queue[0]
queue = queue[1:]
for _, childID := range childrenOf[current] {
if !includedSet[childID] {
includedSet[childID] = true
queue = append(queue, childID)
}
}
}
included = make([]notionPage, 0, len(includedSet))
for id := range includedSet {
included = append(included, pageByID[id])
}
excluded = make(map[string]bool, len(allPages)-len(includedSet))
for id := range allIDs {
if !includedSet[id] {
excluded[id] = true
}
}
return included, excluded
}
// --- File upload resolution ---
// resolveFileUploads walks the block tree and resolves file_upload type files
// by re-fetching the block to get a temporary download URL.
// This mutates blocks in place, replacing file_upload RawContent with resolved URLs.
func resolveFileUploads(ctx context.Context, client *notionClient, blocks []notionBlock) {
for i := range blocks {
if isFileBlock(blocks[i].Type) && blocks[i].RawContent != nil {
var file notionFile
json.Unmarshal(blocks[i].RawContent, &file)
if file.GetFileUploadID() != "" {
resolved, err := client.ResolveBlock(ctx, blocks[i].ID)
if err != nil {
logger.Warnf(ctx, "[Notion] failed to resolve file_upload in block %s: %v", blocks[i].ID, err)
continue
}
blocks[i].RawContent = resolved.RawContent
}
}
if len(blocks[i].Children) > 0 {
resolveFileUploads(ctx, client, blocks[i].Children)
}
}
}
func isFileBlock(blockType string) bool {
switch blockType {
case "image", "file", "pdf", "video", "audio":
return true
}
return false
}
// --- Property extraction ---
// propertyToString recursively extracts a string value from a Notion property.
// Follows the type chain generically, handling all 22 property types without hardcoding.
func propertyToString(value map[string]interface{}) string {
if value == nil {
return ""
}
typeName, _ := value["type"].(string)
if typeName == "" {
return extractLeafValue(value)
}
inner, exists := value[typeName]
if !exists || inner == nil {
return ""
}
return extractValue(inner)
}
func extractValue(v interface{}) string {
switch val := v.(type) {
case string:
return val
case float64:
if val == float64(int64(val)) {
return fmt.Sprintf("%d", int64(val))
}
return fmt.Sprintf("%g", val)
case bool:
return fmt.Sprintf("%t", val)
case map[string]interface{}:
return extractLeafValue(val)
case []interface{}:
var parts []string
for _, item := range val {
s := extractValue(item)
if s != "" {
parts = append(parts, s)
}
}
return strings.Join(parts, ", ")
case nil:
return ""
default:
return fmt.Sprint(v)
}
}
func extractLeafValue(m map[string]interface{}) string {
if name, ok := m["name"].(string); ok {
return name
}
if content, ok := m["content"].(string); ok {
return content
}
if pt, ok := m["plain_text"].(string); ok {
return pt
}
if start, ok := m["start"].(string); ok {
if end, ok := m["end"].(string); ok && end != "" {
return start + " ~ " + end
}
return start
}
if expr, ok := m["expression"].(string); ok {
return expr
}
// Recurse if there's a type field
if typeName, ok := m["type"].(string); ok {
if inner, ok := m[typeName]; ok {
return extractValue(inner)
}
}
return ""
}
// extractPropertySchema extracts non-title property names from a record.
// Names are sorted alphabetically to ensure deterministic order
// (prevents false-positive content changes on incremental sync due to map iteration randomness).
func extractPropertySchema(record notionPage) (propNames []string) {
if record.RawProperties == nil {
return nil
}
var props map[string]json.RawMessage
if err := json.Unmarshal(record.RawProperties, &props); err != nil {
return nil
}
for name, propRaw := range props {
var prop struct {
Type string `json:"type"`
}
json.Unmarshal(propRaw, &prop)
if prop.Type != "title" {
propNames = append(propNames, name)
}
}
sort.Strings(propNames)
return
}
// resolveParentID determines where a page/data_source sits in the workspace hierarchy.
// For data_source objects, uses database_parent which points directly at the
// containing page; the regular `parent` field on a data_source points at its
// database container, not at the workspace location, so it's the wrong source.
func resolveParentID(p notionPage, allIDs map[string]bool) string {
// Data sources: use database_parent which shows the workspace location
// (e.g., which page contains this database), not the database container ID.
if p.isDatabase() && p.DatabaseParent != nil {
pid := p.DatabaseParent.GetParentID()
if pid != "" && allIDs[pid] {
return pid
}
return ""
}
// Pages: use direct parent
if p.Parent.Type == parentTypeWorkspace {
return ""
}
pid := p.Parent.GetParentID()
if pid != "" && allIDs[pid] {
return pid
}
return ""
}
// getDatabaseOrDataSourceInfo retrieves metadata for a database or data_source.
// Tries GET /v1/data_sources/{id} first (search returns data_source IDs),
// falls back to GET /v1/databases/{id} (child_database blocks use database IDs).
func getDatabaseOrDataSourceInfo(ctx context.Context, client *notionClient, id string) (*databaseInfo, error) {
ds, err := client.GetDataSourceInfo(ctx, id)
if err == nil {
return &databaseInfo{Page: *ds, DataSourceID: id}, nil
}
return client.GetDatabaseInfo(ctx, id)
}
// computeExcludedSet returns IDs the user has explicitly deselected: visible in
// the picker but neither selected nor a descendant of a selected node. Used to
// seed `visited` so recursive child_page/child_database traversal skips them.
// Pages the user has never seen (e.g. added after the last config save) are
// not excluded, so a selected parent still picks them up automatically.
func computeExcludedSet(visibleIDs []string, parentOf map[string]string, selectedIDs []string) map[string]bool {
selected := make(map[string]bool, len(selectedIDs))
for _, id := range selectedIDs {
selected[id] = true
}
excluded := make(map[string]bool, len(visibleIDs))
for _, id := range visibleIDs {
hasSelectedAncestor := false
for cur := id; cur != ""; cur = parentOf[cur] {
if selected[cur] {
hasSelectedAncestor = true
break
}
}
if !hasSelectedAncestor {
excluded[id] = true
}
}
return excluded
}
// excludedSetFromListResources fetches the picker hierarchy via ListResources
// and computes the deselected set. Used by FetchAll where no other code path
// already has the page list in hand.
func (c *Connector) excludedSetFromListResources(ctx context.Context, config *types.DataSourceConfig, selectedIDs []string) map[string]bool {
visible, err := c.ListResources(ctx, config, "")
if err != nil {
logger.Warnf(ctx, "[Notion] failed to list visible resources for exclusion: %v", err)
return map[string]bool{}
}
ids := make([]string, 0, len(visible))
parentOf := make(map[string]string, len(visible))
for _, r := range visible {
ids = append(ids, r.ExternalID)
parentOf[r.ExternalID] = r.ParentID
}
return computeExcludedSet(ids, parentOf, selectedIDs)
}
// --- Helpers ---
func extractBaseURL(config *types.DataSourceConfig) string {
if config.Settings != nil {
if url, ok := config.Settings["base_url"].(string); ok && url != "" {
return url
}
}
return DefaultBaseURL
}
func mimeTypeForAttachment(attType string) string {
switch attType {
case "image":
return "image/png"
case "pdf":
return "application/pdf"
case "video":
return "video/mp4"
case "audio":
return "audio/mpeg"
default:
return "application/octet-stream"
}
}
@@ -0,0 +1,301 @@
package notion
import (
"context"
"strings"
"testing"
"github.com/Tencent/WeKnora/internal/types"
secutils "github.com/Tencent/WeKnora/internal/utils"
)
func allowNotionTestServer(t *testing.T) {
t.Helper()
t.Setenv("SSRF_WHITELIST", "127.0.0.1,localhost")
secutils.ResetSSRFWhitelistForTest()
t.Cleanup(secutils.ResetSSRFWhitelistForTest)
}
func makeNotionConfig(cfg *Config, baseURL string, resourceIDs []string) *types.DataSourceConfig {
return &types.DataSourceConfig{
Type: types.ConnectorTypeNotion,
Credentials: map[string]interface{}{
"api_key": cfg.APIKey,
},
ResourceIDs: resourceIDs,
Settings: map[string]interface{}{
"base_url": baseURL,
},
}
}
func TestConnectorType(t *testing.T) {
c := NewConnector()
if c.Type() != types.ConnectorTypeNotion {
t.Errorf("Type() = %q, want %q", c.Type(), types.ConnectorTypeNotion)
}
}
func TestConnectorValidate(t *testing.T) {
allowNotionTestServer(t)
ts, cfg := fakeNotion()
defer ts.Close()
c := NewConnector()
err := c.Validate(context.Background(), makeNotionConfig(cfg, ts.URL, nil))
if err != nil {
t.Fatalf("Validate() error: %v", err)
}
}
func TestConnectorValidate_BadToken(t *testing.T) {
ts, _ := fakeNotion()
defer ts.Close()
c := NewConnector()
err := c.Validate(context.Background(), makeNotionConfig(
&Config{APIKey: "bad-token"}, ts.URL, nil,
))
if err == nil {
t.Fatal("expected error for bad token")
}
}
func TestConnectorListResources(t *testing.T) {
allowNotionTestServer(t)
ts, cfg := fakeNotion()
defer ts.Close()
c := NewConnector()
resources, err := c.ListResources(context.Background(), makeNotionConfig(cfg, ts.URL, nil), "")
if err != nil {
t.Fatalf("ListResources() error: %v", err)
}
if len(resources) != 2 {
t.Fatalf("expected 2 resources, got %d", len(resources))
}
if resources[0].ExternalID != "page-1" || resources[0].Type != "page" || resources[0].Name != "Test Page" {
t.Errorf("resource[0] = %+v", resources[0])
}
if resources[1].ExternalID != "db-1" || resources[1].Type != "database" || resources[1].Name != "Test Database" {
t.Errorf("resource[1] = %+v, want Name=%q", resources[1], "Test Database")
}
}
func TestConnectorFetchAll(t *testing.T) {
allowNotionTestServer(t)
ts, cfg := fakeNotion()
defer ts.Close()
c := NewConnector()
items, err := c.FetchAll(context.Background(), makeNotionConfig(cfg, ts.URL, []string{"page-1"}), []string{"page-1"})
if err != nil {
t.Fatalf("FetchAll() error: %v", err)
}
if len(items) == 0 {
t.Fatal("expected at least 1 item")
}
// First item should be the page with markdown content
found := false
for _, item := range items {
if item.ExternalID == "page-1" {
found = true
if item.ContentType != "text/markdown" {
t.Errorf("ContentType = %q", item.ContentType)
}
if item.Metadata["channel"] != "notion" {
t.Errorf("channel = %q", item.Metadata["channel"])
}
if len(item.Content) == 0 {
t.Error("expected non-empty content")
}
}
}
if !found {
t.Error("page-1 not found in items")
}
}
func TestConnectorFetchAll_Database(t *testing.T) {
allowNotionTestServer(t)
ts, cfg := fakeNotion()
defer ts.Close()
c := NewConnector()
items, err := c.FetchAll(context.Background(), makeNotionConfig(cfg, ts.URL, []string{"db-1"}), []string{"db-1"})
if err != nil {
t.Fatalf("FetchAll() error: %v", err)
}
if len(items) == 0 {
t.Fatal("expected at least 1 item")
}
// The entire database should be synced as a single table knowledge item
found := false
for _, item := range items {
if item.ExternalID == "db-1" {
found = true
if item.Metadata["object_type"] != "database" {
t.Errorf("object_type = %q, want %q", item.Metadata["object_type"], "database")
}
if item.ContentType != "text/markdown" {
t.Errorf("ContentType = %q", item.ContentType)
}
if len(item.Content) == 0 {
t.Error("expected non-empty content")
}
// Verify it contains table headers
contentStr := string(item.Content)
if !strings.Contains(contentStr, "| Title |") {
t.Errorf("expected table header in content, got: %s", contentStr)
}
}
}
if !found {
t.Error("db-1 not found in items")
}
}
// TestConnectorFetchAll_SingleRecord verifies that selecting a single database
// row by ID routes through fetchPage's record-detection branch and produces an
// item via buildRecordItem (instead of being silently dropped as an empty page).
func TestConnectorFetchAll_SingleRecord(t *testing.T) {
allowNotionTestServer(t)
ts, cfg := fakeNotion()
defer ts.Close()
c := NewConnector()
items, err := c.FetchAll(context.Background(), makeNotionConfig(cfg, ts.URL, []string{"record-1"}), []string{"record-1"})
if err != nil {
t.Fatalf("FetchAll() error: %v", err)
}
if len(items) != 1 {
t.Fatalf("expected 1 item, got %d", len(items))
}
if items[0].ExternalID != "record-1" {
t.Errorf("ExternalID = %q, want record-1", items[0].ExternalID)
}
if items[0].Metadata["object_type"] != "page" {
t.Errorf("object_type = %q, want page", items[0].Metadata["object_type"])
}
if len(items[0].Content) == 0 {
t.Error("expected non-empty content from buildRecordItem")
}
}
func TestConnectorFetchIncremental_NoChanges(t *testing.T) {
allowNotionTestServer(t)
ts, cfg := fakeNotion()
defer ts.Close()
c := NewConnector()
config := makeNotionConfig(cfg, ts.URL, []string{"page-1"})
// First: full fetch to establish baseline
_, err := c.FetchAll(context.Background(), config, []string{"page-1"})
if err != nil {
t.Fatalf("FetchAll() error: %v", err)
}
// Build a cursor that matches current state
cursor := &types.SyncCursor{
ConnectorCursor: map[string]interface{}{
"page_edit_times": map[string]interface{}{
"page-1": "2026-01-15T10:00:00Z",
},
},
}
items, newCursor, err := c.FetchIncremental(context.Background(), config, cursor)
if err != nil {
t.Fatalf("FetchIncremental() error: %v", err)
}
// No changes expected (timestamps match)
if len(items) != 0 {
t.Errorf("expected 0 items for no changes, got %d", len(items))
}
if newCursor == nil {
t.Fatal("expected non-nil cursor")
}
}
func TestPropertyToString(t *testing.T) {
tests := []struct {
name string
value map[string]interface{}
expected string
}{
{
name: "select",
value: map[string]interface{}{
"type": "select",
"select": map[string]interface{}{"name": "Done"},
},
expected: "Done",
},
{
name: "rich_text",
value: map[string]interface{}{
"type": "rich_text",
"rich_text": []interface{}{
map[string]interface{}{"plain_text": "Hello"},
},
},
expected: "Hello",
},
{
name: "number",
value: map[string]interface{}{
"type": "number",
"number": 42.0,
},
expected: "42",
},
{
name: "checkbox true",
value: map[string]interface{}{
"type": "checkbox",
"checkbox": true,
},
expected: "true",
},
{
name: "date",
value: map[string]interface{}{
"type": "date",
"date": map[string]interface{}{"start": "2026-01-15", "end": "2026-01-20"},
},
expected: "2026-01-15 ~ 2026-01-20",
},
{
name: "multi_select",
value: map[string]interface{}{
"type": "multi_select",
"multi_select": []interface{}{
map[string]interface{}{"name": "Tag1"},
map[string]interface{}{"name": "Tag2"},
},
},
expected: "Tag1, Tag2",
},
{
name: "nil value",
value: map[string]interface{}{"type": "rich_text", "rich_text": nil},
expected: "",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
got := propertyToString(tt.value)
if got != tt.expected {
t.Errorf("propertyToString() = %q, want %q", got, tt.expected)
}
})
}
}
@@ -0,0 +1,497 @@
package notion
import (
"encoding/json"
"fmt"
"strings"
)
// BlocksToMarkdown converts a tree of Notion blocks into Markdown text,
// and collects attachment references (image, file, pdf, video, audio).
func BlocksToMarkdown(blocks []notionBlock) (string, []attachment) {
var b strings.Builder
var attachments []attachment
renderBlocks(&b, blocks, 0, &attachments)
// Collapse 3+ consecutive newlines into 2 (avoids excessive blank lines from empty paragraphs)
result := b.String()
for strings.Contains(result, "\n\n\n") {
result = strings.ReplaceAll(result, "\n\n\n", "\n\n")
}
return strings.TrimSpace(result) + "\n", attachments
}
// renderBlocks renders a list of blocks at the given indentation depth.
func renderBlocks(b *strings.Builder, blocks []notionBlock, depth int, attachments *[]attachment) {
inList := false
for i, block := range blocks {
isList := block.Type == "bulleted_list_item" || block.Type == "numbered_list_item" || block.Type == "to_do"
// End list spacing: if previous was list and current is not
if inList && !isList {
b.WriteString("\n")
inList = false
}
renderBlock(b, block, depth, attachments, i, blocks)
if isList {
inList = true
}
}
// Close trailing list
if inList {
b.WriteString("\n")
}
}
// renderBlock renders a single block to Markdown.
func renderBlock(b *strings.Builder, block notionBlock, depth int, attachments *[]attachment, index int, siblings []notionBlock) {
indent := strings.Repeat(" ", depth)
switch block.Type {
case "paragraph":
rt := extractRichText(block.RawContent)
text := renderRichText(rt)
b.WriteString(indent + text + "\n\n")
case "heading_1", "heading_2", "heading_3", "heading_4":
level := int(block.Type[len(block.Type)-1] - '0')
rt := extractRichText(block.RawContent)
b.WriteString(strings.Repeat("#", level) + " " + renderRichText(rt) + "\n\n")
if len(block.Children) > 0 {
renderBlocks(b, block.Children, depth, attachments)
}
case "bulleted_list_item":
rt := extractRichText(block.RawContent)
b.WriteString(indent + "- " + renderRichText(rt) + "\n")
if len(block.Children) > 0 {
renderBlocks(b, block.Children, depth+1, attachments)
}
case "numbered_list_item":
rt := extractRichText(block.RawContent)
// Find position among consecutive numbered_list_items
num := 1
for j := index - 1; j >= 0 && siblings[j].Type == "numbered_list_item"; j-- {
num++
}
b.WriteString(fmt.Sprintf("%s%d. %s\n", indent, num, renderRichText(rt)))
if len(block.Children) > 0 {
renderBlocks(b, block.Children, depth+1, attachments)
}
case "to_do":
rt := extractRichText(block.RawContent)
checked := extractBool(block.RawContent, "checked")
checkbox := "[ ]"
if checked {
checkbox = "[x]"
}
b.WriteString(fmt.Sprintf("%s- %s %s\n", indent, checkbox, renderRichText(rt)))
case "toggle":
rt := extractRichText(block.RawContent)
b.WriteString(fmt.Sprintf("<details><summary>%s</summary>\n\n", renderRichText(rt)))
if len(block.Children) > 0 {
renderBlocks(b, block.Children, depth, attachments)
}
b.WriteString("</details>\n\n")
case "code":
rt := extractRichText(block.RawContent)
lang := extractString(block.RawContent, "language")
b.WriteString(fmt.Sprintf("```%s\n%s\n```\n\n", lang, renderRichText(rt)))
case "quote", "meeting_notes":
rt := extractRichText(block.RawContent)
text := renderRichText(rt)
for _, line := range strings.Split(text, "\n") {
b.WriteString("> " + line + "\n")
}
b.WriteString("\n")
if len(block.Children) > 0 {
renderBlocks(b, block.Children, depth, attachments)
}
case "callout":
rt := extractRichText(block.RawContent)
icon := extractIcon(block.RawContent)
prefix := ""
if icon != "" {
prefix = icon + " "
}
b.WriteString("> " + prefix + renderRichText(rt) + "\n\n")
if len(block.Children) > 0 {
renderBlocks(b, block.Children, depth, attachments)
}
case "divider":
b.WriteString("---\n\n")
case "equation":
expr := extractString(block.RawContent, "expression")
b.WriteString("$$" + expr + "$$\n\n")
case "table":
renderTable(b, block)
case "image":
file, caption := extractFileAndCaption(block.RawContent)
url := file.GetURL()
b.WriteString(fmt.Sprintf("![%s](%s)\n\n", caption, url))
if url != "" {
*attachments = append(*attachments, attachment{
URL: url,
FileName: fileNameFromURL(url, "image"),
Type: "image",
})
}
case "file", "pdf", "video", "audio":
renderMediaBlock(b, block, attachments)
case "bookmark", "link_preview":
url := extractString(block.RawContent, "url")
caption := ""
if rt := extractCaptionText(block.RawContent); rt != "" {
caption = rt
} else {
caption = url
}
b.WriteString(fmt.Sprintf("[%s](%s)\n\n", caption, url))
case "embed":
url := extractString(block.RawContent, "url")
b.WriteString(fmt.Sprintf("[%s](%s)\n\n", url, url))
case "link_to_page":
// Render as a link; the actual page title is not available in block data
pageID := extractLinkToPageID(block.RawContent)
b.WriteString(fmt.Sprintf("[Page %s](https://notion.so/%s)\n\n", pageID, strings.ReplaceAll(pageID, "-", "")))
case "synced_block":
// Render children (the synced content)
if len(block.Children) > 0 {
renderBlocks(b, block.Children, depth, attachments)
}
case "column_list":
// Render each column sequentially
if len(block.Children) > 0 {
for _, col := range block.Children {
if len(col.Children) > 0 {
renderBlocks(b, col.Children, depth, attachments)
}
}
}
case "column":
// Handled by column_list
if len(block.Children) > 0 {
renderBlocks(b, block.Children, depth, attachments)
}
case "tab_list":
// Render each tab's content sequentially (similar to column_list)
if len(block.Children) > 0 {
for _, tab := range block.Children {
if len(tab.Children) > 0 {
renderBlocks(b, tab.Children, depth, attachments)
}
}
}
case "tab":
// Render tab content sequentially
if len(block.Children) > 0 {
renderBlocks(b, block.Children, depth, attachments)
}
case "child_page":
// Render as a link in the parent page so it's not empty.
// The child page itself is fetched as a separate knowledge item by connector layer.
title := extractString(block.RawContent, "title")
if title == "" {
title = "Untitled"
}
b.WriteString(fmt.Sprintf("- [%s](https://notion.so/%s)\n", title, strings.ReplaceAll(block.ID, "-", "")))
case "child_database":
title := extractString(block.RawContent, "title")
if title == "" {
title = "Database"
}
b.WriteString(fmt.Sprintf("- [%s](https://notion.so/%s)\n", title, strings.ReplaceAll(block.ID, "-", "")))
case "table_of_contents", "breadcrumb", "template", "table_row", "unsupported":
// Skip — no meaningful content
default:
// Unknown block type — skip silently for forward compatibility
}
}
// renderMediaBlock renders a file/pdf/video/audio block as a Markdown link and collects the attachment.
func renderMediaBlock(b *strings.Builder, block notionBlock, attachments *[]attachment) {
file, _ := extractFileAndCaption(block.RawContent)
url := file.GetURL()
name := file.Name
if name == "" {
name = fileNameFromURL(url, block.Type)
}
b.WriteString(fmt.Sprintf("[%s](%s)\n\n", name, url))
if url != "" {
*attachments = append(*attachments, attachment{
URL: url,
FileName: name,
Type: block.Type,
})
}
}
// renderTable converts a table block (with table_row children) to Markdown.
func renderTable(b *strings.Builder, block notionBlock) {
if len(block.Children) == 0 {
return
}
for i, row := range block.Children {
cells := extractTableCells(row.RawContent)
var parts []string
for _, cell := range cells {
parts = append(parts, strings.ReplaceAll(renderRichText(cell), "|", "\\|"))
}
b.WriteString("| " + strings.Join(parts, " | ") + " |\n")
// Markdown tables always need a separator after the first row
if i == 0 {
var sep []string
for range parts {
sep = append(sep, "---")
}
b.WriteString("| " + strings.Join(sep, " | ") + " |\n")
}
}
b.WriteString("\n")
}
// --- Rich text rendering ---
// renderRichText converts a slice of rich text elements to Markdown.
func renderRichText(texts []notionRichText) string {
var b strings.Builder
for _, rt := range texts {
text := richTextToString(rt)
text = applyAnnotations(text, rt.Annotations)
if rt.Href != "" && !rt.Annotations.Code {
text = fmt.Sprintf("[%s](%s)", text, rt.Href)
}
b.WriteString(text)
}
return b.String()
}
// richTextToString extracts the base text from a rich text element.
func richTextToString(rt notionRichText) string {
switch rt.Type {
case "text":
if rt.Text != nil {
return rt.Text.Content
}
return rt.PlainText
case "mention":
if rt.Mention == nil {
return rt.PlainText
}
switch rt.Mention.Type {
case "date":
if rt.Mention.Date != nil {
if rt.Mention.Date.End != "" {
return rt.Mention.Date.Start + " \u2192 " + rt.Mention.Date.End
}
return rt.Mention.Date.Start
}
case "page":
if rt.Mention.Page != nil {
return rt.PlainText // Notion fills PlainText with page title
}
case "database", "data_source":
if rt.Mention.Database != nil {
return rt.PlainText
}
case "link_preview":
if rt.Mention.LinkPreview != nil {
return rt.Mention.LinkPreview.URL
}
}
return rt.PlainText
case "equation":
if rt.Equation != nil {
return "$" + rt.Equation.Expression + "$"
}
return rt.PlainText
default:
return rt.PlainText
}
}
// applyAnnotations wraps text with Markdown formatting based on annotations.
func applyAnnotations(text string, ann notionAnnotations) string {
if text == "" {
return text
}
if ann.Code {
text = "`" + text + "`"
}
if ann.Bold && ann.Italic {
text = "***" + text + "***"
} else if ann.Bold {
text = "**" + text + "**"
} else if ann.Italic {
text = "*" + text + "*"
}
if ann.Strikethrough {
text = "~~" + text + "~~"
}
if ann.Underline {
text = "<u>" + text + "</u>"
}
return text
}
// --- Raw content extraction helpers ---
func extractRichText(raw json.RawMessage) []notionRichText {
if raw == nil {
return nil
}
var content struct {
RichText []notionRichText `json:"rich_text"`
}
json.Unmarshal(raw, &content)
return content.RichText
}
func extractString(raw json.RawMessage, key string) string {
if raw == nil {
return ""
}
var m map[string]json.RawMessage
if err := json.Unmarshal(raw, &m); err != nil {
return ""
}
val, ok := m[key]
if !ok {
return ""
}
var s string
json.Unmarshal(val, &s)
return s
}
func extractBool(raw json.RawMessage, key string) bool {
if raw == nil {
return false
}
var m map[string]json.RawMessage
if err := json.Unmarshal(raw, &m); err != nil {
return false
}
val, ok := m[key]
if !ok {
return false
}
var b bool
json.Unmarshal(val, &b)
return b
}
func extractIcon(raw json.RawMessage) string {
if raw == nil {
return ""
}
var content struct {
Icon *struct {
Type string `json:"type"`
Emoji string `json:"emoji,omitempty"`
} `json:"icon"`
}
json.Unmarshal(raw, &content)
if content.Icon != nil && content.Icon.Type == "emoji" {
return content.Icon.Emoji
}
return ""
}
func extractFileAndCaption(raw json.RawMessage) (notionFile, string) {
var file notionFile
if raw != nil {
json.Unmarshal(raw, &file)
}
captionText := ""
if len(file.Caption) > 0 {
captionText = renderRichText(file.Caption)
}
return file, captionText
}
func extractCaptionText(raw json.RawMessage) string {
if raw == nil {
return ""
}
var content struct {
Caption []notionRichText `json:"caption"`
}
json.Unmarshal(raw, &content)
if len(content.Caption) > 0 {
return renderRichText(content.Caption)
}
return ""
}
func extractTableCells(raw json.RawMessage) [][]notionRichText {
if raw == nil {
return nil
}
var content struct {
Cells [][]notionRichText `json:"cells"`
}
json.Unmarshal(raw, &content)
return content.Cells
}
func extractLinkToPageID(raw json.RawMessage) string {
if raw == nil {
return ""
}
var content struct {
Type string `json:"type"`
PageID string `json:"page_id,omitempty"`
}
json.Unmarshal(raw, &content)
return content.PageID
}
func fileNameFromURL(url, fallbackType string) string {
if url == "" {
return fallbackType
}
// Extract filename from URL path (before query params)
path := url
if idx := strings.Index(path, "?"); idx >= 0 {
path = path[:idx]
}
if idx := strings.LastIndex(path, "/"); idx >= 0 {
name := path[idx+1:]
if name != "" {
return name
}
}
return fallbackType
}
@@ -0,0 +1,352 @@
package notion
import (
"encoding/json"
"testing"
)
func TestRenderRichText(t *testing.T) {
tests := []struct {
name string
texts []notionRichText
expected string
}{
{
name: "plain text",
texts: []notionRichText{
{Type: "text", PlainText: "hello", Text: &notionTextContent{Content: "hello"}, Annotations: notionAnnotations{}},
},
expected: "hello",
},
{
name: "bold text",
texts: []notionRichText{
{Type: "text", PlainText: "bold", Text: &notionTextContent{Content: "bold"}, Annotations: notionAnnotations{Bold: true}},
},
expected: "**bold**",
},
{
name: "italic text",
texts: []notionRichText{
{Type: "text", PlainText: "italic", Text: &notionTextContent{Content: "italic"}, Annotations: notionAnnotations{Italic: true}},
},
expected: "*italic*",
},
{
name: "code text",
texts: []notionRichText{
{Type: "text", PlainText: "code", Text: &notionTextContent{Content: "code"}, Annotations: notionAnnotations{Code: true}},
},
expected: "`code`",
},
{
name: "strikethrough",
texts: []notionRichText{
{Type: "text", PlainText: "del", Text: &notionTextContent{Content: "del"}, Annotations: notionAnnotations{Strikethrough: true}},
},
expected: "~~del~~",
},
{
name: "underline",
texts: []notionRichText{
{Type: "text", PlainText: "ul", Text: &notionTextContent{Content: "ul"}, Annotations: notionAnnotations{Underline: true}},
},
expected: "<u>ul</u>",
},
{
name: "bold italic combo",
texts: []notionRichText{
{Type: "text", PlainText: "bi", Text: &notionTextContent{Content: "bi"}, Annotations: notionAnnotations{Bold: true, Italic: true}},
},
expected: "***bi***",
},
{
name: "text with link",
texts: []notionRichText{
{Type: "text", PlainText: "click", Href: "https://example.com", Text: &notionTextContent{Content: "click"}, Annotations: notionAnnotations{}},
},
expected: "[click](https://example.com)",
},
{
name: "equation",
texts: []notionRichText{
{Type: "equation", PlainText: "E=mc^2", Equation: &notionEquation{Expression: "E=mc^2"}, Annotations: notionAnnotations{}},
},
expected: "$E=mc^2$",
},
{
name: "date mention",
texts: []notionRichText{
{Type: "mention", PlainText: "2026-01-15", Mention: &notionMention{
Type: "date", Date: &notionDateMention{Start: "2026-01-15"},
}, Annotations: notionAnnotations{}},
},
expected: "2026-01-15",
},
{
name: "date mention with end",
texts: []notionRichText{
{Type: "mention", PlainText: "range", Mention: &notionMention{
Type: "date", Date: &notionDateMention{Start: "2026-01-15", End: "2026-01-20"},
}, Annotations: notionAnnotations{}},
},
expected: "2026-01-15 \u2192 2026-01-20",
},
{
name: "multiple segments",
texts: []notionRichText{
{Type: "text", PlainText: "hello ", Text: &notionTextContent{Content: "hello "}, Annotations: notionAnnotations{}},
{Type: "text", PlainText: "world", Text: &notionTextContent{Content: "world"}, Annotations: notionAnnotations{Bold: true}},
},
expected: "hello **world**",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
got := renderRichText(tt.texts)
if got != tt.expected {
t.Errorf("renderRichText() = %q, want %q", got, tt.expected)
}
})
}
}
func makeBlock(t *testing.T, blockType string, content interface{}) notionBlock {
t.Helper()
raw, err := json.Marshal(content)
if err != nil {
t.Fatalf("marshal block content: %v", err)
}
return notionBlock{
ID: "blk-test",
Type: blockType,
RawContent: raw,
}
}
func TestBlocksToMarkdown_Paragraph(t *testing.T) {
blocks := []notionBlock{
makeBlock(t, "paragraph", map[string]interface{}{
"rich_text": []notionRichText{
{Type: "text", PlainText: "Hello world", Text: &notionTextContent{Content: "Hello world"}, Annotations: notionAnnotations{}},
},
}),
}
md, attachments := BlocksToMarkdown(blocks)
if md != "Hello world\n" {
t.Errorf("markdown = %q", md)
}
if len(attachments) != 0 {
t.Errorf("expected 0 attachments, got %d", len(attachments))
}
}
func TestBlocksToMarkdown_Headings(t *testing.T) {
blocks := []notionBlock{
makeBlock(t, "heading_1", map[string]interface{}{
"rich_text": []notionRichText{
{Type: "text", PlainText: "H1", Text: &notionTextContent{Content: "H1"}, Annotations: notionAnnotations{}},
},
}),
makeBlock(t, "heading_2", map[string]interface{}{
"rich_text": []notionRichText{
{Type: "text", PlainText: "H2", Text: &notionTextContent{Content: "H2"}, Annotations: notionAnnotations{}},
},
}),
makeBlock(t, "heading_3", map[string]interface{}{
"rich_text": []notionRichText{
{Type: "text", PlainText: "H3", Text: &notionTextContent{Content: "H3"}, Annotations: notionAnnotations{}},
},
}),
}
md, _ := BlocksToMarkdown(blocks)
expected := "# H1\n\n## H2\n\n### H3\n"
if md != expected {
t.Errorf("markdown = %q, want %q", md, expected)
}
}
func TestBlocksToMarkdown_Code(t *testing.T) {
blocks := []notionBlock{
makeBlock(t, "code", map[string]interface{}{
"language": "go",
"rich_text": []notionRichText{
{Type: "text", PlainText: "fmt.Println()", Text: &notionTextContent{Content: "fmt.Println()"}, Annotations: notionAnnotations{}},
},
}),
}
md, _ := BlocksToMarkdown(blocks)
expected := "```go\nfmt.Println()\n```\n"
if md != expected {
t.Errorf("markdown = %q, want %q", md, expected)
}
}
func TestBlocksToMarkdown_List(t *testing.T) {
blocks := []notionBlock{
makeBlock(t, "bulleted_list_item", map[string]interface{}{
"rich_text": []notionRichText{
{Type: "text", PlainText: "Item 1", Text: &notionTextContent{Content: "Item 1"}, Annotations: notionAnnotations{}},
},
}),
makeBlock(t, "bulleted_list_item", map[string]interface{}{
"rich_text": []notionRichText{
{Type: "text", PlainText: "Item 2", Text: &notionTextContent{Content: "Item 2"}, Annotations: notionAnnotations{}},
},
}),
}
md, _ := BlocksToMarkdown(blocks)
expected := "- Item 1\n- Item 2\n"
if md != expected {
t.Errorf("markdown = %q, want %q", md, expected)
}
}
func TestBlocksToMarkdown_ToDo(t *testing.T) {
blocks := []notionBlock{
makeBlock(t, "to_do", map[string]interface{}{
"checked": false,
"rich_text": []notionRichText{
{Type: "text", PlainText: "Unchecked", Text: &notionTextContent{Content: "Unchecked"}, Annotations: notionAnnotations{}},
},
}),
makeBlock(t, "to_do", map[string]interface{}{
"checked": true,
"rich_text": []notionRichText{
{Type: "text", PlainText: "Checked", Text: &notionTextContent{Content: "Checked"}, Annotations: notionAnnotations{}},
},
}),
}
md, _ := BlocksToMarkdown(blocks)
expected := "- [ ] Unchecked\n- [x] Checked\n"
if md != expected {
t.Errorf("markdown = %q, want %q", md, expected)
}
}
func TestBlocksToMarkdown_Quote(t *testing.T) {
blocks := []notionBlock{
makeBlock(t, "quote", map[string]interface{}{
"rich_text": []notionRichText{
{Type: "text", PlainText: "A wise quote", Text: &notionTextContent{Content: "A wise quote"}, Annotations: notionAnnotations{}},
},
}),
}
md, _ := BlocksToMarkdown(blocks)
expected := "> A wise quote\n"
if md != expected {
t.Errorf("markdown = %q, want %q", md, expected)
}
}
func TestBlocksToMarkdown_Divider(t *testing.T) {
blocks := []notionBlock{
{ID: "blk-div", Type: "divider"},
}
md, _ := BlocksToMarkdown(blocks)
if md != "---\n" {
t.Errorf("markdown = %q", md)
}
}
func TestBlocksToMarkdown_Equation(t *testing.T) {
blocks := []notionBlock{
makeBlock(t, "equation", map[string]interface{}{
"expression": "E = mc^2",
}),
}
md, _ := BlocksToMarkdown(blocks)
expected := "$$E = mc^2$$\n"
if md != expected {
t.Errorf("markdown = %q, want %q", md, expected)
}
}
func TestBlocksToMarkdown_Image(t *testing.T) {
blocks := []notionBlock{
makeBlock(t, "image", map[string]interface{}{
"type": "file",
"file": map[string]interface{}{
"url": "https://s3.example.com/img.png",
"expiry_time": "2026-01-15T11:00:00.000Z",
},
"caption": []notionRichText{
{Type: "text", PlainText: "A photo", Text: &notionTextContent{Content: "A photo"}, Annotations: notionAnnotations{}},
},
}),
}
md, attachments := BlocksToMarkdown(blocks)
if md != "![A photo](https://s3.example.com/img.png)\n" {
t.Errorf("markdown = %q", md)
}
if len(attachments) != 1 {
t.Fatalf("expected 1 attachment, got %d", len(attachments))
}
if attachments[0].URL != "https://s3.example.com/img.png" {
t.Errorf("attachment URL = %q", attachments[0].URL)
}
if attachments[0].Type != "image" {
t.Errorf("attachment Type = %q", attachments[0].Type)
}
}
func TestBlocksToMarkdown_UnknownType(t *testing.T) {
blocks := []notionBlock{
{ID: "blk-unknown", Type: "future_block_type"},
}
md, _ := BlocksToMarkdown(blocks)
if md != "\n" {
t.Errorf("unknown block should produce empty output, got %q", md)
}
}
func TestBlocksToMarkdown_Table(t *testing.T) {
tableBlock := notionBlock{
ID: "blk-table",
Type: "table",
HasChildren: true,
RawContent: json.RawMessage(`{"table_width": 2, "has_column_header": true}`),
Children: []notionBlock{
makeBlock(t, "table_row", map[string]interface{}{
"cells": [][]notionRichText{
{{Type: "text", PlainText: "Name", Text: &notionTextContent{Content: "Name"}, Annotations: notionAnnotations{}}},
{{Type: "text", PlainText: "Age", Text: &notionTextContent{Content: "Age"}, Annotations: notionAnnotations{}}},
},
}),
makeBlock(t, "table_row", map[string]interface{}{
"cells": [][]notionRichText{
{{Type: "text", PlainText: "Alice", Text: &notionTextContent{Content: "Alice"}, Annotations: notionAnnotations{}}},
{{Type: "text", PlainText: "30", Text: &notionTextContent{Content: "30"}, Annotations: notionAnnotations{}}},
},
}),
},
}
md, _ := BlocksToMarkdown([]notionBlock{tableBlock})
expected := "| Name | Age |\n| --- | --- |\n| Alice | 30 |\n"
if md != expected {
t.Errorf("markdown = %q, want %q", md, expected)
}
}
func TestBlocksToMarkdown_NestedList(t *testing.T) {
parent := makeBlock(t, "bulleted_list_item", map[string]interface{}{
"rich_text": []notionRichText{
{Type: "text", PlainText: "Parent", Text: &notionTextContent{Content: "Parent"}, Annotations: notionAnnotations{}},
},
})
parent.HasChildren = true
parent.Children = []notionBlock{
makeBlock(t, "bulleted_list_item", map[string]interface{}{
"rich_text": []notionRichText{
{Type: "text", PlainText: "Child", Text: &notionTextContent{Content: "Child"}, Annotations: notionAnnotations{}},
},
}),
}
md, _ := BlocksToMarkdown([]notionBlock{parent})
expected := "- Parent\n - Child\n"
if md != expected {
t.Errorf("markdown = %q, want %q", md, expected)
}
}
@@ -0,0 +1,272 @@
// Package notion implements the Notion data source connector for WeKnora.
//
// It syncs pages, databases, and attachments from Notion workspaces into WeKnora knowledge bases.
//
// Notion API docs:
// - Authentication: https://developers.notion.com/docs/authorization
// - Search: https://developers.notion.com/reference/post-search
// - Pages: https://developers.notion.com/reference/retrieve-a-page
// - Blocks: https://developers.notion.com/reference/retrieve-a-block
// - Databases: https://developers.notion.com/reference/retrieve-a-database
package notion
import (
"encoding/json"
"fmt"
"time"
"github.com/Tencent/WeKnora/internal/datasource"
"github.com/Tencent/WeKnora/internal/types"
)
// NotionAPIVersion is the Notion API version used by this connector.
const NotionAPIVersion = "2026-03-11"
// DefaultBaseURL is the Notion API base URL.
const DefaultBaseURL = "https://api.notion.com"
// Config holds Notion-specific configuration for the data source connector.
type Config struct {
APIKey string `json:"api_key"` // Internal Integration Token
}
// parseNotionConfig extracts and validates Notion config from DataSourceConfig.
func parseNotionConfig(config *types.DataSourceConfig) (*Config, error) {
if config == nil {
return nil, datasource.ErrInvalidConfig
}
tokenVal, ok := config.Credentials["api_key"]
if !ok {
return nil, fmt.Errorf("%w: missing api_key", datasource.ErrInvalidCredentials)
}
token, ok := tokenVal.(string)
if !ok || token == "" {
return nil, fmt.Errorf("%w: api_key must be a non-empty string", datasource.ErrInvalidCredentials)
}
return &Config{APIKey: token}, nil
}
// --- API response types ---
// notionPage represents a Notion page or database object.
type notionPage struct {
ID string `json:"id"`
Object string `json:"object"` // "page" | "database" | "data_source" (2025-09-03+)
Parent notionParent `json:"parent"`
URL string `json:"url"`
LastEditedTime time.Time `json:"last_edited_time"`
InTrash bool `json:"in_trash"`
// Title is extracted from properties after unmarshaling; not directly from JSON.
Title string `json:"-"`
// RawTitle holds the top-level title array (used by database objects).
RawTitle json.RawMessage `json:"title,omitempty"`
// RawProperties holds the raw JSON for property extraction.
RawProperties json.RawMessage `json:"properties"`
// DatabaseParent indicates where the database lives in the workspace hierarchy.
// Only present on data_source objects (API 2025-09-03+). For example, if a database
// is inside a page, this will be {type: "page_id", page_id: "..."}.
DatabaseParent *notionParent `json:"database_parent,omitempty"`
}
// Parent type constants for notionParent.Type
const (
parentTypeWorkspace = "workspace"
parentTypePageID = "page_id"
parentTypeDatabaseID = "database_id"
parentTypeDataSourceID = "data_source_id"
parentTypeBlockID = "block_id"
)
// notionParent represents the parent relationship of a page/database.
type notionParent struct {
Type string `json:"type"`
PageID string `json:"page_id,omitempty"`
DatabaseID string `json:"database_id,omitempty"`
DataSourceID string `json:"data_source_id,omitempty"`
BlockID string `json:"block_id,omitempty"`
}
// GetParentID returns the effective parent ID regardless of parent type.
func (p *notionParent) GetParentID() string {
switch p.Type {
case parentTypePageID:
return p.PageID
case parentTypeDatabaseID:
return p.DatabaseID
case parentTypeDataSourceID:
return p.DataSourceID
case parentTypeBlockID:
return p.BlockID
default:
return ""
}
}
// isDatabase returns true if the notionPage represents a database or data source.
func (p *notionPage) isDatabase() bool {
return p.Object == "database" || p.Object == "data_source"
}
// notionBlock represents a single content block in a Notion page.
// Requires custom UnmarshalJSON to extract the type-specific content.
type notionBlock struct {
ID string `json:"id"`
Type string `json:"type"`
HasChildren bool `json:"has_children"`
RawContent json.RawMessage `json:"-"` // Extracted from the type-named field (e.g., "paragraph")
Children []notionBlock `json:"-"` // Populated by client.GetBlockChildrenAll, not from API directly
}
// UnmarshalJSON implements custom JSON unmarshaling for notionBlock.
// It extracts the dynamic type-keyed field (e.g., {"type":"paragraph","paragraph":{...}})
// into RawContent.
func (b *notionBlock) UnmarshalJSON(data []byte) error {
// Use a helper struct to avoid infinite recursion
type blockAlias struct {
ID string `json:"id"`
Type string `json:"type"`
HasChildren bool `json:"has_children"`
}
var alias blockAlias
if err := json.Unmarshal(data, &alias); err != nil {
return err
}
b.ID = alias.ID
b.Type = alias.Type
b.HasChildren = alias.HasChildren
// Extract the type-named field into RawContent
if alias.Type != "" {
var raw map[string]json.RawMessage
if err := json.Unmarshal(data, &raw); err != nil {
return err
}
if content, ok := raw[alias.Type]; ok {
b.RawContent = content
}
}
return nil
}
// notionRichText represents a rich text element in Notion.
type notionRichText struct {
Type string `json:"type"` // "text" | "mention" | "equation"
PlainText string `json:"plain_text"`
Href string `json:"href,omitempty"`
Annotations notionAnnotations `json:"annotations"`
Text *notionTextContent `json:"text,omitempty"`
Mention *notionMention `json:"mention,omitempty"`
Equation *notionEquation `json:"equation,omitempty"`
}
// notionAnnotations holds text style information.
type notionAnnotations struct {
Bold bool `json:"bold"`
Italic bool `json:"italic"`
Strikethrough bool `json:"strikethrough"`
Underline bool `json:"underline"`
Code bool `json:"code"`
Color string `json:"color"`
}
// notionTextContent holds the content for "text" type rich text.
type notionTextContent struct {
Content string `json:"content"`
Link *struct {
URL string `json:"url"`
} `json:"link,omitempty"`
}
// notionMention holds the content for "mention" type rich text.
type notionMention struct {
Type string `json:"type"` // "page" | "database" | "date" | "user" | "link_preview"
Page *struct {
ID string `json:"id"`
} `json:"page,omitempty"`
Database *struct {
ID string `json:"id"`
} `json:"database,omitempty"`
Date *notionDateMention `json:"date,omitempty"`
LinkPreview *struct {
URL string `json:"url"`
} `json:"link_preview,omitempty"`
}
// notionDateMention holds date information for a mention.
type notionDateMention struct {
Start string `json:"start"`
End string `json:"end,omitempty"`
}
// notionEquation holds the content for "equation" type rich text.
type notionEquation struct {
Expression string `json:"expression"`
}
// notionFile represents a file object in Notion (used by image, file, pdf blocks).
type notionFile struct {
Type string `json:"type"` // "file" | "external" | "file_upload"
File *struct {
URL string `json:"url"`
ExpiryTime time.Time `json:"expiry_time"`
} `json:"file,omitempty"`
External *struct {
URL string `json:"url"`
} `json:"external,omitempty"`
FileUpload *struct {
ID string `json:"id"`
} `json:"file_upload,omitempty"`
Caption []notionRichText `json:"caption,omitempty"`
Name string `json:"name,omitempty"`
}
// GetURL returns the file URL regardless of type (hosted or external).
// For file_upload type, returns empty — caller must resolve via client.GetFileUploadURL().
func (f *notionFile) GetURL() string {
if f.File != nil {
return f.File.URL
}
if f.External != nil {
return f.External.URL
}
return ""
}
// GetFileUploadID returns the file upload ID for file_upload type, empty otherwise.
func (f *notionFile) GetFileUploadID() string {
if f.FileUpload != nil {
return f.FileUpload.ID
}
return ""
}
// --- Cursor ---
// notionCursor tracks state for incremental sync.
// Only PageEditTimes is used for diffing; LastSyncTime lives on types.SyncCursor.
type notionCursor struct {
PageEditTimes map[string]time.Time `json:"page_edit_times"` // page_id → last_edited_time
}
// --- Attachment ---
// attachment represents a file to be downloaded from a Notion page.
type attachment struct {
URL string // Notion S3 signed URL (expires in 1 hour)
FileName string
Type string // "image" | "file" | "pdf" | "video" | "audio"
}
// --- Pagination ---
// paginatedResponse is the common response wrapper for paginated Notion API responses.
type paginatedResponse struct {
Object string `json:"object"` // "list"
Results json.RawMessage `json:"results"`
HasMore bool `json:"has_more"`
NextCursor string `json:"next_cursor,omitempty"`
}
@@ -0,0 +1,144 @@
package notion
import (
"encoding/json"
"testing"
"time"
"github.com/Tencent/WeKnora/internal/types"
)
func TestParseNotionConfig(t *testing.T) {
t.Run("valid", func(t *testing.T) {
cfg, err := parseNotionConfig(&types.DataSourceConfig{
Credentials: map[string]interface{}{
"api_key": "ntn_test123",
},
})
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if cfg.APIKey != "ntn_test123" {
t.Errorf("APIKey = %q, want %q", cfg.APIKey, "ntn_test123")
}
})
t.Run("nil config", func(t *testing.T) {
_, err := parseNotionConfig(nil)
if err == nil {
t.Fatal("expected error for nil config")
}
})
t.Run("missing api_key", func(t *testing.T) {
_, err := parseNotionConfig(&types.DataSourceConfig{
Credentials: map[string]interface{}{},
})
if err == nil {
t.Fatal("expected error for missing api_key")
}
})
}
func TestNotionBlockUnmarshalJSON(t *testing.T) {
t.Run("paragraph block", func(t *testing.T) {
data := []byte(`{
"id": "block-1",
"type": "paragraph",
"has_children": false,
"paragraph": {"rich_text": [{"type": "text", "plain_text": "hello"}]}
}`)
var block notionBlock
if err := json.Unmarshal(data, &block); err != nil {
t.Fatalf("unmarshal error: %v", err)
}
if block.ID != "block-1" {
t.Errorf("ID = %q, want %q", block.ID, "block-1")
}
if block.Type != "paragraph" {
t.Errorf("Type = %q, want %q", block.Type, "paragraph")
}
if block.HasChildren != false {
t.Error("HasChildren should be false")
}
if block.RawContent == nil {
t.Fatal("RawContent should not be nil")
}
// Verify RawContent contains the paragraph object
var content struct {
RichText []notionRichText `json:"rich_text"`
}
if err := json.Unmarshal(block.RawContent, &content); err != nil {
t.Fatalf("unmarshal RawContent: %v", err)
}
if len(content.RichText) != 1 || content.RichText[0].PlainText != "hello" {
t.Errorf("unexpected rich_text content: %+v", content.RichText)
}
})
t.Run("block with children flag", func(t *testing.T) {
data := []byte(`{
"id": "block-2",
"type": "toggle",
"has_children": true,
"toggle": {"rich_text": [{"type": "text", "plain_text": "Details"}]}
}`)
var block notionBlock
if err := json.Unmarshal(data, &block); err != nil {
t.Fatalf("unmarshal error: %v", err)
}
if !block.HasChildren {
t.Error("HasChildren should be true")
}
// Children is populated by client, not by unmarshal
if block.Children != nil {
t.Error("Children should be nil after unmarshal")
}
})
t.Run("unknown block type", func(t *testing.T) {
data := []byte(`{"id": "block-3", "type": "future_type", "has_children": false}`)
var block notionBlock
if err := json.Unmarshal(data, &block); err != nil {
t.Fatalf("unmarshal error: %v", err)
}
if block.Type != "future_type" {
t.Errorf("Type = %q", block.Type)
}
// RawContent may be nil for unknown types without matching field
})
}
func TestNotionFileGetURL(t *testing.T) {
t.Run("hosted file", func(t *testing.T) {
f := &notionFile{
Type: "file",
File: &struct {
URL string `json:"url"`
ExpiryTime time.Time `json:"expiry_time"`
}{URL: "https://s3.example.com/file.pdf"},
}
if got := f.GetURL(); got != "https://s3.example.com/file.pdf" {
t.Errorf("GetURL() = %q", got)
}
})
t.Run("external file", func(t *testing.T) {
f := &notionFile{
Type: "external",
External: &struct {
URL string `json:"url"`
}{URL: "https://example.com/image.png"},
}
if got := f.GetURL(); got != "https://example.com/image.png" {
t.Errorf("GetURL() = %q", got)
}
})
t.Run("empty file", func(t *testing.T) {
f := &notionFile{Type: "file"}
if got := f.GetURL(); got != "" {
t.Errorf("GetURL() = %q, want empty", got)
}
})
}
+123
View File
@@ -0,0 +1,123 @@
package rss
import (
"bytes"
"context"
"fmt"
"io"
"net/http"
"net/url"
"time"
readability "codeberg.org/readeck/go-readability/v2"
"github.com/Tencent/WeKnora/internal/utils"
)
const (
// requestTimeout bounds a single feed or article fetch.
requestTimeout = 20 * time.Second
// maxFeedSize caps the feed body to avoid memory blowups on hostile feeds.
maxFeedSize = 10 * 1024 * 1024 // 10 MB
// maxArticleSize caps a single article page body.
maxArticleSize = 5 * 1024 * 1024 // 5 MB
// defaultUserAgent is sent on every request; some feeds reject empty UAs.
defaultUserAgent = "Mozilla/5.0 (compatible; WeKnora-RSS/1.0; +https://weknora.weixin.qq.com)"
)
// client performs SSRF-safe HTTP fetches with optional custom headers.
type client struct {
httpClient *http.Client
headers map[string]string
}
func newClient(headers map[string]string) *client {
cfg := utils.DefaultSSRFSafeHTTPClientConfig()
cfg.Timeout = requestTimeout
return &client{
httpClient: utils.NewSSRFSafeHTTPClient(cfg),
headers: headers,
}
}
// fetch retrieves rawURL with SSRF validation and size limiting. Custom auth
// headers are only attached when withAuthHeaders is true (feed fetches); article
// pages on third-party domains must not receive feed credentials.
func (c *client) fetch(ctx context.Context, rawURL string, maxSize int64, withAuthHeaders bool) ([]byte, error) {
if err := utils.ValidateURLForSSRF(rawURL); err != nil {
return nil, fmt.Errorf("URL rejected: %w", err)
}
if _, err := url.Parse(rawURL); err != nil {
return nil, fmt.Errorf("invalid URL: %w", err)
}
ctx, cancel := context.WithTimeout(ctx, requestTimeout)
defer cancel()
req, err := http.NewRequestWithContext(ctx, http.MethodGet, rawURL, nil)
if err != nil {
return nil, err
}
if withAuthHeaders {
for k, v := range c.headers {
req.Header.Set(k, v)
}
}
if req.Header.Get("User-Agent") == "" {
req.Header.Set("User-Agent", defaultUserAgent)
}
if req.Header.Get("Accept") == "" {
req.Header.Set("Accept",
"application/rss+xml, application/atom+xml, application/xml, text/xml, application/json, text/html;q=0.9, */*;q=0.8")
}
resp, err := c.httpClient.Do(req)
if err != nil {
return nil, fmt.Errorf("fetch failed: %w", err)
}
defer resp.Body.Close()
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
return nil, fmt.Errorf("HTTP %d %s", resp.StatusCode, resp.Status)
}
body, err := io.ReadAll(io.LimitReader(resp.Body, maxSize))
if err != nil {
return nil, fmt.Errorf("read body failed: %w", err)
}
return body, nil
}
// fetchFeed retrieves the raw bytes of a feed document.
func (c *client) fetchFeed(ctx context.Context, feedURL string) ([]byte, error) {
return c.fetch(ctx, feedURL, maxFeedSize, true)
}
// extractArticle fetches an article page and returns the readability-cleaned
// main content as HTML, plus the extracted title (may be empty). Returns an
// error if the page can't be fetched or no readable content is found, so the
// caller can fall back to feed-provided content.
func (c *client) extractArticle(ctx context.Context, articleURL string) (contentHTML, title string, err error) {
body, err := c.fetch(ctx, articleURL, maxArticleSize, false)
if err != nil {
return "", "", err
}
pageURL, _ := url.Parse(articleURL)
article, err := readability.FromReader(bytes.NewReader(body), pageURL)
if err != nil {
return "", "", fmt.Errorf("readability parse: %w", err)
}
if article.Node == nil {
return "", "", fmt.Errorf("no readable content extracted")
}
var buf bytes.Buffer
if err := article.RenderHTML(&buf); err != nil {
return "", "", fmt.Errorf("render article html: %w", err)
}
return buf.String(), article.Title(), nil
}
@@ -0,0 +1,366 @@
package rss
import (
"bytes"
"context"
"encoding/json"
"fmt"
"strings"
"time"
htmltomd "github.com/JohannesKaufmann/html-to-markdown/v2"
"github.com/Tencent/WeKnora/internal/datasource"
"github.com/Tencent/WeKnora/internal/logger"
"github.com/Tencent/WeKnora/internal/types"
"github.com/mmcdole/gofeed"
)
// Compile-time proof that *Connector satisfies the datasource.Connector interface.
var _ datasource.Connector = (*Connector)(nil)
// Connector implements datasource.Connector for RSS/Atom/JSON feeds.
type Connector struct{}
// NewConnector creates a new RSS connector.
func NewConnector() *Connector { return &Connector{} }
// Type returns the connector type identifier.
func (c *Connector) Type() string { return types.ConnectorTypeRSS }
// Validate verifies that every configured feed URL is reachable and parses as
// a valid feed.
func (c *Connector) Validate(ctx context.Context, config *types.DataSourceConfig) error {
cfg, err := parseConfig(config)
if err != nil {
return err
}
cli := newClient(cfg.parseHeaders())
parser := gofeed.NewParser()
for _, feedURL := range cfg.feedURLList() {
data, err := cli.fetchFeed(ctx, feedURL)
if err != nil {
return fmt.Errorf("fetch feed %s: %w", feedURL, err)
}
if _, err := parser.Parse(bytes.NewReader(data)); err != nil {
return fmt.Errorf("parse feed %s: %w", feedURL, err)
}
}
return nil
}
// ResolveResourceAncestors has nothing to do: feeds are a flat list with no
// nesting, so a selection has no ancestors to reveal.
func (c *Connector) ResolveResourceAncestors(
ctx context.Context, config *types.DataSourceConfig, resourceIDs []string,
) ([]string, error) {
return []string{}, nil
}
// ListResources returns one resource per configured feed URL. The feed is
// fetched so the resource can carry its real title; a feed that fails to fetch
// still appears (named by URL) with an error note, so the user can deselect it
// instead of the whole listing failing.
func (c *Connector) ListResources(
ctx context.Context, config *types.DataSourceConfig, parentID string,
) ([]types.Resource, error) {
// Feeds are flat: a lazy-load request for a specific parent has nothing extra.
if parentID != "" {
return []types.Resource{}, nil
}
cfg, err := parseConfig(config)
if err != nil {
return nil, err
}
cli := newClient(cfg.parseHeaders())
parser := gofeed.NewParser()
feedURLs := cfg.feedURLList()
out := make([]types.Resource, 0, len(feedURLs))
for _, feedURL := range feedURLs {
res := types.Resource{
ExternalID: feedURL,
Type: "feed",
Name: feedURL,
URL: feedURL,
}
data, err := cli.fetchFeed(ctx, feedURL)
if err != nil {
logger.Warnf(ctx, "[RSS] list: fetch %s failed: %v", feedURL, err)
res.Description = "fetch failed: " + err.Error()
out = append(out, res)
continue
}
feed, err := parser.Parse(bytes.NewReader(data))
if err != nil {
logger.Warnf(ctx, "[RSS] list: parse %s failed: %v", feedURL, err)
res.Description = "parse failed: " + err.Error()
out = append(out, res)
continue
}
if title := strings.TrimSpace(feed.Title); title != "" {
res.Name = title
}
res.Description = strings.TrimSpace(feed.Description)
if feed.Link != "" {
res.URL = feed.Link
}
if feed.UpdatedParsed != nil {
res.ModifiedAt = *feed.UpdatedParsed
}
res.Metadata = map[string]interface{}{"item_count": len(feed.Items)}
out = append(out, res)
}
return out, nil
}
// FetchAll performs a full sync of the specified feeds (or all configured feeds
// when resourceIDs is empty).
func (c *Connector) FetchAll(
ctx context.Context, config *types.DataSourceConfig, resourceIDs []string,
) ([]types.FetchedItem, error) {
items, _, err := c.walk(ctx, config, resourceIDs, nil, false)
return items, err
}
// FetchIncremental returns only items whose content fingerprint changed since
// the prior cursor. Deletions are intentionally not emitted (feeds drop old
// items as a matter of course).
func (c *Connector) FetchIncremental(
ctx context.Context, config *types.DataSourceConfig, cursor *types.SyncCursor,
) ([]types.FetchedItem, *types.SyncCursor, error) {
var prev *rssCursor
if cursor != nil && cursor.ConnectorCursor != nil {
var p rssCursor
b, err := json.Marshal(cursor.ConnectorCursor)
if err != nil {
logger.Warnf(ctx, "[RSS] marshal connector cursor: %v", err)
} else if err := json.Unmarshal(b, &p); err != nil {
logger.Warnf(ctx, "[RSS] unmarshal connector cursor: %v", err)
} else {
prev = &p
}
}
items, newCursor, err := c.walk(ctx, config, config.ResourceIDs, prev, true)
if err != nil && newCursor == nil {
return nil, nil, err
}
cursorMap := make(map[string]interface{})
if newCursor != nil {
b, marshalErr := json.Marshal(newCursor)
if marshalErr != nil {
logger.Warnf(ctx, "[RSS] marshal new cursor: %v", marshalErr)
} else if unmarshalErr := json.Unmarshal(b, &cursorMap); unmarshalErr != nil {
logger.Warnf(ctx, "[RSS] unmarshal new cursor to map: %v", unmarshalErr)
}
}
var syncCursor *types.SyncCursor
if newCursor != nil {
syncCursor = &types.SyncCursor{
LastSyncTime: newCursor.LastSyncTime,
ConnectorCursor: cursorMap,
}
}
return items, syncCursor, err
}
// walk is the shared implementation for FetchAll / FetchIncremental.
//
// When incremental is true, items whose feed signal and content fingerprint are
// both unchanged are omitted from the result (ingest skipped) without fetching
// article pages. The returned cursor always reflects the best-known item state
// so a later sync can detect changes.
func (c *Connector) walk(
ctx context.Context,
config *types.DataSourceConfig,
resourceIDs []string,
prev *rssCursor,
incremental bool,
) ([]types.FetchedItem, *rssCursor, error) {
cfg, err := parseConfig(config)
if err != nil {
return nil, nil, err
}
// Default to all configured feeds when no explicit selection was made.
feedURLs := resourceIDs
if len(feedURLs) == 0 {
feedURLs = cfg.feedURLList()
}
cli := newClient(cfg.parseHeaders())
parser := gofeed.NewParser()
newCursor := &rssCursor{
LastSyncTime: time.Now().UTC(),
FeedItems: make(map[string]map[string]string),
FeedSignals: make(map[string]map[string]string),
}
var out []types.FetchedItem
var feedErrors []string
for _, feedURL := range feedURLs {
data, err := cli.fetchFeed(ctx, feedURL)
if err != nil {
logger.Warnf(ctx, "[RSS] sync: fetch %s failed: %v", feedURL, err)
feedErrors = append(feedErrors, fmt.Sprintf("%s: %v", feedURL, err))
copyFeedCursor(newCursor, prev, feedURL)
continue
}
feed, err := parser.Parse(bytes.NewReader(data))
if err != nil {
logger.Warnf(ctx, "[RSS] sync: parse %s failed: %v", feedURL, err)
feedErrors = append(feedErrors, fmt.Sprintf("%s: %v", feedURL, err))
copyFeedCursor(newCursor, prev, feedURL)
continue
}
newCursor.FeedItems[feedURL] = make(map[string]string)
newCursor.FeedSignals[feedURL] = make(map[string]string)
var prevItems map[string]string
var prevSignals map[string]string
if incremental && prev != nil {
prevItems = prev.FeedItems[feedURL]
if prev.FeedSignals != nil {
prevSignals = prev.FeedSignals[feedURL]
}
}
var kept, skipped int
for _, item := range feed.Items {
if item == nil {
continue
}
itemID := firstNonEmpty(item.GUID, item.Link, item.Title)
if itemID == "" {
continue
}
feedContent := firstNonEmpty(item.Content, item.Description)
feedSig := feedSignalFingerprint(item, feedContent)
if incremental && prevItems != nil {
prevFP := prevItems[itemID]
prevSig := ""
if prevSignals != nil {
prevSig = prevSignals[itemID]
}
if prevFP != "" && feedSig == prevSig {
newCursor.FeedItems[feedURL][itemID] = prevFP
newCursor.FeedSignals[feedURL][itemID] = feedSig
skipped++
continue
}
}
resolved := c.resolveItem(ctx, cli, feed, item, feedURL, itemID, feedContent)
newCursor.FeedItems[feedURL][itemID] = resolved.fingerprint
newCursor.FeedSignals[feedURL][itemID] = feedSig
if incremental && prevItems != nil && prevItems[itemID] == resolved.fingerprint {
skipped++
continue
}
kept++
out = append(out, resolved.item)
}
logger.Infof(ctx, "[RSS] feed %s: items=%d fetched=%d skipped=%d",
feedURL, len(feed.Items), kept, skipped)
}
if len(feedErrors) > 0 {
if len(out) == 0 && len(feedErrors) == len(feedURLs) {
return nil, newCursor, fmt.Errorf("all feeds failed: %s", strings.Join(feedErrors, "; "))
}
return out, newCursor, &datasource.PartialFetchError{Details: feedErrors}
}
return out, newCursor, nil
}
type resolvedFeedItem struct {
item types.FetchedItem
fingerprint string
}
// resolveItem assembles a FetchedItem for a single feed entry, resolving the
// best available content (full-text article > feed content) and converting it
// to Markdown.
func (c *Connector) resolveItem(
ctx context.Context,
cli *client,
feed *gofeed.Feed,
item *gofeed.Item,
feedURL, itemID, feedContent string,
) resolvedFeedItem {
title := firstNonEmpty(item.Title, "untitled")
// Prefer full article text; fall back to feed-provided content on failure.
contentHTML := feedContent
if strings.TrimSpace(item.Link) != "" {
if articleHTML, articleTitle, err := cli.extractArticle(ctx, item.Link); err == nil {
contentHTML = articleHTML
if item.Title == "" && articleTitle != "" {
title = articleTitle
}
} else {
logger.Warnf(ctx, "[RSS] full-text fetch failed for %s (using feed content): %v", item.Link, err)
}
}
content := htmlToMarkdown(contentHTML)
updatedAt := time.Now().UTC()
switch {
case item.UpdatedParsed != nil && !item.UpdatedParsed.IsZero():
updatedAt = *item.UpdatedParsed
case item.PublishedParsed != nil && !item.PublishedParsed.IsZero():
updatedAt = *item.PublishedParsed
}
author := ""
if item.Author != nil {
author = item.Author.Name
}
return resolvedFeedItem{
fingerprint: contentFingerprint(content),
item: types.FetchedItem{
ExternalID: itemExternalID(feedURL, itemID),
Title: title,
Content: []byte(content),
ContentType: "text/markdown",
FileName: sanitizeFileName(title) + ".md",
URL: item.Link,
UpdatedAt: updatedAt,
SourceResourceID: feedURL,
Metadata: map[string]string{
"channel": types.ChannelRSS,
"feed_url": feedURL,
"feed_title": feed.Title,
"guid": item.GUID,
"link": item.Link,
"author": author,
},
},
}
}
// htmlToMarkdown converts HTML to Markdown, returning the trimmed original on
// conversion failure so we never silently drop content.
func htmlToMarkdown(html string) string {
if strings.TrimSpace(html) == "" {
return ""
}
md, err := htmltomd.ConvertString(html)
if err != nil || strings.TrimSpace(md) == "" {
return strings.TrimSpace(html)
}
return strings.TrimSpace(md)
}
@@ -0,0 +1,415 @@
package rss
import (
"context"
"encoding/json"
"errors"
"fmt"
"net/http"
"net/http/httptest"
"os"
"strings"
"sync/atomic"
"testing"
"github.com/Tencent/WeKnora/internal/datasource"
"github.com/Tencent/WeKnora/internal/types"
"github.com/Tencent/WeKnora/internal/utils"
"github.com/mmcdole/gofeed"
)
// TestMain whitelists loopback for SSRF so the httptest servers (127.0.0.1)
// are reachable. Production keeps the default strict SSRF policy.
func TestMain(m *testing.M) {
_ = os.Setenv("SSRF_WHITELIST", "127.0.0.1,::1")
utils.ResetSSRFWhitelistForTest()
code := m.Run()
os.Exit(code)
}
const longArticleBody = `<p>This is the first paragraph of a reasonably long article that the ` +
`readability extractor should detect as the main content of the page. It contains ` +
`enough words to clear the minimum content threshold used by the algorithm.</p>` +
`<p>The second paragraph continues the discussion with more sentences so that the ` +
`scoring heuristics confidently select this block over navigation and footer noise.</p>` +
`<p>A third paragraph adds further substance, ensuring the article is unmistakably ` +
`the dominant readable region of the document under test.</p>`
// fakeFeed spins up an httptest server serving an RSS feed and article pages.
type fakeFeed struct {
server *httptest.Server
feedTitle string
itemContent string // optional <description>/content for items
articleAuthHeaders []string
articleFetches atomic.Int32
failFeed atomic.Bool
}
func newFakeFeed(t *testing.T) *fakeFeed {
t.Helper()
f := &fakeFeed{feedTitle: "Test Feed"}
mux := http.NewServeMux()
mux.HandleFunc("/article/", func(w http.ResponseWriter, r *http.Request) {
f.articleFetches.Add(1)
for k, vals := range r.Header {
if strings.EqualFold(k, "X-Test-Auth") && len(vals) > 0 && vals[0] != "" {
f.articleAuthHeaders = append(f.articleAuthHeaders, vals[0])
}
}
w.Header().Set("Content-Type", "text/html; charset=utf-8")
fmt.Fprintf(w, `<!DOCTYPE html><html><head><title>%s</title></head>`+
`<body><nav>menu</nav><article><h1>Heading</h1>%s</article><footer>foot</footer></body></html>`,
"Full "+r.URL.Path, longArticleBody)
})
mux.HandleFunc("/feed.xml", func(w http.ResponseWriter, r *http.Request) {
if f.failFeed.Load() {
w.WriteHeader(http.StatusServiceUnavailable)
return
}
if auth := r.Header.Get("X-Test-Auth"); f.itemContent == "needs-auth" && auth != "secret" {
w.WriteHeader(http.StatusUnauthorized)
return
}
w.Header().Set("Content-Type", "application/rss+xml")
base := "http://" + r.Host
desc := "summary fallback"
fmt.Fprintf(w, `<?xml version="1.0"?>
<rss version="2.0"><channel>
<title>%s</title>
<link>%s</link>
<description>A test feed</description>
<item>
<title>Article One</title>
<link>%s/article/a1</link>
<guid>guid-1</guid>
<pubDate>Mon, 02 Jan 2006 15:04:05 GMT</pubDate>
<description>%s</description>
</item>
<item>
<title>Article Two</title>
<link>%s/article/a2</link>
<guid>guid-2</guid>
<pubDate>Tue, 03 Jan 2006 15:04:05 GMT</pubDate>
<description>%s</description>
</item>
</channel></rss>`, f.feedTitle, base, base, desc, base, desc)
})
f.server = httptest.NewServer(mux)
t.Cleanup(f.server.Close)
return f
}
func (f *fakeFeed) feedURL() string { return f.server.URL + "/feed.xml" }
func makeConfig(feedURLs string, headers string) *types.DataSourceConfig {
cfg := &types.DataSourceConfig{
Type: types.ConnectorTypeRSS,
Settings: map[string]interface{}{
"feed_urls": feedURLs,
},
Credentials: map[string]interface{}{},
}
if headers != "" {
cfg.Credentials["auth_headers"] = headers
}
return cfg
}
func TestConnector_Type(t *testing.T) {
if NewConnector().Type() != types.ConnectorTypeRSS {
t.Fatalf("Type() = %q, want %q", NewConnector().Type(), types.ConnectorTypeRSS)
}
}
func TestParseConfig_RequiresFeedURLs(t *testing.T) {
if _, err := parseConfig(makeConfig(" ", "")); err == nil {
t.Fatal("expected error when feed_urls is blank")
}
}
func TestParseConfig_FeedURLsFromSettings(t *testing.T) {
cfg, err := parseConfig(makeConfig("https://example.com/feed.xml", ""))
if err != nil {
t.Fatalf("parseConfig error: %v", err)
}
got := cfg.feedURLList()
if len(got) != 1 || got[0] != "https://example.com/feed.xml" {
t.Fatalf("feedURLList = %v", got)
}
}
func TestParseConfig_LegacyFeedURLsInCredentials(t *testing.T) {
legacy := &types.DataSourceConfig{
Type: types.ConnectorTypeRSS,
Credentials: map[string]interface{}{
"feed_urls": "https://legacy.example/feed.xml",
},
}
cfg, err := parseConfig(legacy)
if err != nil {
t.Fatalf("parseConfig error: %v", err)
}
if got := cfg.feedURLList(); len(got) != 1 || got[0] != "https://legacy.example/feed.xml" {
t.Fatalf("feedURLList = %v", got)
}
}
func TestConfig_FeedURLList_SplitsAndDedupes(t *testing.T) {
cfg := &Config{FeedURLs: "https://a.com/f, https://b.com/f\nhttps://a.com/f\n"}
got := cfg.feedURLList()
want := []string{"https://a.com/f", "https://b.com/f"}
if strings.Join(got, "|") != strings.Join(want, "|") {
t.Fatalf("feedURLList = %v, want %v", got, want)
}
}
func TestConfig_ParseHeaders(t *testing.T) {
cfg := &Config{AuthHeaders: "Authorization: Bearer x\nX-Foo: bar \nbroken-line\n: noname"}
got := cfg.parseHeaders()
if got["Authorization"] != "Bearer x" {
t.Errorf("Authorization = %q", got["Authorization"])
}
if got["X-Foo"] != "bar" {
t.Errorf("X-Foo = %q", got["X-Foo"])
}
if len(got) != 2 {
t.Errorf("expected 2 headers, got %d: %v", len(got), got)
}
}
func TestFeedSignalFingerprint(t *testing.T) {
item := &gofeed.Item{GUID: "g1", Link: "https://example.com/a", Title: "t"}
sig1 := feedSignalFingerprint(item, "body")
sig2 := feedSignalFingerprint(item, "body")
if sig1 == "" || sig1 != sig2 {
t.Fatalf("feed signal unstable: %q vs %q", sig1, sig2)
}
sig3 := feedSignalFingerprint(item, "changed")
if sig1 == sig3 {
t.Fatal("expected different signal when feed content changes")
}
}
func TestItemExternalID_ScopesByFeed(t *testing.T) {
got := itemExternalID("https://a.com/feed", "guid-1")
want := "https://a.com/feed:guid-1"
if got != want {
t.Fatalf("itemExternalID = %q, want %q", got, want)
}
}
func TestConnector_Validate_Success(t *testing.T) {
f := newFakeFeed(t)
if err := NewConnector().Validate(context.Background(), makeConfig(f.feedURL(), "")); err != nil {
t.Fatalf("Validate error: %v", err)
}
}
func TestConnector_Validate_PrivateFeedWithHeader(t *testing.T) {
f := newFakeFeed(t)
f.itemContent = "needs-auth"
// Without the header → 401.
if err := NewConnector().Validate(context.Background(), makeConfig(f.feedURL(), "")); err == nil {
t.Fatal("expected error without auth header")
}
// With the header → success.
if err := NewConnector().Validate(
context.Background(), makeConfig(f.feedURL(), "X-Test-Auth: secret"),
); err != nil {
t.Fatalf("Validate with header error: %v", err)
}
}
func TestConnector_FetchAll_DoesNotSendAuthHeadersToArticles(t *testing.T) {
f := newFakeFeed(t)
f.itemContent = "needs-auth"
_, err := NewConnector().FetchAll(
context.Background(), makeConfig(f.feedURL(), "X-Test-Auth: secret"), nil,
)
if err != nil {
t.Fatalf("FetchAll error: %v", err)
}
if len(f.articleAuthHeaders) != 0 {
t.Fatalf("article requests must not carry feed auth headers, got %v", f.articleAuthHeaders)
}
}
func TestConnector_ListResources(t *testing.T) {
f := newFakeFeed(t)
res, err := NewConnector().ListResources(context.Background(), makeConfig(f.feedURL(), ""), "")
if err != nil {
t.Fatalf("ListResources error: %v", err)
}
if len(res) != 1 {
t.Fatalf("expected 1 resource, got %d", len(res))
}
if res[0].ExternalID != f.feedURL() {
t.Errorf("ExternalID = %q, want %q", res[0].ExternalID, f.feedURL())
}
if res[0].Name != "Test Feed" {
t.Errorf("Name = %q, want %q", res[0].Name, "Test Feed")
}
// Non-empty parentID → empty (feeds are flat).
children, err := NewConnector().ListResources(context.Background(), makeConfig(f.feedURL(), ""), "feed-x")
if err != nil {
t.Fatalf("ListResources(parent) error: %v", err)
}
if len(children) != 0 {
t.Fatalf("expected no children, got %d", len(children))
}
}
func TestConnector_FetchAll_FullTextMarkdown(t *testing.T) {
f := newFakeFeed(t)
items, err := NewConnector().FetchAll(context.Background(), makeConfig(f.feedURL(), ""), nil)
if err != nil {
t.Fatalf("FetchAll error: %v", err)
}
if len(items) != 2 {
t.Fatalf("expected 2 items, got %d", len(items))
}
it := items[0]
wantExternalID := itemExternalID(f.feedURL(), "guid-1")
if it.ExternalID != wantExternalID {
t.Errorf("ExternalID = %q, want %q", it.ExternalID, wantExternalID)
}
if it.ContentType != "text/markdown" {
t.Errorf("ContentType = %q, want text/markdown", it.ContentType)
}
if it.Metadata["channel"] != types.ChannelRSS {
t.Errorf("channel = %q, want %q", it.Metadata["channel"], types.ChannelRSS)
}
// Full text from the article page should be present (not the short summary).
if !strings.Contains(string(it.Content), "first paragraph") {
t.Errorf("expected full article text in content, got: %q", string(it.Content))
}
if !strings.HasSuffix(it.FileName, ".md") {
t.Errorf("FileName = %q, want .md suffix", it.FileName)
}
}
func TestConnector_FetchIncremental_SkipsWithoutArticleFetch(t *testing.T) {
f := newFakeFeed(t)
cfg := makeConfig(f.feedURL(), "")
cfg.ResourceIDs = []string{f.feedURL()}
items, cursor, err := NewConnector().FetchIncremental(context.Background(), cfg, nil)
if err != nil {
t.Fatalf("first FetchIncremental error: %v", err)
}
if len(items) != 2 {
t.Fatalf("expected 2 items on first sync, got %d", len(items))
}
if got := f.articleFetches.Load(); got != 2 {
t.Fatalf("first sync article fetches = %d, want 2", got)
}
f.articleFetches.Store(0)
items2, _, err := NewConnector().FetchIncremental(context.Background(), cfg, cursor)
if err != nil {
t.Fatalf("second FetchIncremental error: %v", err)
}
if len(items2) != 0 {
t.Fatalf("expected 0 items on unchanged second sync, got %d", len(items2))
}
if got := f.articleFetches.Load(); got != 0 {
t.Fatalf("unchanged incremental sync must not refetch articles, got %d fetches", got)
}
}
func TestConnector_FetchIncremental_SkipsUnchanged(t *testing.T) {
f := newFakeFeed(t)
cfg := makeConfig(f.feedURL(), "")
cfg.ResourceIDs = []string{f.feedURL()}
// First sync: everything is new.
items, cursor, err := NewConnector().FetchIncremental(context.Background(), cfg, nil)
if err != nil {
t.Fatalf("first FetchIncremental error: %v", err)
}
if len(items) != 2 {
t.Fatalf("expected 2 items on first sync, got %d", len(items))
}
if cursor == nil {
t.Fatal("expected non-nil cursor")
}
// Second sync with the returned cursor: nothing changed → no items.
items2, _, err := NewConnector().FetchIncremental(context.Background(), cfg, cursor)
if err != nil {
t.Fatalf("second FetchIncremental error: %v", err)
}
if len(items2) != 0 {
t.Fatalf("expected 0 items on unchanged second sync, got %d", len(items2))
}
}
func TestConnector_Walk_PartialFeedFailure(t *testing.T) {
f := newFakeFeed(t)
cfg := makeConfig(f.feedURL()+", https://invalid.invalid/feed.xml", "")
items, err := NewConnector().FetchAll(context.Background(), cfg, nil)
var partial *datasource.PartialFetchError
if !errors.As(err, &partial) {
t.Fatalf("expected PartialFetchError, got %v", err)
}
if len(partial.Details) != 1 {
t.Fatalf("expected 1 feed error, got %v", partial.Details)
}
if len(items) != 2 {
t.Fatalf("expected 2 items from healthy feed, got %d", len(items))
}
}
func TestConnector_Walk_PreservesCursorOnFeedFailure(t *testing.T) {
f := newFakeFeed(t)
cfg := makeConfig(f.feedURL(), "")
cfg.ResourceIDs = []string{f.feedURL()}
_, cursor, err := NewConnector().FetchIncremental(context.Background(), cfg, nil)
if err != nil {
t.Fatalf("first sync error: %v", err)
}
f.failFeed.Store(true)
_, newCursor, err := NewConnector().FetchIncremental(context.Background(), cfg, cursor)
if err == nil {
t.Fatal("expected error when sole feed is unavailable")
}
if newCursor == nil || newCursor.ConnectorCursor == nil {
t.Fatal("expected cursor to be preserved on feed failure")
}
var restored rssCursor
b, err := json.Marshal(newCursor.ConnectorCursor)
if err != nil {
t.Fatalf("marshal cursor: %v", err)
}
if err := json.Unmarshal(b, &restored); err != nil {
t.Fatalf("unmarshal cursor: %v", err)
}
items := restored.FeedItems[f.feedURL()]
if len(items) != 2 {
t.Fatalf("expected preserved fingerprints for 2 items, got %d", len(items))
}
if items["guid-1"] == "" || items["guid-2"] == "" {
t.Fatalf("expected non-empty preserved fingerprints, got %+v", items)
}
}
func TestConnector_ResolveResourceAncestors_Empty(t *testing.T) {
got, err := NewConnector().ResolveResourceAncestors(
context.Background(), makeConfig("https://a.com/f", ""), []string{"x"},
)
if err != nil {
t.Fatalf("ResolveResourceAncestors error: %v", err)
}
if len(got) != 0 {
t.Fatalf("expected empty ancestors, got %v", got)
}
}
+249
View File
@@ -0,0 +1,249 @@
// Package rss implements the RSS/Atom data source connector for WeKnora.
//
// It syncs articles from one or more RSS/Atom/JSON feeds into a WeKnora
// knowledge base. Each configured feed URL is treated as a selectable
// resource; each feed item becomes a knowledge entry whose body is the
// article's full text rendered as Markdown.
//
// Capabilities:
// - Private feeds: optional custom HTTP headers (e.g. "Authorization: Bearer …")
// are attached only to feed fetches, never to third-party article pages.
// - Full text: when an item exposes a link, the article page is fetched and
// run through a readability extractor; the cleaned HTML is converted to
// Markdown. Feed-provided content (content:encoded / description) is used
// as a fallback.
// - Incremental: feed-level signals skip full-text article fetches when a feed
// entry is unchanged; content fingerprints detect feed-body changes. Article-
// only edits without feed updates may be missed until the feed entry changes.
// Deletions are NOT synced — feeds routinely drop old items.
//
// All outbound requests go through the SSRF-safe HTTP client so a malicious
// feed cannot redirect WeKnora to internal services.
package rss
import (
"crypto/sha256"
"encoding/hex"
"encoding/json"
"fmt"
"maps"
"strings"
"time"
"unicode/utf8"
"github.com/Tencent/WeKnora/internal/datasource"
"github.com/Tencent/WeKnora/internal/types"
"github.com/mmcdole/gofeed"
)
// Config holds RSS-specific configuration.
//
// FeedURLs are stored in DataSourceConfig.Settings (non-secret, editable in
// the UI without replacing credentials). AuthHeaders live in Credentials
// because they may carry secrets that must be encrypted at rest. Credentials
// may still carry feed_urls for backward compatibility with older rows.
type Config struct {
// FeedURLs is a newline- or comma-separated list of feed URLs.
FeedURLs string `json:"feed_urls"`
// AuthHeaders is an optional newline-separated list of custom request
// headers in "Name: Value" form, applied only to feed fetches.
AuthHeaders string `json:"auth_headers,omitempty"`
}
// parseConfig extracts and validates RSS configuration.
func parseConfig(config *types.DataSourceConfig) (*Config, error) {
if config == nil {
return nil, fmt.Errorf("%w: config is nil", datasource.ErrInvalidConfig)
}
credBytes, err := json.Marshal(config.Credentials)
if err != nil {
return nil, fmt.Errorf("marshal credentials: %w", err)
}
var cfg Config
if err := json.Unmarshal(credBytes, &cfg); err != nil {
return nil, fmt.Errorf("parse rss credentials: %w", err)
}
if urls := feedURLsFromSettings(config.Settings); urls != "" {
cfg.FeedURLs = urls
}
if len(cfg.feedURLList()) == 0 {
return nil, fmt.Errorf("%w: feed_urls is required", datasource.ErrInvalidCredentials)
}
return &cfg, nil
}
func feedURLsFromSettings(settings map[string]interface{}) string {
if len(settings) == 0 {
return ""
}
raw, ok := settings["feed_urls"]
if !ok {
return ""
}
s, ok := raw.(string)
if !ok {
return ""
}
return strings.TrimSpace(s)
}
// feedURLList splits FeedURLs on newlines and commas, trims, dedupes (order
// preserved), and drops blanks.
func (c *Config) feedURLList() []string {
if c == nil {
return nil
}
raw := strings.FieldsFunc(c.FeedURLs, func(r rune) bool {
return r == '\n' || r == '\r' || r == ','
})
seen := make(map[string]struct{}, len(raw))
out := make([]string, 0, len(raw))
for _, u := range raw {
u = strings.TrimSpace(u)
if u == "" {
continue
}
if _, ok := seen[u]; ok {
continue
}
seen[u] = struct{}{}
out = append(out, u)
}
return out
}
// parseHeaders turns the newline-separated "Name: Value" AuthHeaders blob into
// a map. Lines without a colon, or with an empty name, are skipped.
func (c *Config) parseHeaders() map[string]string {
if c == nil || strings.TrimSpace(c.AuthHeaders) == "" {
return nil
}
headers := make(map[string]string)
for _, line := range strings.Split(c.AuthHeaders, "\n") {
line = strings.TrimSpace(line)
if line == "" {
continue
}
idx := strings.Index(line, ":")
if idx <= 0 {
continue
}
name := strings.TrimSpace(line[:idx])
value := strings.TrimSpace(line[idx+1:])
if name == "" {
continue
}
headers[name] = value
}
if len(headers) == 0 {
return nil
}
return headers
}
// rssCursor stores incremental sync state.
//
// FeedItems maps feedURL → itemID → fingerprint, where fingerprint is a
// "h:<sha256-prefix>" hash of the final Markdown body that would be ingested.
// FeedSignals maps feedURL → itemID → feed-only signal used to skip full-text
// article fetches when the feed entry itself has not changed.
type rssCursor struct {
LastSyncTime time.Time `json:"last_sync_time"`
FeedItems map[string]map[string]string `json:"feed_items,omitempty"`
FeedSignals map[string]map[string]string `json:"feed_signals,omitempty"`
}
// contentFingerprint hashes ingested Markdown for incremental change detection.
func contentFingerprint(markdown string) string {
sum := sha256.Sum256([]byte(markdown))
return "h:" + hex.EncodeToString(sum[:])[:16]
}
// feedSignalFingerprint hashes feed-visible fields so incremental sync can skip
// article-page fetches when the entry has not changed in the feed.
func feedSignalFingerprint(item *gofeed.Item, feedContent string) string {
if item == nil {
return ""
}
var b strings.Builder
b.WriteString(item.GUID)
b.WriteByte('\n')
b.WriteString(item.Link)
b.WriteByte('\n')
b.WriteString(item.Title)
b.WriteByte('\n')
if item.UpdatedParsed != nil && !item.UpdatedParsed.IsZero() {
b.WriteString(item.UpdatedParsed.UTC().Format(time.RFC3339))
}
b.WriteByte('\n')
if item.PublishedParsed != nil && !item.PublishedParsed.IsZero() {
b.WriteString(item.PublishedParsed.UTC().Format(time.RFC3339))
}
b.WriteByte('\n')
b.WriteString(feedContent)
sum := sha256.Sum256([]byte(b.String()))
return "s:" + hex.EncodeToString(sum[:])[:16]
}
func copyFeedCursor(dst, prev *rssCursor, feedURL string) {
if dst == nil || prev == nil {
return
}
if src, ok := prev.FeedItems[feedURL]; ok && len(src) > 0 {
dst.FeedItems[feedURL] = maps.Clone(src)
}
if prev.FeedSignals != nil {
if src, ok := prev.FeedSignals[feedURL]; ok && len(src) > 0 {
if dst.FeedSignals == nil {
dst.FeedSignals = make(map[string]map[string]string)
}
dst.FeedSignals[feedURL] = maps.Clone(src)
}
}
}
// itemExternalID scopes an item ID to its feed so GUIDs cannot collide across feeds.
func itemExternalID(feedURL, itemID string) string {
return feedURL + ":" + itemID
}
// firstNonEmpty returns the first non-empty trimmed string among the args.
func firstNonEmpty(values ...string) string {
for _, v := range values {
if strings.TrimSpace(v) != "" {
return v
}
}
return ""
}
// sanitizeFileName removes characters invalid in filenames and truncates to a
// safe length at a UTF-8 rune boundary (mirrors the Yuque connector).
func sanitizeFileName(name string) string {
name = strings.TrimSpace(name)
if name == "" {
return "untitled"
}
replacer := strings.NewReplacer(
"/", "_", "\\", "_", ":", "_", "*", "_",
"?", "_", "\"", "_", "<", "_", ">", "_", "|", "_",
"\n", " ", "\r", " ", "\t", " ",
)
result := strings.TrimSpace(replacer.Replace(name))
if result == "" {
return "untitled"
}
const maxBytes = 200
if len(result) > maxBytes {
result = result[:maxBytes]
for len(result) > 0 {
r, size := utf8.DecodeLastRuneInString(result)
if r != utf8.RuneError || size != 1 {
break
}
result = result[:len(result)-1]
}
}
return result
}
@@ -0,0 +1,296 @@
package yuque
import (
"context"
"encoding/json"
"fmt"
"io"
"net/http"
"net/url"
"sync"
"time"
"github.com/Tencent/WeKnora/internal/datasource"
"github.com/Tencent/WeKnora/internal/logger"
)
const (
defaultTimeout = 30 * time.Second
defaultPageSize = 100
userAgent = "WeKnora-Yuque-Connector/1.0"
)
// client wraps the Yuque Open API.
type client struct {
baseURL string
token string
httpClient *http.Client
// logTokenOnce ensures the redacted token identity is logged at most once
// per client lifetime (first real request), rather than on every call.
logTokenOnce sync.Once
}
// newClient constructs a client with a normalized base URL.
func newClient(cfg *Config) *client {
return &client{
baseURL: cfg.GetBaseURL(),
token: cfg.APIToken,
httpClient: datasource.NewConnectorHTTPClient(defaultTimeout),
}
}
// doRequest executes an authenticated request and decodes JSON, with retry logic
// for transient errors (429, 5xx, transport failures).
// The raw X-Auth-Token is never logged; only a redacted form is emitted once per
// client lifetime (not per request) to keep sync logs readable at thousand-doc scale.
func (c *client) doRequest(ctx context.Context, method, path string, result interface{}) error {
const (
maxRetries = 3
max5xxRetries = 1
retry5xxDelay = 2 * time.Second
)
var lastErr error
backoff := []time.Duration{2 * time.Second, 4 * time.Second, 8 * time.Second}
c.logTokenOnce.Do(func() {
logger.Infof(ctx, "[Yuque] client configured token=%s base=%s", redactToken(c.token), c.baseURL)
})
for attempt := 0; attempt <= maxRetries; attempt++ {
reqURL := c.baseURL + path
req, err := http.NewRequestWithContext(ctx, method, reqURL, nil)
if err != nil {
return fmt.Errorf("create request: %w", err)
}
req.Header.Set("X-Auth-Token", c.token)
req.Header.Set("User-Agent", userAgent)
req.Header.Set("Content-Type", "application/json; charset=utf-8")
if attempt == 0 {
logger.Infof(ctx, "[Yuque] %s %s", method, path)
} else {
logger.Infof(ctx, "[Yuque] %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
}
body, 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
}
bodyPreview := truncate(string(body), 500)
logger.Infof(ctx, "[Yuque] %s %s → status=%d bodyLen=%d body=%s",
method, path, resp.StatusCode, len(body), bodyPreview)
if resp.StatusCode == 429 {
wait := parseRetryAfter(resp.Header.Get("Retry-After"), backoff[min(attempt, len(backoff)-1)])
lastErr = fmt.Errorf("yuque rate limited: status=429 body=%s", bodyPreview)
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("yuque server error: status=%d body=%s", resp.StatusCode, bodyPreview)
if attempt < max5xxRetries {
if sErr := sleepCtx(ctx, retry5xxDelay); sErr != nil {
return sErr
}
continue
}
return lastErr
}
// 401/403 → surface as ErrInvalidCredentials so DataSourceService can
// distinguish bad-token from transient failures and auto-flag the source.
if resp.StatusCode == http.StatusUnauthorized || resp.StatusCode == http.StatusForbidden {
return fmt.Errorf("%w: status=%d body=%s", datasource.ErrInvalidCredentials, resp.StatusCode, bodyPreview)
}
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
var apiErr apiErrorBody
_ = json.Unmarshal(body, &apiErr)
if apiErr.Message != "" {
return fmt.Errorf("yuque api error: status=%d msg=%s", resp.StatusCode, apiErr.Message)
}
return fmt.Errorf("yuque api error: status=%d body=%s", resp.StatusCode, bodyPreview)
}
if result != nil {
if err := json.Unmarshal(body, result); err != nil {
return fmt.Errorf("decode response: %w", err)
}
}
return nil
}
return lastErr
}
// parseRetryAfter returns the Retry-After duration from the header, or fallback if unparseable.
// Retry-After: "0" (or negative) is coerced to 100ms so we still yield and don't busy-retry.
// Note: only integer-seconds form is supported (RFC 7231 also allows HTTP-date — not seen from Yuque).
func parseRetryAfter(header string, fallback time.Duration) time.Duration {
if header == "" {
return fallback
}
if secs, err := time.ParseDuration(header + "s"); err == nil {
if secs <= 0 {
return 100 * time.Millisecond
}
return secs
}
return fallback
}
// sleepCtx pauses for d, returning early if ctx is cancelled.
func sleepCtx(ctx context.Context, d time.Duration) error {
t := time.NewTimer(d)
defer t.Stop()
select {
case <-t.C:
return nil
case <-ctx.Done():
return ctx.Err()
}
}
// Ping verifies the credentials by calling GET /api/v2/user.
func (c *client) Ping(ctx context.Context) error {
var resp v2UserResponse
return c.doRequest(ctx, http.MethodGet, "/api/v2/user", &resp)
}
// truncate returns s truncated to maxLen with "..." appended if longer.
func truncate(s string, maxLen int) string {
if len(s) <= maxLen {
return s
}
return s[:maxLen] + "..."
}
// buildQuery encodes query parameters, omitting empty values.
func buildQuery(params map[string]string) string {
values := url.Values{}
for k, v := range params {
if v != "" {
values.Set(k, v)
}
}
if len(values) == 0 {
return ""
}
return "?" + values.Encode()
}
// GetCurrentUser returns the user associated with the current token.
func (c *client) GetCurrentUser(ctx context.Context) (v2User, error) {
var resp v2UserResponse
if err := c.doRequest(ctx, http.MethodGet, "/api/v2/user", &resp); err != nil {
return v2User{}, err
}
return resp.Data, nil
}
// ListUserGroups returns the groups the given user belongs to.
// Note: userID is Yuque's numeric user ID (not the login) — the /users/{id}/groups
// endpoint requires the integer ID form.
func (c *client) ListUserGroups(ctx context.Context, userID int64) ([]v2Group, error) {
path := fmt.Sprintf("/api/v2/users/%d/groups", userID)
var resp v2GroupListResponse
if err := c.doRequest(ctx, http.MethodGet, path, &resp); err != nil {
return nil, err
}
return resp.Data, nil
}
// ListUserRepos returns document-type (Book) repos owned by the given user login.
func (c *client) ListUserRepos(ctx context.Context, login string) ([]v2Repo, error) {
return c.listReposPaginated(ctx, fmt.Sprintf("/api/v2/users/%s/repos", login))
}
// ListGroupRepos returns document-type (Book) repos owned by the given group login.
func (c *client) ListGroupRepos(ctx context.Context, login string) ([]v2Repo, error) {
return c.listReposPaginated(ctx, fmt.Sprintf("/api/v2/groups/%s/repos", login))
}
// listReposPaginated walks the pagination for a repo listing endpoint.
// Filters to type=Book (document-type knowledge bases only; design/sheet/resource skipped).
func (c *client) listReposPaginated(ctx context.Context, basePath string) ([]v2Repo, error) {
var all []v2Repo
offset := 0
for {
q := buildQuery(map[string]string{
"type": "Book",
"offset": fmt.Sprintf("%d", offset),
"limit": fmt.Sprintf("%d", defaultPageSize),
})
var resp v2RepoListResponse
if err := c.doRequest(ctx, http.MethodGet, basePath+q, &resp); err != nil {
return nil, err
}
all = append(all, resp.Data...)
if len(resp.Data) < defaultPageSize {
break
}
offset += defaultPageSize
}
return all, nil
}
// ListBookDocs lists all documents in a book, handling pagination.
// Returns summaries only (body not included — use GetDocDetail to fetch body per doc).
func (c *client) ListBookDocs(ctx context.Context, bookID int64) ([]v2Doc, error) {
basePath := fmt.Sprintf("/api/v2/repos/%d/docs", bookID)
var all []v2Doc
offset := 0
for {
q := buildQuery(map[string]string{
"offset": fmt.Sprintf("%d", offset),
"limit": fmt.Sprintf("%d", defaultPageSize),
})
var resp v2DocListResponse
if err := c.doRequest(ctx, http.MethodGet, basePath+q, &resp); err != nil {
return nil, err
}
all = append(all, resp.Data...)
if len(resp.Data) < defaultPageSize {
break
}
offset += defaultPageSize
}
return all, nil
}
// GetDocDetail fetches the full document detail (including body) by doc ID.
func (c *client) GetDocDetail(ctx context.Context, docID int64) (v2DocDetail, error) {
path := fmt.Sprintf("/api/v2/repos/docs/%d", docID)
var resp v2DocDetailResponse
if err := c.doRequest(ctx, http.MethodGet, path, &resp); err != nil {
return v2DocDetail{}, err
}
return resp.Data, nil
}
@@ -0,0 +1,422 @@
package yuque
import (
"bytes"
"context"
"encoding/json"
"errors"
"fmt"
"io"
"net/http"
"net/http/httptest"
"os"
"strings"
"testing"
"time"
"github.com/Tencent/WeKnora/internal/datasource"
"github.com/Tencent/WeKnora/internal/logger"
)
// fakeYuque sets up an httptest server emulating Yuque API endpoints.
// Handlers can be overridden per test.
type fakeYuque struct {
server *httptest.Server
mux *http.ServeMux
calls []string // methodPath history
}
func newFakeYuque() *fakeYuque {
f := &fakeYuque{mux: http.NewServeMux()}
f.server = httptest.NewServer(f.mux)
// Default /api/v2/user handler (most tests hit it at least once).
f.handleJSON("/api/v2/user", 200, v2UserResponse{Data: v2User{ID: 1, Login: "me", Name: "Me"}})
return f
}
func (f *fakeYuque) Close() { f.server.Close() }
func (f *fakeYuque) cfg() *Config {
return &Config{APIToken: "tok-super-secret-value-1234", BaseURL: f.server.URL}
}
func (f *fakeYuque) handleJSON(path string, status int, body interface{}) {
f.mux.HandleFunc(path, func(w http.ResponseWriter, r *http.Request) {
f.calls = append(f.calls, r.Method+" "+r.URL.Path+"?"+r.URL.RawQuery)
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(status)
_ = json.NewEncoder(w).Encode(body)
})
}
func TestClient_Ping_Success(t *testing.T) {
f := newFakeYuque()
defer f.Close()
c := newClient(f.cfg())
if err := c.Ping(context.Background()); err != nil {
t.Fatalf("Ping error: %v", err)
}
}
func TestClient_Ping_SendsAuthHeader(t *testing.T) {
f := newFakeYuque()
defer f.Close()
var gotToken string
f.mux = http.NewServeMux()
f.mux.HandleFunc("/api/v2/user", func(w http.ResponseWriter, r *http.Request) {
gotToken = r.Header.Get("X-Auth-Token")
w.Header().Set("Content-Type", "application/json")
_ = json.NewEncoder(w).Encode(v2UserResponse{Data: v2User{ID: 1, Login: "me"}})
})
f.server.Config.Handler = f.mux
c := newClient(f.cfg())
if err := c.Ping(context.Background()); err != nil {
t.Fatalf("Ping error: %v", err)
}
if gotToken != "tok-super-secret-value-1234" {
t.Errorf("X-Auth-Token = %q, want the token", gotToken)
}
}
func TestClient_Ping_401(t *testing.T) {
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.WriteHeader(401)
_, _ = io.WriteString(w, `{"message":"Unauthorized"}`)
}))
defer srv.Close()
c := newClient(&Config{APIToken: "bad", BaseURL: srv.URL})
err := c.Ping(context.Background())
if err == nil {
t.Fatal("expected error on 401")
}
if !errors.Is(err, datasource.ErrInvalidCredentials) {
t.Errorf("401 should wrap ErrInvalidCredentials, got: %v", err)
}
}
func TestClient_Ping_403WrapsInvalidCredentials(t *testing.T) {
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.WriteHeader(403)
_, _ = io.WriteString(w, `{"message":"Forbidden"}`)
}))
defer srv.Close()
c := newClient(&Config{APIToken: "insufficient", BaseURL: srv.URL})
err := c.Ping(context.Background())
if err == nil {
t.Fatal("expected error on 403")
}
if !errors.Is(err, datasource.ErrInvalidCredentials) {
t.Errorf("403 should wrap ErrInvalidCredentials, got: %v", err)
}
}
func TestClient_TokenNeverLoggedInFull(t *testing.T) {
t.Setenv("LOG_FORMAT", "")
logger.ConfigureFromEnv()
t.Cleanup(func() { logger.ConfigureFromEnv() })
// Redirect the project's internal logger to an in-memory buffer so we can
// assert the raw token never appears in log output. Using stdlib `log`
// would be vacuous — the real logger is a private logrus instance.
var buf bytes.Buffer
logger.SetOutput(&buf)
defer logger.SetOutput(os.Stdout)
f := newFakeYuque()
defer f.Close()
rawToken := "tok-super-secret-value-1234"
cfg := &Config{APIToken: rawToken, BaseURL: f.server.URL}
c := newClient(cfg)
_ = c.Ping(context.Background())
out := buf.String()
if out == "" {
t.Fatal("expected logger output (sanity check — SetOutput wiring may be broken)")
}
if strings.Contains(out, rawToken) {
t.Errorf("raw token leaked into logs:\n%s", out)
}
// Verify the redacted form was emitted (positive assertion, not just absence).
if !strings.Contains(out, redactToken(rawToken)) {
t.Errorf("expected redacted token %q in logs, got:\n%s", redactToken(rawToken), out)
}
}
func TestRedactToken(t *testing.T) {
tests := []struct{ in, want string }{
{"short", "***"},
{"abcdef1234567890", "abcdef...7890"},
}
for _, tt := range tests {
if got := redactToken(tt.in); got != tt.want {
t.Errorf("redactToken(%q) = %q, want %q", tt.in, got, tt.want)
}
}
}
func TestClient_429WithRetryAfter_Retries(t *testing.T) {
attempt := 0
mux := http.NewServeMux()
mux.HandleFunc("/api/v2/user", func(w http.ResponseWriter, r *http.Request) {
attempt++
if attempt == 1 {
// "0" is coerced to 100ms inside the client so the test stays fast.
w.Header().Set("Retry-After", "0")
w.WriteHeader(429)
_, _ = io.WriteString(w, `{"message":"rate limited"}`)
return
}
w.Header().Set("Content-Type", "application/json")
_ = json.NewEncoder(w).Encode(v2UserResponse{Data: v2User{ID: 1, Login: "me"}})
})
srv := httptest.NewServer(mux)
defer srv.Close()
c := newClient(&Config{APIToken: "t", BaseURL: srv.URL})
if err := c.Ping(context.Background()); err != nil {
t.Fatalf("expected success after retry, got %v", err)
}
if attempt < 2 {
t.Errorf("expected at least 2 attempts, got %d", attempt)
}
}
func TestClient_429ExhaustsRetries(t *testing.T) {
attempts := 0
mux := http.NewServeMux()
mux.HandleFunc("/api/v2/user", func(w http.ResponseWriter, r *http.Request) {
attempts++
w.Header().Set("Retry-After", "0")
w.WriteHeader(429)
_, _ = io.WriteString(w, `{"message":"rate limited"}`)
})
srv := httptest.NewServer(mux)
defer srv.Close()
c := newClient(&Config{APIToken: "t", BaseURL: srv.URL})
err := c.Ping(context.Background())
if err == nil {
t.Fatal("expected error when 429s exceed retry budget")
}
if attempts != 4 { // initial + 3 retries
t.Errorf("attempts = %d, want 4 (1 + 3 retries)", attempts)
}
}
func TestClient_5xxRetriesOnce(t *testing.T) {
attempts := 0
mux := http.NewServeMux()
mux.HandleFunc("/api/v2/user", func(w http.ResponseWriter, r *http.Request) {
attempts++
w.WriteHeader(500)
_, _ = io.WriteString(w, `{"message":"internal error"}`)
})
srv := httptest.NewServer(mux)
defer srv.Close()
// Make the fixed 2-second 5xx sleep negligible by using a context whose
// deadline fires after just one retry cycle. We still expect exactly 2
// attempts because the implementation should retry 5xx once and then stop.
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
c := newClient(&Config{APIToken: "t", BaseURL: srv.URL})
if err := c.Ping(ctx); 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 TestClient_4xxNotRetried(t *testing.T) {
attempts := 0
mux := http.NewServeMux()
mux.HandleFunc("/api/v2/user", func(w http.ResponseWriter, r *http.Request) {
attempts++
w.WriteHeader(400)
_, _ = io.WriteString(w, `{"message":"bad request"}`)
})
srv := httptest.NewServer(mux)
defer srv.Close()
c := newClient(&Config{APIToken: "t", BaseURL: srv.URL})
if err := c.Ping(context.Background()); 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 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
{"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)
}
}
}
func TestClient_GetCurrentUser(t *testing.T) {
f := newFakeYuque()
defer f.Close()
// Override the default /api/v2/user handler via a fresh mux.
f.mux = http.NewServeMux()
f.mux.HandleFunc("/api/v2/user", func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "application/json")
_ = json.NewEncoder(w).Encode(v2UserResponse{Data: v2User{ID: 42, Login: "alice", Name: "Alice"}})
})
f.server.Config.Handler = f.mux
c := newClient(f.cfg())
u, err := c.GetCurrentUser(context.Background())
if err != nil {
t.Fatalf("GetCurrentUser: %v", err)
}
if u.ID != 42 || u.Login != "alice" {
t.Errorf("got %+v, want {42 alice}", u)
}
}
func TestClient_ListUserGroups(t *testing.T) {
f := newFakeYuque()
defer f.Close()
f.handleJSON("/api/v2/users/42/groups", 200, v2GroupListResponse{Data: []v2Group{
{ID: 100, Login: "team-a", Name: "Team A"},
{ID: 101, Login: "team-b", Name: "Team B"},
}})
c := newClient(f.cfg())
gs, err := c.ListUserGroups(context.Background(), 42)
if err != nil {
t.Fatalf("ListUserGroups: %v", err)
}
if len(gs) != 2 || gs[0].Login != "team-a" {
t.Errorf("got %+v", gs)
}
}
func TestClient_ListUserRepos_FiltersType(t *testing.T) {
f := newFakeYuque()
defer f.Close()
var gotQuery string
f.mux.HandleFunc("/api/v2/users/alice/repos", func(w http.ResponseWriter, r *http.Request) {
gotQuery = r.URL.RawQuery
w.Header().Set("Content-Type", "application/json")
_ = json.NewEncoder(w).Encode(v2RepoListResponse{Data: []v2Repo{
{ID: 9, Slug: "book1", Name: "Book 1", Type: "Book", Namespace: "alice/book1"},
}})
})
c := newClient(f.cfg())
repos, err := c.ListUserRepos(context.Background(), "alice")
if err != nil {
t.Fatalf("ListUserRepos: %v", err)
}
if !strings.Contains(gotQuery, "type=Book") {
t.Errorf("expected type=Book in query, got %q", gotQuery)
}
if !strings.Contains(gotQuery, "offset=0") {
t.Errorf("expected offset=0 in query, got %q", gotQuery)
}
if !strings.Contains(gotQuery, "limit=100") {
t.Errorf("expected limit=100 in query, got %q", gotQuery)
}
if len(repos) != 1 || repos[0].Namespace != "alice/book1" {
t.Errorf("got %+v", repos)
}
}
func TestClient_ListGroupRepos(t *testing.T) {
f := newFakeYuque()
defer f.Close()
f.handleJSON("/api/v2/groups/team-a/repos", 200, v2RepoListResponse{Data: []v2Repo{
{ID: 11, Slug: "tb", Name: "Team Book", Type: "Book", Namespace: "team-a/tb"},
}})
c := newClient(f.cfg())
repos, err := c.ListGroupRepos(context.Background(), "team-a")
if err != nil {
t.Fatalf("ListGroupRepos: %v", err)
}
if len(repos) != 1 || repos[0].ID != 11 {
t.Errorf("got %+v", repos)
}
}
func TestClient_ListBookDocs_Pagination(t *testing.T) {
f := newFakeYuque()
defer f.Close()
// Page 1: 100 docs, page 2: 30 docs, then stop.
callCount := 0
f.mux.HandleFunc("/api/v2/repos/555/docs", func(w http.ResponseWriter, r *http.Request) {
callCount++
offset := r.URL.Query().Get("offset")
docs := make([]v2Doc, 0)
switch offset {
case "0":
for i := 0; i < 100; i++ {
docs = append(docs, v2Doc{ID: int64(i + 1), Type: "Doc", Status: "1", Title: fmt.Sprintf("D%d", i+1), ContentUpdatedAt: "2026-04-20T00:00:00Z"})
}
case "100":
for i := 0; i < 30; i++ {
docs = append(docs, v2Doc{ID: int64(i + 101), Type: "Doc", Status: "1", Title: fmt.Sprintf("D%d", i+101), ContentUpdatedAt: "2026-04-20T00:00:00Z"})
}
}
w.Header().Set("Content-Type", "application/json")
_ = json.NewEncoder(w).Encode(v2DocListResponse{Data: docs})
})
c := newClient(f.cfg())
docs, err := c.ListBookDocs(context.Background(), 555)
if err != nil {
t.Fatalf("ListBookDocs: %v", err)
}
if len(docs) != 130 {
t.Errorf("len(docs) = %d, want 130", len(docs))
}
if callCount != 2 {
t.Errorf("callCount = %d, want 2 (one full page + one partial)", callCount)
}
}
func TestClient_GetDocDetail(t *testing.T) {
f := newFakeYuque()
defer f.Close()
f.handleJSON("/api/v2/repos/docs/7", 200, v2DocDetailResponse{
Data: v2DocDetail{
ID: 7,
Title: "Hello",
Body: "# Hello\n\nworld",
Format: "markdown",
Status: "1",
ContentUpdatedAt: "2026-04-20T12:00:00Z",
},
})
c := newClient(f.cfg())
d, err := c.GetDocDetail(context.Background(), 7)
if err != nil {
t.Fatalf("GetDocDetail: %v", err)
}
if d.Body != "# Hello\n\nworld" {
t.Errorf("Body = %q", d.Body)
}
}
@@ -0,0 +1,350 @@
package yuque
import (
"context"
"encoding/json"
"fmt"
"sort"
"strconv"
"time"
"github.com/Tencent/WeKnora/internal/datasource"
"github.com/Tencent/WeKnora/internal/logger"
"github.com/Tencent/WeKnora/internal/types"
)
// Compile-time proof that *Connector satisfies the datasource.Connector interface.
// Catches signature drift as soon as the interface changes, rather than at
// container wiring or runtime.
var _ datasource.Connector = (*Connector)(nil)
// Connector implements datasource.Connector for Yuque.
type Connector struct{}
// NewConnector creates a new Yuque connector.
func NewConnector() *Connector { return &Connector{} }
// Type returns the connector type identifier.
func (c *Connector) Type() string { return types.ConnectorTypeYuque }
// Validate verifies the given credentials by pinging the current-user endpoint.
func (c *Connector) Validate(ctx context.Context, config *types.DataSourceConfig) error {
cfg, err := parseYuqueConfig(config)
if err != nil {
return err
}
cli := newClient(cfg)
if err := cli.Ping(ctx); err != nil {
return fmt.Errorf("yuque connection failed: %w", err)
}
return nil
}
// ResolveResourceAncestors has nothing to do for Yuque: repositories are a flat
// list with no nesting, so a selection has no ancestors to reveal.
func (c *Connector) ResolveResourceAncestors(
ctx context.Context, config *types.DataSourceConfig, resourceIDs []string,
) ([]string, error) {
return []string{}, nil
}
// ListResources returns all repos (personal + team) accessible to the token.
// Serial fetch for v1 (user groups typically <10). TODO(perf): parallelize if slow.
func (c *Connector) ListResources(
ctx context.Context, config *types.DataSourceConfig, parentID string,
) ([]types.Resource, error) {
// Yuque resources are a flat list of repositories (no nesting), so a
// lazy-load request for a specific parent has nothing extra to return.
if parentID != "" {
return []types.Resource{}, nil
}
cfg, err := parseYuqueConfig(config)
if err != nil {
return nil, err
}
cli := newClient(cfg)
me, err := cli.GetCurrentUser(ctx)
if err != nil {
return nil, fmt.Errorf("get current user: %w", err)
}
repos := make(map[int64]v2Repo)
// Team token: /api/v2/user returns type="Group" — the token represents a team,
// not a personal user. In this case, directly list the team's own repos instead
// of going through the personal-repos + user-groups flow.
if me.Type == "Group" {
logger.Infof(ctx, "[Yuque] detected team token (type=Group, login=%s), listing team repos directly", me.Login)
teamRepos, err := cli.ListGroupRepos(ctx, me.Login)
if err != nil {
return nil, fmt.Errorf("list team repos: %w", err)
}
for _, r := range teamRepos {
repos[r.ID] = r
}
} else {
// Personal token flow: list user's own repos + repos from joined groups.
personal, err := cli.ListUserRepos(ctx, me.Login)
if err != nil {
return nil, fmt.Errorf("list personal repos: %w", err)
}
for _, r := range personal {
if _, ok := repos[r.ID]; !ok {
repos[r.ID] = r
}
}
groups, err := cli.ListUserGroups(ctx, me.ID)
if err != nil {
// Yuque returns 404 when the user has not joined any groups (teams),
// instead of an empty list. Treat this as "no groups" and continue
// — personal repos were already fetched above.
logger.Warnf(ctx, "[Yuque] list user groups failed (treating as empty): %v", err)
groups = nil
}
for _, g := range groups {
teamRepos, err := cli.ListGroupRepos(ctx, g.Login)
if err != nil {
// Skip this group but continue others (e.g., 403 on a restricted group).
logger.Warnf(ctx, "[Yuque] skip group %s: %v", g.Login, err)
continue
}
for _, r := range teamRepos {
if _, ok := repos[r.ID]; !ok {
repos[r.ID] = r
}
}
}
}
out := make([]types.Resource, 0, len(repos))
for _, r := range repos {
out = append(out, types.Resource{
ExternalID: strconv.FormatInt(r.ID, 10),
Name: r.Name,
Type: "book",
URL: cfg.GetBaseURL() + "/" + r.Namespace,
Description: r.Namespace,
ModifiedAt: parseContentUpdatedAt(r.UpdatedAt),
Metadata: map[string]interface{}{
"public": r.Public,
"book_type": r.Type,
},
})
}
// Stable, deterministic order for UI rendering and response-body caching.
sort.Slice(out, func(i, j int) bool { return out[i].ExternalID < out[j].ExternalID })
return out, nil
}
// FetchAll performs a full sync of all books specified in resourceIDs.
func (c *Connector) FetchAll(ctx context.Context, config *types.DataSourceConfig, resourceIDs []string) ([]types.FetchedItem, error) {
items, _, err := c.walk(ctx, config, resourceIDs, nil, false)
return items, err
}
// walk is the shared implementation for FetchAll / FetchIncremental.
// If incremental is false, prev is ignored and no cursor is returned (returns nil for cursor).
func (c *Connector) walk(
ctx context.Context,
config *types.DataSourceConfig,
resourceIDs []string,
prev *yuqueCursor,
incremental bool,
) ([]types.FetchedItem, *yuqueCursor, error) {
cfg, err := parseYuqueConfig(config)
if err != nil {
return nil, nil, err
}
cli := newClient(cfg)
newCursor := &yuqueCursor{LastSyncTime: time.Now(), BookDocTimes: make(map[string]map[string]string)}
var out []types.FetchedItem
for _, bookIDStr := range resourceIDs {
bookID, err := strconv.ParseInt(bookIDStr, 10, 64)
if err != nil {
return nil, nil, fmt.Errorf("invalid book id %q: %w", bookIDStr, err)
}
docs, err := cli.ListBookDocs(ctx, bookID)
if err != nil {
return nil, nil, fmt.Errorf("list docs for book %d: %w", bookID, err)
}
currentDocs := make(map[string]bool)
newCursor.BookDocTimes[bookIDStr] = make(map[string]string)
var skippedType, skippedDraft, kept int
var sampleSkipType, sampleSkipDraft string
for _, d := range docs {
// Empty type/status is treated as acceptable — forward-compat with
// API variations that omit the field.
if d.Type != "" && d.Type != "Doc" {
skippedType++
if sampleSkipType == "" {
sampleSkipType = fmt.Sprintf("id=%d type=%q title=%q", d.ID, d.Type, d.Title)
}
continue
}
if d.Status != "" && d.Status != "1" {
skippedDraft++
if sampleSkipDraft == "" {
sampleSkipDraft = fmt.Sprintf("id=%d status=%q title=%q", d.ID, d.Status, d.Title)
}
continue
}
kept++
docIDStr := strconv.FormatInt(d.ID, 10)
currentDocs[docIDStr] = true
newCursor.BookDocTimes[bookIDStr][docIDStr] = d.ContentUpdatedAt
// Incremental: skip if content hasn't changed.
if incremental && prev != nil && prev.BookDocTimes != nil {
if prevTimes, ok := prev.BookDocTimes[bookIDStr]; ok {
if prevTimes[docIDStr] == d.ContentUpdatedAt {
continue
}
}
}
// Rate-limit: pause between GetDocDetail calls to avoid hitting
// Yuque's API rate limit (typically ~100 req/5min for personal tokens).
if err := sleepCtx(ctx, 300*time.Millisecond); err != nil {
return nil, nil, err
}
detail, err := cli.GetDocDetail(ctx, d.ID)
if err != nil {
// Record failure but continue (placeholder item with error metadata).
// Keep doc_id/book_id/slug for observability pipelines that join on these.
out = append(out, types.FetchedItem{
ExternalID: docIDStr,
Title: d.Title,
SourceResourceID: bookIDStr,
Metadata: map[string]string{
"error": err.Error(),
"channel": types.ChannelYuque,
"doc_id": docIDStr,
"book_id": bookIDStr,
"slug": d.Slug,
},
})
continue
}
// Yuque serializes `body` as Markdown when format is "markdown" or
// "lake" (Lake XML lives separately in `body_lake`) — empirically
// verified against the v2 API. Any other format (e.g. "html") may
// put non-Markdown content in `body`, so skip defensively.
if detail.Format != "" && detail.Format != "markdown" && detail.Format != "lake" {
logger.Warnf(ctx, "[Yuque] skip doc %d (%q): unsupported format %q",
d.ID, d.Title, detail.Format)
out = append(out, types.FetchedItem{
ExternalID: docIDStr,
Title: d.Title,
SourceResourceID: bookIDStr,
Metadata: map[string]string{
"channel": types.ChannelYuque,
"doc_id": docIDStr,
"book_id": bookIDStr,
"slug": d.Slug,
"skip_reason": "unsupported format: " + detail.Format,
},
})
continue
}
out = append(out, types.FetchedItem{
ExternalID: docIDStr,
Title: d.Title,
Content: []byte(detail.Body),
ContentType: "text/markdown",
FileName: sanitizeFileName(d.Title) + ".md",
URL: buildDocURL(cfg.GetBaseURL(), detail.Book.Namespace, d.Slug),
UpdatedAt: parseContentUpdatedAt(d.ContentUpdatedAt),
SourceResourceID: bookIDStr,
Metadata: map[string]string{
"doc_id": docIDStr,
"book_id": bookIDStr,
"slug": d.Slug,
"creator": strconv.FormatInt(d.UserID, 10),
"word_count": strconv.Itoa(d.WordCount),
"channel": types.ChannelYuque,
},
})
}
logger.Infof(ctx, "[Yuque] book %d: total=%d kept=%d skipped_non_doc=%d skipped_draft=%d non_doc_sample={%s} draft_sample={%s}",
bookID, len(docs), kept, skippedType, skippedDraft, sampleSkipType, sampleSkipDraft)
// Deletion detection (incremental only): previous doc IDs not in current → IsDeleted=true
if incremental && prev != nil && prev.BookDocTimes != nil {
if prevTimes, ok := prev.BookDocTimes[bookIDStr]; ok {
for prevDocID := range prevTimes {
if !currentDocs[prevDocID] {
out = append(out, types.FetchedItem{
ExternalID: prevDocID,
IsDeleted: true,
SourceResourceID: bookIDStr,
})
}
}
}
}
}
if !incremental {
return out, nil, nil
}
return out, newCursor, nil
}
// buildDocURL constructs a browser URL for a Yuque doc.
// Namespace may be empty on some responses; fall back to base URL only.
func buildDocURL(baseURL, namespace, slug string) string {
if namespace == "" {
return baseURL
}
return baseURL + "/" + namespace + "/" + slug
}
// FetchIncremental returns items changed (or deleted) since the prior cursor.
// Deletion detection: docs present in the prior cursor but absent from the
// current list are emitted as IsDeleted=true placeholder items.
func (c *Connector) FetchIncremental(
ctx context.Context,
config *types.DataSourceConfig,
cursor *types.SyncCursor,
) ([]types.FetchedItem, *types.SyncCursor, error) {
resourceIDs := config.ResourceIDs
if len(resourceIDs) == 0 {
return nil, nil, fmt.Errorf("no resource IDs (book IDs) configured")
}
// Decode prior cursor (if any).
var prev *yuqueCursor
if cursor != nil && cursor.ConnectorCursor != nil {
var p yuqueCursor
b, _ := json.Marshal(cursor.ConnectorCursor)
_ = json.Unmarshal(b, &p)
prev = &p
}
items, newCursor, err := c.walk(ctx, config, resourceIDs, prev, true)
if err != nil {
return nil, nil, err
}
// Marshal newCursor into a generic map for the SyncCursor wrapper.
cursorMap := make(map[string]interface{})
b, _ := json.Marshal(newCursor)
_ = json.Unmarshal(b, &cursorMap)
return items, &types.SyncCursor{
LastSyncTime: newCursor.LastSyncTime,
ConnectorCursor: cursorMap,
}, nil
}
@@ -0,0 +1,483 @@
package yuque
import (
"context"
"net/http"
"net/http/httptest"
"os"
"testing"
"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 makeDSConfig(f *fakeYuque, resourceIDs []string) *types.DataSourceConfig {
return &types.DataSourceConfig{
Type: types.ConnectorTypeYuque,
Credentials: map[string]interface{}{"api_token": f.cfg().APIToken, "base_url": f.cfg().BaseURL},
ResourceIDs: resourceIDs,
}
}
func TestConnector_Type(t *testing.T) {
if NewConnector().Type() != types.ConnectorTypeYuque {
t.Errorf("Type() = %q, want %q", NewConnector().Type(), types.ConnectorTypeYuque)
}
}
func TestConnector_Validate_Success(t *testing.T) {
f := newFakeYuque()
defer f.Close()
if err := NewConnector().Validate(context.Background(), makeDSConfig(f, nil)); err != nil {
t.Fatalf("Validate error: %v", err)
}
}
func TestConnector_Validate_Bad401(t *testing.T) {
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.WriteHeader(401)
}))
defer srv.Close()
err := NewConnector().Validate(context.Background(), &types.DataSourceConfig{
Credentials: map[string]interface{}{"api_token": "bad", "base_url": srv.URL},
})
if err == nil {
t.Fatal("expected error for 401")
}
}
func TestConnector_ListResources_Aggregates(t *testing.T) {
f := newFakeYuque()
defer f.Close()
// user = alice (id=1)
// Override default /api/v2/user handler to return alice (f.handleJSON panics on duplicate path registration).
f.mux = http.NewServeMux()
f.server.Config.Handler = f.mux
f.handleJSON("/api/v2/user", 200, v2UserResponse{Data: v2User{ID: 1, Login: "alice"}})
// personal repos
f.handleJSON("/api/v2/users/alice/repos", 200, v2RepoListResponse{Data: []v2Repo{
{ID: 10, Slug: "personal", Name: "Personal", Type: "Book", Namespace: "alice/personal"},
}})
// 2 groups
f.handleJSON("/api/v2/users/1/groups", 200, v2GroupListResponse{Data: []v2Group{
{ID: 100, Login: "team-a"}, {ID: 101, Login: "team-b"},
}})
f.handleJSON("/api/v2/groups/team-a/repos", 200, v2RepoListResponse{Data: []v2Repo{
{ID: 20, Slug: "ab", Name: "A Book", Type: "Book", Namespace: "team-a/ab"},
}})
f.handleJSON("/api/v2/groups/team-b/repos", 200, v2RepoListResponse{Data: []v2Repo{
{ID: 30, Slug: "bb", Name: "B Book", Type: "Book", Namespace: "team-b/bb"},
}})
resources, err := NewConnector().ListResources(context.Background(), makeDSConfig(f, nil), "")
if err != nil {
t.Fatalf("ListResources: %v", err)
}
if len(resources) != 3 {
t.Fatalf("len = %d, want 3", len(resources))
}
// Verify fields on one entry
var ab *types.Resource
for i := range resources {
if resources[i].Name == "A Book" {
ab = &resources[i]
}
}
if ab == nil {
t.Fatal("expected A Book resource")
}
if ab.ExternalID != "20" || ab.Description != "team-a/ab" || ab.Type != "book" {
t.Errorf("got %+v", ab)
}
}
func TestConnector_ListResources_DedupByID(t *testing.T) {
f := newFakeYuque()
defer f.Close()
f.mux = http.NewServeMux()
f.server.Config.Handler = f.mux
f.handleJSON("/api/v2/user", 200, v2UserResponse{Data: v2User{ID: 1, Login: "alice"}})
f.handleJSON("/api/v2/users/alice/repos", 200, v2RepoListResponse{Data: []v2Repo{
{ID: 99, Slug: "shared", Type: "Book", Namespace: "alice/shared"},
}})
f.handleJSON("/api/v2/users/1/groups", 200, v2GroupListResponse{Data: []v2Group{
{ID: 200, Login: "team-x"},
}})
// Same repo (ID 99) appears again via a team the user is in.
f.handleJSON("/api/v2/groups/team-x/repos", 200, v2RepoListResponse{Data: []v2Repo{
{ID: 99, Slug: "shared", Type: "Book", Namespace: "alice/shared"},
}})
resources, err := NewConnector().ListResources(context.Background(), makeDSConfig(f, nil), "")
if err != nil {
t.Fatalf("ListResources: %v", err)
}
if len(resources) != 1 {
t.Errorf("expected dedup to 1, got %d", len(resources))
}
}
func TestConnector_FetchAll_Markdown(t *testing.T) {
f := newFakeYuque()
defer f.Close()
// Book 7 → 2 docs; one published Doc, one draft (filtered)
f.handleJSON("/api/v2/repos/7/docs", 200, v2DocListResponse{Data: []v2Doc{
{ID: 101, Type: "Doc", Status: "1", Title: "Hello", Slug: "hello", BookID: 7, ContentUpdatedAt: "2026-04-20T10:00:00Z", WordCount: 42},
{ID: 102, Type: "Doc", Status: "0", Title: "Draft", Slug: "draft", BookID: 7, ContentUpdatedAt: "2026-04-20T11:00:00Z"}, // draft, skipped
}})
f.handleJSON("/api/v2/repos/docs/101", 200, v2DocDetailResponse{
Data: v2DocDetail{ID: 101, Title: "Hello", Body: "# Hello\n\nworld", Format: "markdown", Status: "1",
ContentUpdatedAt: "2026-04-20T10:00:00Z", Book: v2Repo{Namespace: "alice/demo"}},
})
items, err := NewConnector().FetchAll(context.Background(), makeDSConfig(f, []string{"7"}), []string{"7"})
if err != nil {
t.Fatalf("FetchAll: %v", err)
}
if len(items) != 1 {
t.Fatalf("len = %d, want 1 (draft filtered)", len(items))
}
it := items[0]
if it.ExternalID != "101" || it.Title != "Hello" {
t.Errorf("got %+v", it)
}
if string(it.Content) != "# Hello\n\nworld" {
t.Errorf("Content = %q", string(it.Content))
}
if it.ContentType != "text/markdown" {
t.Errorf("ContentType = %q", it.ContentType)
}
if it.FileName != "Hello.md" {
t.Errorf("FileName = %q", it.FileName)
}
if it.Metadata["channel"] != types.ChannelYuque {
t.Errorf("channel = %q", it.Metadata["channel"])
}
if it.UpdatedAt.IsZero() {
t.Error("UpdatedAt should be parsed")
}
}
func TestConnector_FetchAll_DocDetailError_EmitsPlaceholder(t *testing.T) {
f := newFakeYuque()
defer f.Close()
f.handleJSON("/api/v2/repos/9/docs", 200, v2DocListResponse{Data: []v2Doc{
{ID: 301, Type: "Doc", Status: "1", Title: "Broken", Slug: "broken", ContentUpdatedAt: "2026-04-20T10:00:00Z"},
{ID: 302, Type: "Doc", Status: "1", Title: "OK", Slug: "ok", ContentUpdatedAt: "2026-04-20T10:00:00Z"},
}})
f.mux.HandleFunc("/api/v2/repos/docs/301", func(w http.ResponseWriter, r *http.Request) {
w.WriteHeader(http.StatusBadRequest) // 4xx is non-retriable, bubbles up as an error
_, _ = w.Write([]byte(`{"message":"broken"}`))
})
f.handleJSON("/api/v2/repos/docs/302", 200, v2DocDetailResponse{
Data: v2DocDetail{ID: 302, Title: "OK", Body: "ok", Status: "1", ContentUpdatedAt: "2026-04-20T10:00:00Z"},
})
items, err := NewConnector().FetchAll(context.Background(), makeDSConfig(f, []string{"9"}), []string{"9"})
if err != nil {
t.Fatalf("FetchAll must not abort the batch on a single detail error, got %v", err)
}
if len(items) != 2 {
t.Fatalf("want 2 items (placeholder + ok), got %d: %+v", len(items), items)
}
// Find the placeholder — the one with the "error" metadata key.
var placeholder *types.FetchedItem
for i := range items {
if items[i].Metadata["error"] != "" {
placeholder = &items[i]
}
}
if placeholder == nil {
t.Fatal("expected a placeholder item with Metadata[error] set")
}
if placeholder.ExternalID != "301" || placeholder.Title != "Broken" {
t.Errorf("placeholder identity wrong: %+v", placeholder)
}
if placeholder.Metadata["channel"] != types.ChannelYuque {
t.Errorf("placeholder channel = %q", placeholder.Metadata["channel"])
}
if placeholder.Metadata["book_id"] != "9" || placeholder.Metadata["doc_id"] != "301" || placeholder.Metadata["slug"] != "broken" {
t.Errorf("placeholder missing traceability metadata: %+v", placeholder.Metadata)
}
if len(placeholder.Content) != 0 {
t.Errorf("placeholder should have empty Content, got %q", string(placeholder.Content))
}
}
func TestConnector_FetchAll_LakeFormatIngestedAsMarkdown(t *testing.T) {
// Yuque's v2 doc-detail serves `body` as Markdown for type="Doc" regardless
// of authoring format. format="lake" means the doc was edited in the Lake
// editor, but `body` is still the Markdown representation (Lake XML lives
// in `body_lake`). So both should be ingested identically.
f := newFakeYuque()
defer f.Close()
f.handleJSON("/api/v2/repos/12/docs", 200, v2DocListResponse{Data: []v2Doc{
{ID: 401, Type: "Doc", Status: "1", Title: "Lake Doc", Slug: "lake", ContentUpdatedAt: "2026-04-20T10:00:00Z"},
{ID: 402, Type: "Doc", Status: "1", Title: "MD Doc", Slug: "md", ContentUpdatedAt: "2026-04-20T10:00:00Z"},
}})
f.handleJSON("/api/v2/repos/docs/401", 200, v2DocDetailResponse{
Data: v2DocDetail{ID: 401, Title: "Lake Doc", Format: "lake", Body: "# lake body as markdown", Status: "1", ContentUpdatedAt: "2026-04-20T10:00:00Z"},
})
f.handleJSON("/api/v2/repos/docs/402", 200, v2DocDetailResponse{
Data: v2DocDetail{ID: 402, Title: "MD Doc", Format: "markdown", Body: "# MD", Status: "1", ContentUpdatedAt: "2026-04-20T10:00:00Z"},
})
items, err := NewConnector().FetchAll(context.Background(), makeDSConfig(f, []string{"12"}), []string{"12"})
if err != nil {
t.Fatalf("FetchAll: %v", err)
}
if len(items) != 2 {
t.Fatalf("want 2 items, got %d", len(items))
}
wantBody := map[string]string{
"401": "# lake body as markdown",
"402": "# MD",
}
for _, it := range items {
if string(it.Content) != wantBody[it.ExternalID] {
t.Errorf("item %s Content = %q, want %q", it.ExternalID, it.Content, wantBody[it.ExternalID])
}
if it.ContentType != "text/markdown" {
t.Errorf("item %s ContentType = %q, want text/markdown", it.ExternalID, it.ContentType)
}
}
}
func TestConnector_FetchAll_SkipsUnsupportedFormats(t *testing.T) {
f := newFakeYuque()
defer f.Close()
f.handleJSON("/api/v2/repos/13/docs", 200, v2DocListResponse{Data: []v2Doc{
{ID: 501, Type: "Doc", Status: "1", Title: "HTML Doc", Slug: "h", ContentUpdatedAt: "2026-04-20T10:00:00Z"},
{ID: 502, Type: "Doc", Status: "1", Title: "OK Doc", Slug: "ok", ContentUpdatedAt: "2026-04-20T10:00:00Z"},
}})
f.handleJSON("/api/v2/repos/docs/501", 200, v2DocDetailResponse{
Data: v2DocDetail{ID: 501, Title: "HTML Doc", Format: "html", Body: "<p>raw html</p>", Status: "1", ContentUpdatedAt: "2026-04-20T10:00:00Z"},
})
f.handleJSON("/api/v2/repos/docs/502", 200, v2DocDetailResponse{
Data: v2DocDetail{ID: 502, Title: "OK Doc", Format: "lake", Body: "# ok", Status: "1", ContentUpdatedAt: "2026-04-20T10:00:00Z"},
})
items, err := NewConnector().FetchAll(context.Background(), makeDSConfig(f, []string{"13"}), []string{"13"})
if err != nil {
t.Fatalf("FetchAll: %v", err)
}
if len(items) != 2 {
t.Fatalf("want 2 items (placeholder + ok), got %d", len(items))
}
var htmlPlaceholder, okItem *types.FetchedItem
for i := range items {
if items[i].ExternalID == "501" {
htmlPlaceholder = &items[i]
}
if items[i].ExternalID == "502" {
okItem = &items[i]
}
}
if htmlPlaceholder == nil || okItem == nil {
t.Fatal("missing one of the items")
}
if htmlPlaceholder.Metadata["skip_reason"] == "" {
t.Errorf("html doc should have skip_reason metadata, got %+v", htmlPlaceholder.Metadata)
}
if len(htmlPlaceholder.Content) != 0 {
t.Errorf("html placeholder should have empty Content")
}
if string(okItem.Content) != "# ok" {
t.Errorf("lake doc Content = %q, want %q", string(okItem.Content), "# ok")
}
}
func TestConnector_FetchAll_SkipsNonDocTypes(t *testing.T) {
f := newFakeYuque()
defer f.Close()
f.handleJSON("/api/v2/repos/8/docs", 200, v2DocListResponse{Data: []v2Doc{
{ID: 201, Type: "Sheet", Status: "1", Title: "S"},
{ID: 202, Type: "Board", Status: "1", Title: "B"},
{ID: 203, Type: "Thread", Status: "1", Title: "T"},
{ID: 204, Type: "Doc", Status: "1", Title: "D", ContentUpdatedAt: "2026-04-20T10:00:00Z"},
}})
f.handleJSON("/api/v2/repos/docs/204", 200, v2DocDetailResponse{
Data: v2DocDetail{ID: 204, Title: "D", Body: "text", Status: "1", ContentUpdatedAt: "2026-04-20T10:00:00Z"},
})
items, err := NewConnector().FetchAll(context.Background(), makeDSConfig(f, []string{"8"}), []string{"8"})
if err != nil {
t.Fatalf("FetchAll: %v", err)
}
if len(items) != 1 || items[0].ExternalID != "204" {
t.Errorf("expected only the Doc, got %+v", items)
}
}
func TestConnector_FetchIncremental_FirstSync(t *testing.T) {
f := newFakeYuque()
defer f.Close()
f.handleJSON("/api/v2/repos/10/docs", 200, v2DocListResponse{Data: []v2Doc{
{ID: 1, Type: "Doc", Status: "1", Title: "A", ContentUpdatedAt: "2026-04-20T10:00:00Z"},
{ID: 2, Type: "Doc", Status: "1", Title: "B", ContentUpdatedAt: "2026-04-20T11:00:00Z"},
}})
f.handleJSON("/api/v2/repos/docs/1", 200, v2DocDetailResponse{Data: v2DocDetail{ID: 1, Title: "A", Body: "a", Status: "1", ContentUpdatedAt: "2026-04-20T10:00:00Z"}})
f.handleJSON("/api/v2/repos/docs/2", 200, v2DocDetailResponse{Data: v2DocDetail{ID: 2, Title: "B", Body: "b", Status: "1", ContentUpdatedAt: "2026-04-20T11:00:00Z"}})
items, cursor, err := NewConnector().FetchIncremental(context.Background(), makeDSConfig(f, []string{"10"}), nil)
if err != nil {
t.Fatalf("FetchIncremental: %v", err)
}
if len(items) != 2 {
t.Errorf("items = %d, want 2 on first sync", len(items))
}
if cursor == nil {
t.Fatal("cursor must not be nil")
}
}
func TestConnector_FetchIncremental_NoChanges(t *testing.T) {
f := newFakeYuque()
defer f.Close()
f.handleJSON("/api/v2/repos/10/docs", 200, v2DocListResponse{Data: []v2Doc{
{ID: 1, Type: "Doc", Status: "1", Title: "A", ContentUpdatedAt: "2026-04-20T10:00:00Z"},
}})
f.handleJSON("/api/v2/repos/docs/1", 200, v2DocDetailResponse{Data: v2DocDetail{ID: 1, Title: "A", Body: "a", Status: "1", ContentUpdatedAt: "2026-04-20T10:00:00Z"}})
// First sync → establish cursor
_, cursor1, err := NewConnector().FetchIncremental(context.Background(), makeDSConfig(f, []string{"10"}), nil)
if err != nil {
t.Fatalf("first sync: %v", err)
}
// Second sync with same cursor → 0 items
items, _, err := NewConnector().FetchIncremental(context.Background(), makeDSConfig(f, []string{"10"}), cursor1)
if err != nil {
t.Fatalf("second sync: %v", err)
}
if len(items) != 0 {
t.Errorf("expected 0 changed items, got %d", len(items))
}
}
func TestConnector_FetchIncremental_ReturnsOnlyChanged(t *testing.T) {
// First sync: 2 docs, both at t0.
f1 := newFakeYuque()
f1.handleJSON("/api/v2/repos/11/docs", 200, v2DocListResponse{Data: []v2Doc{
{ID: 1, Type: "Doc", Status: "1", Title: "A", Slug: "a", ContentUpdatedAt: "2026-04-20T10:00:00Z"},
{ID: 2, Type: "Doc", Status: "1", Title: "B", Slug: "b", ContentUpdatedAt: "2026-04-20T10:00:00Z"},
}})
f1.handleJSON("/api/v2/repos/docs/1", 200, v2DocDetailResponse{Data: v2DocDetail{ID: 1, Title: "A", Body: "a", Status: "1", ContentUpdatedAt: "2026-04-20T10:00:00Z"}})
f1.handleJSON("/api/v2/repos/docs/2", 200, v2DocDetailResponse{Data: v2DocDetail{ID: 2, Title: "B", Body: "b", Status: "1", ContentUpdatedAt: "2026-04-20T10:00:00Z"}})
_, cursor1, err := NewConnector().FetchIncremental(context.Background(), makeDSConfig(f1, []string{"11"}), nil)
if err != nil {
t.Fatalf("first sync: %v", err)
}
f1.Close()
// Second sync: doc 1 unchanged, doc 2 has a newer ContentUpdatedAt.
f2 := newFakeYuque()
defer f2.Close()
f2.handleJSON("/api/v2/repos/11/docs", 200, v2DocListResponse{Data: []v2Doc{
{ID: 1, Type: "Doc", Status: "1", Title: "A", Slug: "a", ContentUpdatedAt: "2026-04-20T10:00:00Z"},
{ID: 2, Type: "Doc", Status: "1", Title: "B2", Slug: "b", ContentUpdatedAt: "2026-04-20T12:00:00Z"},
}})
// Only doc 2's detail should be fetched; if the implementation incorrectly
// fetches doc 1, the test will still pass metadata-wise — but this handler
// makes the expected single call clear.
f2.handleJSON("/api/v2/repos/docs/2", 200, v2DocDetailResponse{Data: v2DocDetail{ID: 2, Title: "B2", Body: "b2 updated", Status: "1", ContentUpdatedAt: "2026-04-20T12:00:00Z"}})
items, _, err := NewConnector().FetchIncremental(context.Background(), makeDSConfig(f2, []string{"11"}), cursor1)
if err != nil {
t.Fatalf("second sync: %v", err)
}
if len(items) != 1 {
t.Fatalf("expected exactly 1 changed item, got %d: %+v", len(items), items)
}
if items[0].ExternalID != "2" {
t.Errorf("expected changed item ExternalID=2, got %q", items[0].ExternalID)
}
if string(items[0].Content) != "b2 updated" {
t.Errorf("expected refreshed body, got %q", string(items[0].Content))
}
}
func TestConnector_FetchIncremental_DetectsDeletion(t *testing.T) {
// First sync: 2 docs
f1 := newFakeYuque()
f1.handleJSON("/api/v2/repos/10/docs", 200, v2DocListResponse{Data: []v2Doc{
{ID: 1, Type: "Doc", Status: "1", Title: "A", ContentUpdatedAt: "2026-04-20T10:00:00Z"},
{ID: 2, Type: "Doc", Status: "1", Title: "B", ContentUpdatedAt: "2026-04-20T11:00:00Z"},
}})
f1.handleJSON("/api/v2/repos/docs/1", 200, v2DocDetailResponse{Data: v2DocDetail{ID: 1, Title: "A", Body: "a", Status: "1", ContentUpdatedAt: "2026-04-20T10:00:00Z"}})
f1.handleJSON("/api/v2/repos/docs/2", 200, v2DocDetailResponse{Data: v2DocDetail{ID: 2, Title: "B", Body: "b", Status: "1", ContentUpdatedAt: "2026-04-20T11:00:00Z"}})
_, cursor1, err := NewConnector().FetchIncremental(context.Background(), makeDSConfig(f1, []string{"10"}), nil)
if err != nil {
t.Fatalf("first sync: %v", err)
}
f1.Close()
// Second sync: doc 2 removed
f2 := newFakeYuque()
defer f2.Close()
f2.handleJSON("/api/v2/repos/10/docs", 200, v2DocListResponse{Data: []v2Doc{
{ID: 1, Type: "Doc", Status: "1", Title: "A", ContentUpdatedAt: "2026-04-20T10:00:00Z"},
}})
items, _, err := NewConnector().FetchIncremental(context.Background(), makeDSConfig(f2, []string{"10"}), cursor1)
if err != nil {
t.Fatalf("second sync: %v", err)
}
deleted := 0
for _, it := range items {
if it.IsDeleted && it.ExternalID == "2" {
deleted++
}
}
if deleted != 1 {
t.Errorf("expected 1 deletion for doc 2, got %d; items=%+v", deleted, items)
}
}
func TestConnector_ListResources_ContinuesOnGroupFailure(t *testing.T) {
f := newFakeYuque()
defer f.Close()
f.mux = http.NewServeMux()
f.server.Config.Handler = f.mux
f.handleJSON("/api/v2/user", 200, v2UserResponse{Data: v2User{ID: 1, Login: "alice"}})
f.handleJSON("/api/v2/users/alice/repos", 200, v2RepoListResponse{Data: []v2Repo{
{ID: 10, Slug: "me", Name: "Mine", Type: "Book", Namespace: "alice/me"},
}})
f.handleJSON("/api/v2/users/1/groups", 200, v2GroupListResponse{Data: []v2Group{
{ID: 100, Login: "team-ok"}, {ID: 101, Login: "team-forbidden"},
}})
f.handleJSON("/api/v2/groups/team-ok/repos", 200, v2RepoListResponse{Data: []v2Repo{
{ID: 20, Slug: "ok", Name: "OK Book", Type: "Book", Namespace: "team-ok/ok"},
}})
// team-forbidden returns 403 — should be skipped, not abort the whole call.
f.mux.HandleFunc("/api/v2/groups/team-forbidden/repos", func(w http.ResponseWriter, r *http.Request) {
w.WriteHeader(http.StatusForbidden)
_, _ = w.Write([]byte(`{"message":"forbidden"}`))
})
resources, err := NewConnector().ListResources(context.Background(), makeDSConfig(f, nil), "")
if err != nil {
t.Fatalf("ListResources should not fail on per-group error, got %v", err)
}
// Expect the 2 successful repos (personal + team-ok). team-forbidden is skipped.
if len(resources) != 2 {
t.Fatalf("expected 2 resources (personal + team-ok), got %d: %+v", len(resources), resources)
}
names := map[string]bool{}
for _, r := range resources {
names[r.Name] = true
}
if !names["Mine"] || !names["OK Book"] {
t.Errorf("expected Mine + OK Book, got %v", names)
}
}
@@ -0,0 +1,265 @@
// Package yuque implements the Yuque (语雀) data source connector for WeKnora.
//
// It syncs documents from personal and group knowledge bases (books/repos) into WeKnora
// knowledge bases, preserving Markdown formatting.
//
// Yuque API docs:
// - Authentication: X-Auth-Token header, personal token from https://www.yuque.com/settings/tokens
// - User: GET /api/v2/user
// - Groups: GET /api/v2/users/{id}/groups
// - Repos: GET /api/v2/users/{login}/repos, GET /api/v2/groups/{login}/repos
// - Docs: GET /api/v2/repos/{book_id}/docs (list), GET /api/v2/repos/docs/{id} (detail)
//
// Known limitations (v1):
// - Only syncs type=Doc (Sheet/Thread/Board/Table skipped)
// - Only syncs status="1" (published), drafts skipped
// - Private-book images (CDN URLs with auth) may fail to load
// - Lake editor may leave non-standard markdown (anchors, sized image attrs)
package yuque
import (
"encoding/json"
"fmt"
"strconv"
"strings"
"time"
"unicode/utf8"
"github.com/Tencent/WeKnora/internal/datasource"
"github.com/Tencent/WeKnora/internal/types"
)
// DefaultBaseURL is the Yuque public cloud API base URL.
const DefaultBaseURL = "https://www.yuque.com"
// Config holds Yuque-specific configuration.
type Config struct {
// APIToken is a personal token from Yuque settings → tokens.
APIToken string `json:"api_token"`
// BaseURL is the Yuque deployment base URL (default: https://www.yuque.com).
// For enterprise/private deployments, use the company's yuque domain.
BaseURL string `json:"base_url,omitempty"`
}
// GetBaseURL returns the normalized base URL:
// - empty → DefaultBaseURL
// - missing scheme → prepend "https://"
// - trailing slash → stripped
func (c *Config) GetBaseURL() string {
url := strings.TrimSpace(c.BaseURL)
if url == "" {
return DefaultBaseURL
}
if !strings.Contains(url, "://") {
url = "https://" + url
}
url = strings.TrimRight(url, "/")
return url
}
// parseYuqueConfig extracts and validates Yuque-specific configuration.
// Uses JSON marshal/unmarshal roundtrip (consistent with Feishu's parseFeishuConfig)
// rather than single-field type assertion, because we have multiple fields with
// optional defaults.
func parseYuqueConfig(config *types.DataSourceConfig) (*Config, error) {
if config == nil {
return nil, fmt.Errorf("%w: config is nil", datasource.ErrInvalidConfig)
}
credBytes, err := json.Marshal(config.Credentials)
if err != nil {
return nil, fmt.Errorf("marshal credentials: %w", err)
}
var cfg Config
if err := json.Unmarshal(credBytes, &cfg); err != nil {
return nil, fmt.Errorf("parse yuque credentials: %w", err)
}
if strings.TrimSpace(cfg.APIToken) == "" {
return nil, fmt.Errorf("%w: api_token is required", datasource.ErrInvalidCredentials)
}
if err := datasource.ValidateConnectorBaseURL(cfg.GetBaseURL()); err != nil {
return nil, err
}
return &cfg, nil
}
// --- Yuque API response types ---
// flexibleStatus accepts either a string ("1") or a number (1) for the doc
// `status` field. Yuque's OpenAPI spec declares `status` as string, but the
// runtime API returns it as an integer, so unmarshaling into a plain string
// fails with "cannot unmarshal number into Go struct field ... of type string".
// Normalizing to the textual form lets existing comparisons (e.g. != "1") keep
// working for both response shapes.
type flexibleStatus string
func (s *flexibleStatus) UnmarshalJSON(b []byte) error {
if string(b) == "null" {
*s = ""
return nil
}
if len(b) > 0 && b[0] == '"' {
var str string
if err := json.Unmarshal(b, &str); err != nil {
return err
}
*s = flexibleStatus(str)
return nil
}
// Integer form. We decode to int64 so that floats, booleans, arrays, and
// objects fail loudly instead of being silently stringified — if Yuque
// changes the shape again, we'd rather surface a clear error than feed
// garbage to the Status == "1" comparison.
var i int64
if err := json.Unmarshal(b, &i); err != nil {
return fmt.Errorf("flexibleStatus: expected string or integer, got %s: %w", b, err)
}
*s = flexibleStatus(strconv.FormatInt(i, 10))
return nil
}
// apiErrorBody is the error body shape Yuque sometimes returns on non-2xx.
type apiErrorBody struct {
Message string `json:"message"`
Status int `json:"status"`
}
// v2UserResponse wraps GET /api/v2/user.
type v2UserResponse struct {
Data v2User `json:"data"`
}
type v2User struct {
ID int64 `json:"id"`
Type string `json:"type"` // "User" for personal token, "Group" for team token
Login string `json:"login"`
Name string `json:"name"`
}
// v2GroupListResponse wraps GET /api/v2/users/{id}/groups.
type v2GroupListResponse struct {
Data []v2Group `json:"data"`
}
type v2Group struct {
ID int64 `json:"id"`
Login string `json:"login"`
Name string `json:"name"`
}
// v2RepoListResponse wraps GET /api/v2/users/{login}/repos and /groups/{login}/repos.
type v2RepoListResponse struct {
Data []v2Repo `json:"data"`
}
type v2Repo struct {
ID int64 `json:"id"`
Type string `json:"type"` // "Book" | "Design" (listing filter enum; connector requests type=Book)
Slug string `json:"slug"`
Name string `json:"name"`
UserID int64 `json:"user_id"`
Namespace string `json:"namespace"` // e.g. "group_login/book_slug"
Public int `json:"public"` // 0:private, 1:public, 2:internal
Description string `json:"description"`
UpdatedAt string `json:"updated_at"` // RFC3339 string
}
// v2DocListResponse wraps GET /api/v2/repos/{book_id}/docs.
type v2DocListResponse struct {
Meta struct {
Total int `json:"total"`
} `json:"meta"`
Data []v2Doc `json:"data"`
}
// v2Doc is the document summary returned by the list endpoint (no body).
type v2Doc struct {
ID int64 `json:"id"`
Type string `json:"type"` // Doc / Sheet / Thread / Board / Table
Slug string `json:"slug"`
Title string `json:"title"`
BookID int64 `json:"book_id"`
UserID int64 `json:"user_id"`
Status flexibleStatus `json:"status"` // "0" draft, "1" published — API may return int or string
ContentUpdatedAt string `json:"content_updated_at"` // RFC3339 string — use for change detection
UpdatedAt string `json:"updated_at"`
WordCount int `json:"word_count"`
}
// v2DocDetailResponse wraps GET /api/v2/repos/docs/{id}.
type v2DocDetailResponse struct {
Data v2DocDetail `json:"data"`
}
type v2DocDetail struct {
ID int64 `json:"id"`
Type string `json:"type"`
Slug string `json:"slug"`
Title string `json:"title"`
BookID int64 `json:"book_id"`
Format string `json:"format"` // "markdown" / "lake" / "html"
Body string `json:"body"` // Markdown content
Status flexibleStatus `json:"status"` // API may return int or string
ContentUpdatedAt string `json:"content_updated_at"`
UpdatedAt string `json:"updated_at"`
WordCount int `json:"word_count"`
Book v2Repo `json:"book"`
}
// yuqueCursor stores incremental sync state.
// Key1: book_id (string), Key2: doc_id (string), Value: content_updated_at (raw RFC3339 string)
type yuqueCursor struct {
LastSyncTime time.Time `json:"last_sync_time"`
BookDocTimes map[string]map[string]string `json:"book_doc_times,omitempty"`
}
// parseContentUpdatedAt parses Yuque ISO 8601 timestamp (returns zero time on parse failure).
func parseContentUpdatedAt(ts string) time.Time {
if ts == "" {
return time.Time{}
}
t, err := time.Parse(time.RFC3339, ts)
if err != nil {
return time.Time{}
}
return t
}
// sanitizeFileName removes characters that are invalid in filenames and
// truncates to a safe length at a UTF-8 rune boundary. Raw byte truncation
// would split a multi-byte codepoint (Chinese characters are 3 bytes in UTF-8)
// and produce an invalid UTF-8 string, which downstream filename validation
// (utf8.ValidString) rejects with "文件名包含非法字符".
func sanitizeFileName(name string) string {
if name == "" {
return "untitled"
}
replacer := strings.NewReplacer(
"/", "_", "\\", "_", ":", "_", "*", "_",
"?", "_", "\"", "_", "<", "_", ">", "_", "|", "_",
)
result := replacer.Replace(name)
const maxBytes = 200
if len(result) > maxBytes {
result = result[:maxBytes]
// Peel any trailing bytes that no longer form a complete rune.
// DecodeLastRuneInString returns (RuneError, 1) for a partial UTF-8
// sequence — at most 3 bytes need peeling since a rune is ≤ 4 bytes.
for len(result) > 0 {
r, size := utf8.DecodeLastRuneInString(result)
if r != utf8.RuneError || size != 1 {
break
}
result = result[:len(result)-1]
}
}
return result
}
// redactToken returns a masked form of the token for logging (never log the full token).
func redactToken(t string) string {
if len(t) < 12 {
return "***"
}
return t[:6] + "..." + t[len(t)-4:]
}
@@ -0,0 +1,176 @@
package yuque
import (
"encoding/json"
"strings"
"testing"
"unicode/utf8"
"github.com/Tencent/WeKnora/internal/types"
)
// Yuque's runtime API returns the doc `status` field as an integer (0 or 1),
// even though its OpenAPI spec declares it as string. We tolerate both forms.
func TestV2Doc_Status_AcceptsNumberAndString(t *testing.T) {
cases := []struct {
name string
body string
want flexibleStatus
}{
{"number 1", `{"id":1,"status":1}`, "1"},
{"number 0", `{"id":1,"status":0}`, "0"},
{"string 1", `{"id":1,"status":"1"}`, "1"},
{"null", `{"id":1,"status":null}`, ""},
{"missing", `{"id":1}`, ""},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
var d v2Doc
if err := json.Unmarshal([]byte(tc.body), &d); err != nil {
t.Fatalf("unmarshal: %v", err)
}
if d.Status != tc.want {
t.Errorf("Status = %q, want %q", d.Status, tc.want)
}
})
}
}
func TestV2DocDetail_Status_AcceptsNumber(t *testing.T) {
var d v2DocDetail
if err := json.Unmarshal([]byte(`{"id":1,"status":1}`), &d); err != nil {
t.Fatalf("unmarshal: %v", err)
}
if d.Status != "1" {
t.Errorf("Status = %q, want %q", d.Status, "1")
}
}
func TestV2DocListResponse_Status_Number(t *testing.T) {
// Reproduces the real-world failure:
// decode response: json: cannot unmarshal number into Go struct field v2Doc.data.status of type string
body := `{"meta":{"total":1},"data":[{"id":101,"type":"Doc","status":1,"title":"hello"}]}`
var resp v2DocListResponse
if err := json.Unmarshal([]byte(body), &resp); err != nil {
t.Fatalf("unmarshal: %v", err)
}
if len(resp.Data) != 1 {
t.Fatalf("len(Data) = %d", len(resp.Data))
}
if resp.Data[0].Status != "1" {
t.Errorf("Status = %q, want %q", resp.Data[0].Status, "1")
}
}
// flexibleStatus deliberately rejects unexpected JSON shapes so a future Yuque
// schema drift (e.g. status becoming a float, bool, array) fails loudly rather
// than being silently stringified and bypassing the `!= "1"` filter.
func TestV2Doc_Status_RejectsUnexpectedShapes(t *testing.T) {
cases := []string{
`{"id":1,"status":true}`,
`{"id":1,"status":1.5}`,
`{"id":1,"status":[1]}`,
`{"id":1,"status":{"x":1}}`,
}
for _, body := range cases {
t.Run(body, func(t *testing.T) {
var d v2Doc
if err := json.Unmarshal([]byte(body), &d); err == nil {
t.Errorf("expected error for %s, got Status=%q", body, d.Status)
}
})
}
}
func TestParseYuqueConfig(t *testing.T) {
t.Run("valid full", func(t *testing.T) {
cfg, err := parseYuqueConfig(&types.DataSourceConfig{
Credentials: map[string]interface{}{
"api_token": "tok-123",
"base_url": "https://company.yuque.com",
},
})
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if cfg.APIToken != "tok-123" {
t.Errorf("APIToken = %q, want tok-123", cfg.APIToken)
}
if cfg.BaseURL != "https://company.yuque.com" {
t.Errorf("BaseURL = %q", cfg.BaseURL)
}
})
t.Run("default base_url", func(t *testing.T) {
cfg, err := parseYuqueConfig(&types.DataSourceConfig{
Credentials: map[string]interface{}{"api_token": "tok-abc"},
})
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if cfg.GetBaseURL() != "https://www.yuque.com" {
t.Errorf("GetBaseURL default = %q, want https://www.yuque.com", cfg.GetBaseURL())
}
})
t.Run("missing api_token", func(t *testing.T) {
_, err := parseYuqueConfig(&types.DataSourceConfig{
Credentials: map[string]interface{}{},
})
if err == nil {
t.Fatal("expected error for missing api_token")
}
})
t.Run("nil config", func(t *testing.T) {
_, err := parseYuqueConfig(nil)
if err == nil {
t.Fatal("expected error for nil config")
}
})
}
func TestGetBaseURL_TrimsTrailingSlash(t *testing.T) {
c := &Config{BaseURL: "https://x.yuque.com/"}
if got := c.GetBaseURL(); got != "https://x.yuque.com" {
t.Errorf("GetBaseURL = %q, want https://x.yuque.com", got)
}
}
func TestGetBaseURL_AddsSchemeIfMissing(t *testing.T) {
c := &Config{BaseURL: "company.yuque.com"}
if got := c.GetBaseURL(); got != "https://company.yuque.com" {
t.Errorf("GetBaseURL = %q, want https://company.yuque.com", got)
}
}
func TestGetBaseURL_EmptyUsesDefault(t *testing.T) {
c := &Config{BaseURL: ""}
if got := c.GetBaseURL(); got != "https://www.yuque.com" {
t.Errorf("GetBaseURL = %q, want default", got)
}
}
func TestGetBaseURL_TrimsWhitespace(t *testing.T) {
c := &Config{BaseURL: " https://x.yuque.com/ "}
if got := c.GetBaseURL(); got != "https://x.yuque.com" {
t.Errorf("GetBaseURL = %q, want https://x.yuque.com", got)
}
}
func TestSanitizeFileName_TruncatesAtRuneBoundary(t *testing.T) {
// Long Chinese title (each 测 is 3 bytes in UTF-8). Raw byte slicing at 200
// would split a rune and produce invalid UTF-8 that downstream filename
// validation rejects with "文件名包含非法字符".
long := strings.Repeat("测试", 100) // 600 bytes
got := sanitizeFileName(long)
if !utf8.ValidString(got) {
t.Fatalf("sanitizeFileName produced invalid UTF-8: %q", got)
}
if len(got) > 200 {
t.Errorf("len = %d, want ≤ 200", len(got))
}
if len(got) == 0 {
t.Error("result is empty")
}
}
+17
View File
@@ -0,0 +1,17 @@
package datasource
import (
"testing"
"github.com/Tencent/WeKnora/internal/types"
)
func TestFeishuMetadataDoesNotAdvertiseWebhook(t *testing.T) {
meta := ConnectorMetadataRegistry[types.ConnectorTypeFeishu]
for _, capability := range meta.Capabilities {
if capability == "webhook" {
t.Fatalf("Feishu connector should not advertise webhook until webhook sync is implemented")
}
}
}
+50
View File
@@ -0,0 +1,50 @@
package datasource
import (
"errors"
"fmt"
"strings"
)
// Error definitions for datasource operations
var (
// Connector errors
ErrConnectorNil = errors.New("connector is nil")
ErrConnectorTypeEmpty = errors.New("connector type is empty")
ErrConnectorNotFound = errors.New("connector type not found in registry")
// DataSource errors
ErrDataSourceNotFound = errors.New("data source not found")
ErrDataSourceInvalid = errors.New("data source configuration is invalid")
ErrDataSourceNotActive = errors.New("data source is not active")
// Configuration errors
ErrInvalidConfig = errors.New("invalid configuration")
ErrInvalidCredentials = errors.New("invalid credentials")
// Sync errors
ErrSyncFailed = errors.New("sync operation failed")
ErrSyncCanceled = errors.New("sync operation was canceled")
ErrFetchFailed = errors.New("failed to fetch items from source")
ErrResourceNotFound = errors.New("resource not found in source system")
// Knowledge base errors
ErrKnowledgeBaseNotFound = errors.New("knowledge base not found")
// Sync log errors
ErrSyncLogNotFound = errors.New("sync log not found")
)
// PartialFetchError indicates that some resources were fetched successfully but
// others failed. The caller should process Items (if any), persist an updated
// cursor when provided, and surface Details to the user as a partial sync.
type PartialFetchError struct {
Details []string
}
func (e *PartialFetchError) Error() string {
if e == nil || len(e.Details) == 0 {
return "partial fetch: some resources failed"
}
return fmt.Sprintf("partial fetch: %s", strings.Join(e.Details, "; "))
}
+33
View File
@@ -0,0 +1,33 @@
package datasource
import (
"fmt"
"net/http"
"strings"
"time"
"github.com/Tencent/WeKnora/internal/utils"
)
// ValidateConnectorBaseURL checks a connector API base URL against the SSRF policy.
// Empty rawURL is allowed; callers apply their own default before issuing requests.
func ValidateConnectorBaseURL(rawURL string) error {
url := strings.TrimSpace(rawURL)
if url == "" {
return nil
}
if !strings.Contains(url, "://") {
url = "https://" + url
}
if err := utils.ValidateURLForSSRF(url); err != nil {
return fmt.Errorf("base_url SSRF validation failed: %w", err)
}
return nil
}
// NewConnectorHTTPClient returns an HTTP client with redirect and dial-time SSRF guards.
func NewConnectorHTTPClient(timeout time.Duration) *http.Client {
cfg := utils.DefaultSSRFSafeHTTPClientConfig()
cfg.Timeout = timeout
return utils.NewSSRFSafeHTTPClient(cfg)
}
+26
View File
@@ -0,0 +1,26 @@
package datasource
import (
"testing"
secutils "github.com/Tencent/WeKnora/internal/utils"
)
func TestValidateConnectorBaseURLBlocksLoopback(t *testing.T) {
secutils.ResetSSRFWhitelistForTest()
t.Cleanup(secutils.ResetSSRFWhitelistForTest)
err := ValidateConnectorBaseURL("http://127.0.0.1:8000")
if err == nil {
t.Fatal("expected loopback base_url to be rejected")
}
}
func TestValidateConnectorBaseURLAllowsPublicHTTPS(t *testing.T) {
secutils.ResetSSRFWhitelistForTest()
t.Cleanup(secutils.ResetSSRFWhitelistForTest)
if err := ValidateConnectorBaseURL("https://open.feishu.cn"); err != nil {
t.Fatalf("expected public base_url to pass: %v", err)
}
}
+213
View File
@@ -0,0 +1,213 @@
package datasource
import (
"context"
"encoding/json"
"fmt"
"sync"
"time"
"github.com/Tencent/WeKnora/internal/logger"
"github.com/Tencent/WeKnora/internal/tracing/langfuse"
"github.com/Tencent/WeKnora/internal/types"
"github.com/Tencent/WeKnora/internal/types/interfaces"
"github.com/hibiken/asynq"
"github.com/robfig/cron/v3"
)
// Scheduler manages cron-based periodic sync for data sources.
//
// robfig/cron fires at absolute wall-clock times (e.g. "0 0 * * * *" always fires
// at the top of every hour regardless of when the process started). So multiple
// instances will fire at the same moment. Dedup is handled by two layers:
//
// 1. HasRunningSync — if a previous sync is still running, skip (prevent overlap).
// 2. asynq.TaskID — deterministic ID per (dataSourceID, minute). Redis ensures
// only one task with a given ID is enqueued. Losers get ErrTaskIDConflict.
type Scheduler struct {
cron *cron.Cron
dsRepo interfaces.DataSourceRepository
syncLogRepo interfaces.SyncLogRepository
taskEnqueuer interfaces.TaskEnqueuer
mu sync.Mutex
entries map[string]cron.EntryID // dataSourceID → cron entry ID
}
// NewScheduler creates a new Scheduler.
func NewScheduler(
dsRepo interfaces.DataSourceRepository,
syncLogRepo interfaces.SyncLogRepository,
taskEnqueuer interfaces.TaskEnqueuer,
) *Scheduler {
return &Scheduler{
cron: cron.New(cron.WithSeconds(), cron.WithChain(
cron.Recover(cron.DefaultLogger),
)),
dsRepo: dsRepo,
syncLogRepo: syncLogRepo,
taskEnqueuer: taskEnqueuer,
entries: make(map[string]cron.EntryID),
}
}
// Start loads all active data sources from the database and registers their
// cron schedules. Then starts the cron runner in the background.
func (s *Scheduler) Start(ctx context.Context) error {
dataSources, err := s.dsRepo.FindActive(ctx)
if err != nil {
return fmt.Errorf("load active data sources: %w", err)
}
for _, ds := range dataSources {
if ds.SyncSchedule == "" {
continue
}
if err := s.addEntry(ds); err != nil {
logger.Warnf(ctx, "[Scheduler] failed to register cron for ds=%s schedule=%q: %v",
ds.ID, ds.SyncSchedule, err)
}
}
s.cron.Start()
logger.Infof(ctx, "[Scheduler] started with %d cron entries", len(s.entries))
return nil
}
// Stop gracefully stops the cron runner and waits for running jobs to finish.
func (s *Scheduler) Stop() {
ctx := s.cron.Stop()
<-ctx.Done()
}
// AddOrUpdate registers (or re-registers) a cron entry for the given data source.
func (s *Scheduler) AddOrUpdate(ds *types.DataSource) error {
s.mu.Lock()
defer s.mu.Unlock()
if entryID, ok := s.entries[ds.ID]; ok {
s.cron.Remove(entryID)
delete(s.entries, ds.ID)
}
if ds.Status != types.DataSourceStatusActive || ds.SyncSchedule == "" {
return nil
}
return s.addEntryLocked(ds)
}
// Remove removes the cron entry for a data source.
func (s *Scheduler) Remove(dataSourceID string) {
s.mu.Lock()
defer s.mu.Unlock()
if entryID, ok := s.entries[dataSourceID]; ok {
s.cron.Remove(entryID)
delete(s.entries, dataSourceID)
}
}
func (s *Scheduler) addEntry(ds *types.DataSource) error {
s.mu.Lock()
defer s.mu.Unlock()
return s.addEntryLocked(ds)
}
func (s *Scheduler) addEntryLocked(ds *types.DataSource) error {
dsID := ds.ID
tenantID := ds.TenantID
entryID, err := s.cron.AddFunc(ds.SyncSchedule, func() {
s.triggerSync(dsID, tenantID)
})
if err != nil {
return fmt.Errorf("invalid cron expression %q: %w", ds.SyncSchedule, err)
}
s.entries[dsID] = entryID
return nil
}
// triggerSync is called by the cron runner on each tick.
//
// Layer 1 — DB: if a previous sync is still running, skip. This prevents
// overlap when a sync takes longer than the cron interval.
//
// Layer 2 — Redis: deterministic asynq.TaskID = "dssync:<dsID>:<minute>".
// Since robfig/cron fires at absolute wall-clock times, all instances trigger
// at the same minute. The first Enqueue wins; others get ErrTaskIDConflict.
func (s *Scheduler) triggerSync(dataSourceID string, tenantID uint64) {
ctx := context.Background()
ds, err := s.dsRepo.FindByID(ctx, dataSourceID)
if err != nil || ds == nil || ds.Status != types.DataSourceStatusActive {
logger.Infof(ctx, "[Scheduler] skipping sync for ds=%s (not active or not found)", dataSourceID)
return
}
// Layer 1: prevent overlap with a still-running sync
if running, _ := s.syncLogRepo.HasRunningSync(ctx, dataSourceID); running {
logger.Infof(ctx, "[Scheduler] skipping sync for ds=%s (previous sync still running)", dataSourceID)
return
}
syncLog := &types.SyncLog{
DataSourceID: dataSourceID,
TenantID: tenantID,
Status: types.SyncLogStatusRunning,
StartedAt: time.Now().UTC(),
}
if err := s.syncLogRepo.Create(ctx, syncLog); err != nil {
logger.Errorf(ctx, "[Scheduler] failed to create sync log for ds=%s: %v", dataSourceID, err)
return
}
payload := &types.DataSourceSyncPayload{
DataSourceID: dataSourceID,
TenantID: tenantID,
SyncLogID: syncLog.ID,
ForceFull: false,
Trigger: "schedule",
}
langfuse.InjectTracing(ctx, payload)
payloadJSON, _ := json.Marshal(payload)
task := asynq.NewTask(types.TypeDataSourceSync, payloadJSON)
// Layer 2: deterministic TaskID — all instances in the same minute produce the same ID
taskID := fmt.Sprintf("dssync:%s:%s", dataSourceID, time.Now().UTC().Truncate(time.Minute).Format("200601021504"))
_, err = s.taskEnqueuer.Enqueue(task,
asynq.Queue(types.QueueSync),
asynq.MaxRetry(5),
asynq.Timeout(2*time.Hour),
asynq.TaskID(taskID),
)
if err != nil {
if err == asynq.ErrTaskIDConflict {
logger.Infof(ctx, "[Scheduler] sync already enqueued by another instance for ds=%s", dataSourceID)
syncLog.Status = types.SyncLogStatusCanceled
now := time.Now().UTC()
syncLog.FinishedAt = &now
syncLog.ErrorMessage = "deduplicated: another instance enqueued first"
_ = s.syncLogRepo.Update(ctx, syncLog)
return
}
logger.Errorf(ctx, "[Scheduler] failed to enqueue sync task for ds=%s: %v", dataSourceID, err)
syncLog.Status = types.SyncLogStatusFailed
now := time.Now().UTC()
syncLog.FinishedAt = &now
syncLog.ErrorMessage = fmt.Sprintf("enqueue failed: %v", err)
_ = s.syncLogRepo.Update(ctx, syncLog)
return
}
logger.Infof(ctx, "[Scheduler] sync task enqueued for ds=%s syncLog=%s", dataSourceID, syncLog.ID)
}
// EntryCount returns the number of active cron entries (for testing/monitoring).
func (s *Scheduler) EntryCount() int {
s.mu.Lock()
defer s.mu.Unlock()
return len(s.entries)
}
+379
View File
@@ -0,0 +1,379 @@
package datasource
import (
"context"
"sync"
"sync/atomic"
"testing"
"time"
"github.com/Tencent/WeKnora/internal/types"
"github.com/hibiken/asynq"
)
// ──────────────────────────────────────────────────────────────────────
// Fake implementations for testing
// ──────────────────────────────────────────────────────────────────────
// fakeDataSourceRepo is an in-memory DataSourceRepository.
type fakeDataSourceRepo struct {
mu sync.Mutex
dataSources map[string]*types.DataSource
}
func newFakeDataSourceRepo() *fakeDataSourceRepo {
return &fakeDataSourceRepo{dataSources: make(map[string]*types.DataSource)}
}
func (r *fakeDataSourceRepo) Create(_ context.Context, ds *types.DataSource) error {
r.mu.Lock()
defer r.mu.Unlock()
r.dataSources[ds.ID] = ds
return nil
}
func (r *fakeDataSourceRepo) FindByID(_ context.Context, id string) (*types.DataSource, error) {
r.mu.Lock()
defer r.mu.Unlock()
ds, ok := r.dataSources[id]
if !ok {
return nil, ErrDataSourceNotFound
}
return ds, nil
}
func (r *fakeDataSourceRepo) FindByKnowledgeBase(_ context.Context, kbID string) ([]*types.DataSource, error) {
return nil, nil
}
func (r *fakeDataSourceRepo) Update(_ context.Context, ds *types.DataSource) error {
r.mu.Lock()
defer r.mu.Unlock()
r.dataSources[ds.ID] = ds
return nil
}
func (r *fakeDataSourceRepo) UpdateSyncState(ctx context.Context, ds *types.DataSource) error {
return r.Update(ctx, ds)
}
func (r *fakeDataSourceRepo) Delete(_ context.Context, id string) error {
r.mu.Lock()
defer r.mu.Unlock()
delete(r.dataSources, id)
return nil
}
func (r *fakeDataSourceRepo) FindActive(_ context.Context) ([]*types.DataSource, error) {
r.mu.Lock()
defer r.mu.Unlock()
var result []*types.DataSource
for _, ds := range r.dataSources {
if ds.Status == types.DataSourceStatusActive && ds.SyncSchedule != "" {
result = append(result, ds)
}
}
return result, nil
}
// fakeSyncLogRepo is an in-memory SyncLogRepository.
type fakeSyncLogRepo struct {
mu sync.Mutex
logs map[string]*types.SyncLog
}
func newFakeSyncLogRepo() *fakeSyncLogRepo {
return &fakeSyncLogRepo{logs: make(map[string]*types.SyncLog)}
}
func (r *fakeSyncLogRepo) Create(_ context.Context, log *types.SyncLog) error {
r.mu.Lock()
defer r.mu.Unlock()
if log.ID == "" {
log.ID = "log-" + time.Now().Format("150405.000")
}
r.logs[log.ID] = log
return nil
}
func (r *fakeSyncLogRepo) FindByID(_ context.Context, id string) (*types.SyncLog, error) {
r.mu.Lock()
defer r.mu.Unlock()
l, ok := r.logs[id]
if !ok {
return nil, ErrSyncLogNotFound
}
return l, nil
}
func (r *fakeSyncLogRepo) FindByDataSource(_ context.Context, dsID string, limit, offset int) ([]*types.SyncLog, error) {
return nil, nil
}
func (r *fakeSyncLogRepo) FindLatest(_ context.Context, dsID string) (*types.SyncLog, error) {
return nil, nil
}
func (r *fakeSyncLogRepo) Update(_ context.Context, log *types.SyncLog) error {
r.mu.Lock()
defer r.mu.Unlock()
r.logs[log.ID] = log
return nil
}
func (r *fakeSyncLogRepo) UpdateResult(ctx context.Context, log *types.SyncLog) error {
return r.Update(ctx, log)
}
func (r *fakeSyncLogRepo) CancelPendingByDataSource(_ context.Context, dsID string) error {
return nil
}
func (r *fakeSyncLogRepo) CleanupOldLogs(_ context.Context, retentionDays int) error {
return nil
}
func (r *fakeSyncLogRepo) HasRunningSync(_ context.Context, dsID string) (bool, error) {
r.mu.Lock()
defer r.mu.Unlock()
for _, log := range r.logs {
if log.DataSourceID == dsID && log.Status == types.SyncLogStatusRunning {
return true, nil
}
}
return false, nil
}
// fakeTaskEnqueuer counts how many tasks are enqueued.
type fakeTaskEnqueuer struct {
count atomic.Int64
lastQueue atomic.Value
}
func (e *fakeTaskEnqueuer) Enqueue(task *asynq.Task, opts ...asynq.Option) (*asynq.TaskInfo, error) {
e.count.Add(1)
for _, opt := range opts {
if opt.Type() == asynq.QueueOpt {
if queue, ok := opt.Value().(string); ok {
e.lastQueue.Store(queue)
}
}
}
return &asynq.TaskInfo{ID: "task-fake"}, nil
}
// ──────────────────────────────────────────────────────────────────────
// Tests
// ──────────────────────────────────────────────────────────────────────
func TestScheduler_StartWithActiveDataSources(t *testing.T) {
repo := newFakeDataSourceRepo()
_ = repo.Create(context.Background(), &types.DataSource{
ID: "ds-1",
TenantID: 1,
Status: types.DataSourceStatusActive,
SyncSchedule: "*/2 * * * * *", // every 2 seconds (6-field cron with seconds)
})
_ = repo.Create(context.Background(), &types.DataSource{
ID: "ds-2",
TenantID: 1,
Status: types.DataSourceStatusPaused, // should NOT be scheduled
SyncSchedule: "*/2 * * * * *",
})
enqueuer := &fakeTaskEnqueuer{}
scheduler := NewScheduler(repo, newFakeSyncLogRepo(), enqueuer)
if err := scheduler.Start(context.Background()); err != nil {
t.Fatalf("Start() error: %v", err)
}
defer scheduler.Stop()
// Only ds-1 should be registered (ds-2 is paused, not returned by FindActive)
if scheduler.EntryCount() != 1 {
t.Errorf("EntryCount() = %d, want 1", scheduler.EntryCount())
}
}
func TestScheduler_CronFires(t *testing.T) {
repo := newFakeDataSourceRepo()
_ = repo.Create(context.Background(), &types.DataSource{
ID: "ds-fire",
TenantID: 1,
Status: types.DataSourceStatusActive,
SyncSchedule: "* * * * * *", // every second
})
enqueuer := &fakeTaskEnqueuer{}
scheduler := NewScheduler(repo, newFakeSyncLogRepo(), enqueuer)
if err := scheduler.Start(context.Background()); err != nil {
t.Fatalf("Start() error: %v", err)
}
// Wait for at least one tick
time.Sleep(2500 * time.Millisecond)
scheduler.Stop()
if enqueuer.count.Load() == 0 {
t.Error("expected at least 1 enqueue, got 0")
}
if queue, _ := enqueuer.lastQueue.Load().(string); queue != types.QueueSync {
t.Errorf("scheduled sync queue = %q, want %q", queue, types.QueueSync)
}
}
func TestScheduler_AddOrUpdate(t *testing.T) {
enqueuer := &fakeTaskEnqueuer{}
scheduler := NewScheduler(newFakeDataSourceRepo(), newFakeSyncLogRepo(), enqueuer)
scheduler.cron.Start()
defer scheduler.Stop()
ds := &types.DataSource{
ID: "ds-new",
TenantID: 1,
Status: types.DataSourceStatusActive,
SyncSchedule: "0 0 * * * *", // every hour
}
// Add
if err := scheduler.AddOrUpdate(ds); err != nil {
t.Fatalf("AddOrUpdate() error: %v", err)
}
if scheduler.EntryCount() != 1 {
t.Errorf("EntryCount() = %d, want 1", scheduler.EntryCount())
}
// Update schedule
ds.SyncSchedule = "0 30 * * * *" // every half hour
if err := scheduler.AddOrUpdate(ds); err != nil {
t.Fatalf("AddOrUpdate() (update) error: %v", err)
}
if scheduler.EntryCount() != 1 {
t.Errorf("after update: EntryCount() = %d, want 1 (should replace, not add)", scheduler.EntryCount())
}
}
func TestScheduler_AddOrUpdate_PausedIsNoop(t *testing.T) {
enqueuer := &fakeTaskEnqueuer{}
scheduler := NewScheduler(newFakeDataSourceRepo(), newFakeSyncLogRepo(), enqueuer)
scheduler.cron.Start()
defer scheduler.Stop()
ds := &types.DataSource{
ID: "ds-paused",
TenantID: 1,
Status: types.DataSourceStatusPaused,
SyncSchedule: "0 0 * * * *",
}
if err := scheduler.AddOrUpdate(ds); err != nil {
t.Fatalf("AddOrUpdate() error: %v", err)
}
if scheduler.EntryCount() != 0 {
t.Errorf("paused ds should not be scheduled, EntryCount() = %d", scheduler.EntryCount())
}
}
func TestScheduler_AddOrUpdate_EmptyScheduleIsNoop(t *testing.T) {
enqueuer := &fakeTaskEnqueuer{}
scheduler := NewScheduler(newFakeDataSourceRepo(), newFakeSyncLogRepo(), enqueuer)
scheduler.cron.Start()
defer scheduler.Stop()
ds := &types.DataSource{
ID: "ds-no-sched",
TenantID: 1,
Status: types.DataSourceStatusActive,
SyncSchedule: "",
}
if err := scheduler.AddOrUpdate(ds); err != nil {
t.Fatalf("AddOrUpdate() error: %v", err)
}
if scheduler.EntryCount() != 0 {
t.Errorf("empty schedule should not be registered, EntryCount() = %d", scheduler.EntryCount())
}
}
func TestScheduler_Remove(t *testing.T) {
enqueuer := &fakeTaskEnqueuer{}
scheduler := NewScheduler(newFakeDataSourceRepo(), newFakeSyncLogRepo(), enqueuer)
scheduler.cron.Start()
defer scheduler.Stop()
ds := &types.DataSource{
ID: "ds-rm",
TenantID: 1,
Status: types.DataSourceStatusActive,
SyncSchedule: "0 0 * * * *",
}
_ = scheduler.AddOrUpdate(ds)
if scheduler.EntryCount() != 1 {
t.Fatalf("pre-remove: EntryCount() = %d, want 1", scheduler.EntryCount())
}
scheduler.Remove("ds-rm")
if scheduler.EntryCount() != 0 {
t.Errorf("post-remove: EntryCount() = %d, want 0", scheduler.EntryCount())
}
// Remove non-existent is safe
scheduler.Remove("does-not-exist")
}
func TestScheduler_InvalidCron(t *testing.T) {
enqueuer := &fakeTaskEnqueuer{}
scheduler := NewScheduler(newFakeDataSourceRepo(), newFakeSyncLogRepo(), enqueuer)
scheduler.cron.Start()
defer scheduler.Stop()
ds := &types.DataSource{
ID: "ds-bad",
TenantID: 1,
Status: types.DataSourceStatusActive,
SyncSchedule: "not a cron",
}
err := scheduler.AddOrUpdate(ds)
if err == nil {
t.Fatal("expected error for invalid cron expression")
}
if scheduler.EntryCount() != 0 {
t.Errorf("invalid cron should not be registered, EntryCount() = %d", scheduler.EntryCount())
}
}
func TestScheduler_TriggerSync_InactiveSkipped(t *testing.T) {
repo := newFakeDataSourceRepo()
// Create a data source that is paused
_ = repo.Create(context.Background(), &types.DataSource{
ID: "ds-inactive",
TenantID: 1,
Status: types.DataSourceStatusPaused,
})
enqueuer := &fakeTaskEnqueuer{}
scheduler := NewScheduler(repo, newFakeSyncLogRepo(), enqueuer)
// Directly call triggerSync — it should skip because ds is not active
scheduler.triggerSync("ds-inactive", 1)
if enqueuer.count.Load() != 0 {
t.Error("should not enqueue for inactive data source")
}
}
func TestScheduler_TriggerSync_NotFound(t *testing.T) {
repo := newFakeDataSourceRepo()
enqueuer := &fakeTaskEnqueuer{}
scheduler := NewScheduler(repo, newFakeSyncLogRepo(), enqueuer)
// Should not panic, just skip
scheduler.triggerSync("nonexistent", 1)
if enqueuer.count.Load() != 0 {
t.Error("should not enqueue for non-existent data source")
}
}