mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-08-09 00:47:59 +08:00
f532f27f1f
### Summary As title --------- Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>
377 lines
13 KiB
Go
377 lines
13 KiB
Go
//
|
|
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
|
//
|
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
|
// you may not use this file except in compliance with the License.
|
|
// You may obtain a copy of the License at
|
|
//
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
|
//
|
|
// Unless required by applicable law or agreed to in writing, software
|
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
// See the License for the specific language governing permissions and
|
|
// limitations under the License.
|
|
//
|
|
|
|
package dao
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"ragflow/internal/entity"
|
|
"ragflow/internal/utility"
|
|
"time"
|
|
|
|
"gorm.io/gorm"
|
|
"gorm.io/gorm/clause"
|
|
)
|
|
|
|
const (
|
|
// TaskTypeSync is the Python-compatible SYNC task type.
|
|
TaskTypeSync = "sync"
|
|
// TaskTypePrune is the Python-compatible PRUNE task type.
|
|
TaskTypePrune = "prune"
|
|
// SyncStatusSchedule is the Python TaskStatus.SCHEDULE value for sync_logs.
|
|
SyncStatusSchedule = "5"
|
|
// SyncStatusRunning is the Python TaskStatus.RUNNING value for sync_logs.
|
|
SyncStatusRunning = "1"
|
|
// SyncStatusCancel is the Python TaskStatus.CANCEL value for sync_logs.
|
|
SyncStatusCancel = "2"
|
|
// SyncStatusDone is the Python TaskStatus.DONE value for sync_logs.
|
|
SyncStatusDone = "3"
|
|
// SyncStatusFail is the Python TaskStatus.FAIL value for sync_logs.
|
|
SyncStatusFail = "4"
|
|
)
|
|
|
|
// SyncTask is the Python-compatible sync task row.
|
|
type SyncTask struct {
|
|
entity.SyncLogs `gorm:"embedded"`
|
|
}
|
|
|
|
// TableName returns the Python-compatible sync task table.
|
|
func (SyncTask) TableName() string {
|
|
return "sync_logs"
|
|
}
|
|
|
|
// SyncTaskContext contains every database row required to execute one task.
|
|
type SyncTaskContext struct {
|
|
Task entity.SyncLogs
|
|
Connector entity.Connector
|
|
Connector2Kb entity.Connector2Kb
|
|
Knowledgebase entity.Knowledgebase
|
|
}
|
|
|
|
// ConnectorSource returns the connector source name.
|
|
func (c SyncTaskContext) ConnectorSource() string {
|
|
return c.Connector.Source
|
|
}
|
|
|
|
// SyncTaskDAO reads and updates sync_logs tasks.
|
|
type SyncTaskDAO struct {
|
|
db *gorm.DB
|
|
}
|
|
|
|
// NewSyncTaskDAO creates a syncer task DAO.
|
|
func NewSyncTaskDAO(db *gorm.DB) *SyncTaskDAO {
|
|
if db == nil {
|
|
db = GetDB()
|
|
}
|
|
return &SyncTaskDAO{db: db}
|
|
}
|
|
|
|
// DB returns the DAO database handle.
|
|
func (d *SyncTaskDAO) DB() *gorm.DB {
|
|
return d.db
|
|
}
|
|
|
|
type dueSyncTaskRow struct {
|
|
entity.SyncLogs
|
|
ConnectorRefreshFreq int64 `gorm:"column:connector_refresh_freq"`
|
|
ConnectorPruneFreq int64 `gorm:"column:connector_prune_freq"`
|
|
ConnectorConfig entity.JSONMap `gorm:"column:connector_config"`
|
|
}
|
|
|
|
// ListDueTasks returns due schedule tasks across SYNC and PRUNE.
|
|
func (d *SyncTaskDAO) ListDueTasks(ctx context.Context, now time.Time, limit int) ([]entity.SyncLogs, error) {
|
|
var rows []dueSyncTaskRow
|
|
if err := d.db.WithContext(ctx).
|
|
Model(&entity.SyncLogs{}).
|
|
Select("sync_logs.*, connector.refresh_freq AS connector_refresh_freq, connector.prune_freq AS connector_prune_freq, connector.config AS connector_config").
|
|
Joins("JOIN connector ON sync_logs.connector_id = connector.id").
|
|
Joins("JOIN connector2kb ON sync_logs.connector_id = connector2kb.connector_id AND sync_logs.kb_id = connector2kb.kb_id").
|
|
Joins("JOIN knowledgebase ON sync_logs.kb_id = knowledgebase.id").
|
|
Where("sync_logs.status = ? AND connector.status = ? AND sync_logs.task_type IN ?", SyncStatusSchedule, SyncStatusSchedule, []string{TaskTypeSync, TaskTypePrune}).
|
|
Order("sync_logs.update_time DESC").
|
|
Limit(limit * 4).
|
|
Scan(&rows).Error; err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
tasks := make([]entity.SyncLogs, 0, limit)
|
|
for _, row := range rows {
|
|
if !isDue(row.SyncLogs, row.ConnectorRefreshFreq, row.ConnectorPruneFreq, row.ConnectorConfig, now) {
|
|
continue
|
|
}
|
|
tasks = append(tasks, row.SyncLogs)
|
|
if len(tasks) >= limit {
|
|
break
|
|
}
|
|
}
|
|
return tasks, nil
|
|
}
|
|
|
|
// ClaimTask conditionally marks a scheduled task as running.
|
|
func (d *SyncTaskDAO) ClaimTask(ctx context.Context, taskID string, now time.Time) (bool, error) {
|
|
result := d.db.WithContext(ctx).Model(&entity.SyncLogs{}).
|
|
Where("id = ? AND status = ?", taskID, SyncStatusSchedule).
|
|
Updates(map[string]any{"status": SyncStatusRunning, "time_started": now})
|
|
if result.Error != nil {
|
|
return false, result.Error
|
|
}
|
|
return result.RowsAffected == 1, nil
|
|
}
|
|
|
|
// GetTaskContext loads a task with connector, mapping, and knowledgebase rows.
|
|
func (d *SyncTaskDAO) GetTaskContext(ctx context.Context, taskID string) (SyncTaskContext, error) {
|
|
var task entity.SyncLogs
|
|
if err := d.db.WithContext(ctx).Where("id = ?", taskID).First(&task).Error; err != nil {
|
|
return SyncTaskContext{}, err
|
|
}
|
|
// get connector
|
|
var connector entity.Connector
|
|
if err := d.db.WithContext(ctx).Where("id = ?", task.ConnectorID).First(&connector).Error; err != nil {
|
|
return SyncTaskContext{}, err
|
|
}
|
|
// get relation
|
|
var connector2Kb entity.Connector2Kb
|
|
if err := d.db.WithContext(ctx).Where("connector_id = ? AND kb_id = ?", task.ConnectorID, task.KbID).First(&connector2Kb).Error; err != nil {
|
|
return SyncTaskContext{}, err
|
|
}
|
|
// get KB
|
|
var kb entity.Knowledgebase
|
|
if err := d.db.WithContext(ctx).Where("id = ?", task.KbID).First(&kb).Error; err != nil {
|
|
return SyncTaskContext{}, err
|
|
}
|
|
|
|
return SyncTaskContext{Task: task, Connector: connector, Connector2Kb: connector2Kb, Knowledgebase: kb}, nil
|
|
}
|
|
|
|
// MarkConnectorRunning marks a connector running.
|
|
func (d *SyncTaskDAO) MarkConnectorRunning(ctx context.Context, connectorID string) error {
|
|
return d.db.WithContext(ctx).Model(&entity.Connector{}).Where("id = ?", connectorID).Update("status", SyncStatusRunning).Error
|
|
}
|
|
|
|
// RescheduleClaimed puts a claimed task back into schedule state.
|
|
func (d *SyncTaskDAO) RescheduleClaimed(ctx context.Context, taskID string) error {
|
|
return d.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
|
|
var task entity.SyncLogs
|
|
if err := tx.WithContext(ctx).Where("id = ? AND status = ?", taskID, SyncStatusRunning).First(&task).Error; err != nil {
|
|
if errors.Is(err, gorm.ErrRecordNotFound) {
|
|
return nil
|
|
}
|
|
return err
|
|
}
|
|
if err := tx.Model(&entity.SyncLogs{}).
|
|
Where("id = ? AND status = ?", taskID, SyncStatusRunning).
|
|
Update("status", SyncStatusSchedule).Error; err != nil {
|
|
return err
|
|
}
|
|
return tx.Model(&entity.Connector{}).Where("id = ? AND status = ?", task.ConnectorID, SyncStatusRunning).Update("status", SyncStatusSchedule).Error
|
|
})
|
|
}
|
|
|
|
// FailTask marks a task failed without advancing its poll waterline.
|
|
func (d *SyncTaskDAO) FailTask(ctx context.Context, taskID, connectorID, message string, errorCount int64) error {
|
|
return d.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
|
|
if err := tx.Model(&entity.SyncLogs{}).Where("id = ?", taskID).Updates(map[string]any{
|
|
"status": SyncStatusFail,
|
|
"error_msg": message,
|
|
"error_count": errorCount,
|
|
}).Error; err != nil {
|
|
return err
|
|
}
|
|
if connectorID == "" {
|
|
return nil
|
|
}
|
|
return tx.Model(&entity.Connector{}).Where("id = ?", connectorID).Update("status", SyncStatusFail).Error
|
|
})
|
|
}
|
|
|
|
// CompleteSyncTask marks SYNC done and creates the next schedule task.
|
|
func (d *SyncTaskDAO) CompleteSyncTask(ctx context.Context, taskContext SyncTaskContext, pollRangeEnd time.Time, newDocs, totalDocs, errorCount int64, errorMsg string) error {
|
|
return d.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
|
|
if err := tx.Model(&entity.SyncLogs{}).Where("id = ?", taskContext.Task.ID).Updates(map[string]any{
|
|
"status": SyncStatusDone,
|
|
"poll_range_end": pollRangeEnd,
|
|
"new_docs_indexed": newDocs,
|
|
"total_docs_indexed": gorm.Expr("total_docs_indexed + ?", totalDocs),
|
|
"error_msg": errorMsg,
|
|
"error_count": errorCount,
|
|
}).Error; err != nil {
|
|
return err
|
|
}
|
|
if err := tx.Model(&entity.Connector{}).Where("id = ?", taskContext.Connector.ID).Update("status", SyncStatusDone).Error; err != nil {
|
|
return err
|
|
}
|
|
|
|
return createScheduledTask(ctx, tx, taskContext.Connector.ID, taskContext.Knowledgebase.ID, TaskTypeSync, false, &pollRangeEnd, taskContext.Task.TotalDocsIndexed+totalDocs)
|
|
})
|
|
}
|
|
|
|
// CompletePruneTask marks PRUNE done and creates the next schedule task.
|
|
func (d *SyncTaskDAO) CompletePruneTask(ctx context.Context, taskContext SyncTaskContext, removed int64) error {
|
|
return d.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
|
|
if err := tx.Model(&entity.SyncLogs{}).Where("id = ?", taskContext.Task.ID).Updates(map[string]any{
|
|
"status": SyncStatusDone,
|
|
"docs_removed_from_index": gorm.Expr("docs_removed_from_index + ?", removed),
|
|
}).Error; err != nil {
|
|
return err
|
|
}
|
|
if err := tx.Model(&entity.Connector{}).Where("id = ?", taskContext.Connector.ID).Update("status", SyncStatusDone).Error; err != nil {
|
|
return err
|
|
}
|
|
if !syncConnectorConfigBool(taskContext.Connector.Config, "sync_deleted_files") {
|
|
return nil
|
|
}
|
|
return createScheduledTask(ctx, tx, taskContext.Connector.ID, taskContext.Knowledgebase.ID, TaskTypePrune, false, nil, taskContext.Task.TotalDocsIndexed)
|
|
})
|
|
}
|
|
|
|
// RecoverStaleRunning restores timed-out running tasks to schedule.
|
|
func (d *SyncTaskDAO) RecoverStaleRunning(ctx context.Context, now time.Time) (int64, error) {
|
|
type staleRunningTaskRow struct {
|
|
ID string `gorm:"column:id"`
|
|
ConnectorID string `gorm:"column:connector_id"`
|
|
TimeStarted *time.Time `gorm:"column:time_started"`
|
|
ConnectorTimeoutSecs int64 `gorm:"column:connector_timeout_secs"`
|
|
}
|
|
var rows []staleRunningTaskRow
|
|
if err := d.db.WithContext(ctx).
|
|
Model(&entity.SyncLogs{}).
|
|
Select("sync_logs.id, sync_logs.connector_id, sync_logs.time_started, connector.timeout_secs AS connector_timeout_secs").
|
|
Joins("JOIN connector ON sync_logs.connector_id = connector.id").
|
|
Where("sync_logs.status = ?", SyncStatusRunning).
|
|
Scan(&rows).Error; err != nil {
|
|
return 0, err
|
|
}
|
|
var recovered int64
|
|
connectorIDs := map[string]struct{}{}
|
|
|
|
for _, row := range rows {
|
|
if row.TimeStarted == nil {
|
|
continue
|
|
}
|
|
timeout := time.Duration(row.ConnectorTimeoutSecs) * time.Second
|
|
if timeout <= 0 {
|
|
timeout = time.Hour
|
|
}
|
|
if row.TimeStarted.Add(timeout).After(now) {
|
|
continue
|
|
}
|
|
if err := d.RescheduleClaimed(ctx, row.ID); err != nil {
|
|
return recovered, err
|
|
}
|
|
recovered++
|
|
connectorIDs[row.ConnectorID] = struct{}{}
|
|
}
|
|
|
|
if recovered == 0 {
|
|
return 0, nil
|
|
}
|
|
ids := make([]string, 0, len(connectorIDs))
|
|
|
|
for connectorID := range connectorIDs {
|
|
ids = append(ids, connectorID)
|
|
}
|
|
|
|
return recovered, d.db.WithContext(ctx).Model(&entity.Connector{}).Where("id IN ? AND status = ?", ids, SyncStatusRunning).Update("status", SyncStatusSchedule).Error
|
|
}
|
|
|
|
// createScheduledTask creates the next Python-compatible scheduled task.
|
|
func createScheduledTask(ctx context.Context, tx *gorm.DB, connectorID, kbID, taskType string, fromBeginning bool, pollRangeStart *time.Time, totalDocsIndexed int64) error {
|
|
var lockRow entity.Connector2Kb
|
|
query := tx.WithContext(ctx)
|
|
if tx.Dialector.Name() != "sqlite" {
|
|
query = query.Clauses(clause.Locking{Strength: "UPDATE"})
|
|
}
|
|
if err := query.
|
|
Where("connector_id = ? AND kb_id = ?", connectorID, kbID).
|
|
First(&lockRow).Error; err != nil {
|
|
return err
|
|
}
|
|
|
|
var existing int64
|
|
if err := tx.WithContext(ctx).Model(&entity.SyncLogs{}).
|
|
Where("connector_id = ? AND kb_id = ? AND task_type = ? AND status = ?", connectorID, kbID, taskType, SyncStatusSchedule).
|
|
Count(&existing).Error; err != nil {
|
|
return err
|
|
}
|
|
if existing > 0 {
|
|
return nil
|
|
}
|
|
|
|
reindex := "0"
|
|
if fromBeginning {
|
|
reindex = "1"
|
|
}
|
|
|
|
now := time.Now().Local()
|
|
return tx.WithContext(ctx).Create(&entity.SyncLogs{
|
|
ID: utility.GenerateToken(),
|
|
ConnectorID: connectorID,
|
|
KbID: kbID,
|
|
TaskType: taskType,
|
|
Status: SyncStatusSchedule,
|
|
FromBeginning: &reindex,
|
|
PollRangeStart: pollRangeStart,
|
|
TimeStarted: &now,
|
|
ErrorMsg: "",
|
|
TotalDocsIndexed: totalDocsIndexed,
|
|
}).Error
|
|
}
|
|
|
|
// isDue applies refresh_freq and prune_freq scheduling semantics.
|
|
func isDue(task entity.SyncLogs, refreshFreq, pruneFreq int64, config map[string]any, now time.Time) bool {
|
|
if task.UpdateDate == nil {
|
|
return true
|
|
}
|
|
var freqMinutes int64
|
|
switch task.TaskType {
|
|
case TaskTypeSync:
|
|
freqMinutes = refreshFreq
|
|
case TaskTypePrune:
|
|
if !syncConnectorConfigBool(config, "sync_deleted_files") {
|
|
return false
|
|
}
|
|
freqMinutes = pruneFreq
|
|
default:
|
|
return false
|
|
}
|
|
if freqMinutes <= 0 {
|
|
return true
|
|
}
|
|
return task.UpdateDate.Before(now.Add(-time.Duration(freqMinutes) * time.Minute))
|
|
}
|
|
|
|
// syncConnectorConfigBool reads a Python JSON bool/string flag.
|
|
func syncConnectorConfigBool(config map[string]any, key string) bool {
|
|
value, ok := config[key]
|
|
if !ok {
|
|
return false
|
|
}
|
|
switch typed := value.(type) {
|
|
case bool:
|
|
return typed
|
|
case string:
|
|
return typed == "1" || typed == "true" || typed == "TRUE"
|
|
default:
|
|
return false
|
|
}
|
|
}
|
|
|
|
// IsNotFound reports whether an error is a gorm not found error.
|
|
func IsNotFound(err error) bool {
|
|
return errors.Is(err, gorm.ErrRecordNotFound)
|
|
}
|