From 615c7d8bcb08926ce0631dc7da996ebf7afe4637 Mon Sep 17 00:00:00 2001 From: Haruko386 Date: Fri, 7 Aug 2026 13:35:57 +0800 Subject: [PATCH] feat[Go]: implement data source gmail and test prune (#17969) ### Summary As title --- internal/dao/sync_prune_snapshot.go | 69 --- internal/service/sync_prune_service.go | 42 +- internal/service/sync_task_service.go | 2 + internal/syncer/connector/gmail.go | 731 ++++++++++++++++++++++++ internal/syncer/connector/gmail_test.go | 160 ++++++ internal/syncer/prune_runner.go | 17 +- internal/syncer/syncer.go | 28 +- internal/syncer/syncer_test.go | 15 +- internal/syncer/task_worker.go | 5 + 9 files changed, 947 insertions(+), 122 deletions(-) delete mode 100644 internal/dao/sync_prune_snapshot.go create mode 100644 internal/syncer/connector/gmail.go create mode 100644 internal/syncer/connector/gmail_test.go diff --git a/internal/dao/sync_prune_snapshot.go b/internal/dao/sync_prune_snapshot.go deleted file mode 100644 index bf4d4a964f..0000000000 --- a/internal/dao/sync_prune_snapshot.go +++ /dev/null @@ -1,69 +0,0 @@ -// -// 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" - syncerconnector "ragflow/internal/syncer/connector" - - "gorm.io/gorm" -) - -// SyncPruneSnapshot stores one source ID from a complete prune snapshot. -type SyncPruneSnapshot struct { - TaskID string `gorm:"column:task_id;primaryKey;size:32"` - SourceDocID string `gorm:"column:source_doc_id;primaryKey;size:512"` -} - -// TableName returns the prune snapshot table name. -func (SyncPruneSnapshot) TableName() string { - return "sync_prune_snapshot" -} - -// SyncPruneSnapshotDAO manages temporary prune snapshots. -type SyncPruneSnapshotDAO struct { - db *gorm.DB -} - -// NewSyncPruneSnapshotDAO creates a prune snapshot DAO. -func NewSyncPruneSnapshotDAO(db *gorm.DB) *SyncPruneSnapshotDAO { - return &SyncPruneSnapshotDAO{db: db} -} - -// InsertBatch writes one prune snapshot batch. -func (d *SyncPruneSnapshotDAO) InsertBatch(ctx context.Context, taskID string, documents []syncerconnector.SlimDocument) error { - if len(documents) == 0 { - return nil - } - rows := make([]SyncPruneSnapshot, 0, len(documents)) - for _, doc := range documents { - rows = append(rows, SyncPruneSnapshot{TaskID: taskID, SourceDocID: doc.SourceID}) - } - return d.db.WithContext(ctx).Create(&rows).Error -} - -// ListSourceIDs returns all source IDs for a task snapshot. -func (d *SyncPruneSnapshotDAO) ListSourceIDs(ctx context.Context, taskID string) ([]string, error) { - var ids []string - err := d.db.WithContext(ctx).Model(&SyncPruneSnapshot{}).Where("task_id = ?", taskID).Pluck("source_doc_id", &ids).Error - return ids, err -} - -// Clear removes all snapshot rows for a task. -func (d *SyncPruneSnapshotDAO) Clear(ctx context.Context, taskID string) error { - return d.db.WithContext(ctx).Where("task_id = ?", taskID).Delete(&SyncPruneSnapshot{}).Error -} diff --git a/internal/service/sync_prune_service.go b/internal/service/sync_prune_service.go index 107464ee28..7277ef7568 100644 --- a/internal/service/sync_prune_service.go +++ b/internal/service/sync_prune_service.go @@ -19,8 +19,6 @@ package service import ( "context" "errors" - "ragflow/internal/dao" - syncerconnector "ragflow/internal/syncer/connector" ) // ErrSyncDocumentDeleterNotConfigured reports a missing production delete service. @@ -32,39 +30,23 @@ type SyncDocumentDeleter interface { DeleteDocument(ctx context.Context, docID string) error } -// SyncPruneService owns prune snapshots and stale document deletion. +// SyncPruneService owns stale document deletion for complete prune snapshots. type SyncPruneService struct { - snapshotDAO *dao.SyncPruneSnapshotDAO - deleter SyncDocumentDeleter - store DocumentStore + deleter SyncDocumentDeleter + store DocumentStore } // NewSyncPruneService creates a prune service. -func NewSyncPruneService(snapshotDAO *dao.SyncPruneSnapshotDAO, deleter SyncDocumentDeleter, store DocumentStore) *SyncPruneService { +func NewSyncPruneService(deleter SyncDocumentDeleter, store DocumentStore) *SyncPruneService { if store == nil { store = NewGormDocumentStore() } - return &SyncPruneService{snapshotDAO: snapshotDAO, deleter: deleter, store: store} -} - -// ClearSnapshot removes temporary snapshot rows. -func (s *SyncPruneService) ClearSnapshot(ctx context.Context, taskID string) error { - return s.snapshotDAO.Clear(ctx, taskID) -} - -// AddSnapshotBatch writes one slim snapshot batch. -func (s *SyncPruneService) AddSnapshotBatch(ctx context.Context, taskID string, documents []syncerconnector.SlimDocument) error { - return s.snapshotDAO.InsertBatch(ctx, taskID, documents) + return &SyncPruneService{deleter: deleter, store: store} } // DeleteStale removes documents missing from the complete source snapshot. -func (s *SyncPruneService) DeleteStale(ctx context.Context, taskContext SyncTaskContext) (int64, error) { - sourceIDs, err := s.snapshotDAO.ListSourceIDs(ctx, taskContext.Task.ID) - if err != nil { - return 0, err - } +func (s *SyncPruneService) DeleteStale(ctx context.Context, taskContext SyncTaskContext, retain map[string]struct{}) (int64, error) { sourceType := SourceType(taskContext.Connector.Source, taskContext.Connector.ID) - retain := RetainDocumentIDs(taskContext.Knowledgebase.ID, taskContext.Connector.ID, sourceIDs) existing, err := s.store.ListIDs(ctx, taskContext.Knowledgebase.ID, sourceType) if err != nil { return 0, err @@ -82,9 +64,6 @@ func (s *SyncPruneService) DeleteStale(ctx context.Context, taskContext SyncTask } removed++ } - if err = s.snapshotDAO.Clear(ctx, taskContext.Task.ID); err != nil { - return removed, err - } return removed, nil } @@ -92,8 +71,13 @@ func (s *SyncPruneService) DeleteStale(ctx context.Context, taskContext SyncTask func RetainDocumentIDs(kbID, connectorID string, sourceIDs []string) map[string]struct{} { retain := make(map[string]struct{}, len(sourceIDs)*2) for _, sourceID := range sourceIDs { - retain[Hash128(connectorID+":"+sourceID)] = struct{}{} - retain[Hash128(kbID+":"+connectorID+":"+sourceID)] = struct{}{} + AddRetainDocumentID(retain, kbID, connectorID, sourceID) } return retain } + +// AddRetainDocumentID adds both legacy and current document IDs for one source ID. +func AddRetainDocumentID(retain map[string]struct{}, kbID, connectorID, sourceID string) { + retain[Hash128(connectorID+":"+sourceID)] = struct{}{} + retain[Hash128(kbID+":"+connectorID+":"+sourceID)] = struct{}{} +} diff --git a/internal/service/sync_task_service.go b/internal/service/sync_task_service.go index 5d383ec1d0..c9ee375cc6 100644 --- a/internal/service/sync_task_service.go +++ b/internal/service/sync_task_service.go @@ -107,6 +107,8 @@ func (s *SyncTaskService) Claim(ctx context.Context, taskID string) (bool, error return true, s.taskDAO.MarkConnectorRunning(ctx, taskContext.Connector.ID) } +// TODO: refactor some needless func + // GetContext loads a task execution context. func (s *SyncTaskService) GetContext(ctx context.Context, taskID string) (SyncTaskContext, error) { return s.taskDAO.GetTaskContext(ctx, taskID) diff --git a/internal/syncer/connector/gmail.go b/internal/syncer/connector/gmail.go new file mode 100644 index 0000000000..408112a19c --- /dev/null +++ b/internal/syncer/connector/gmail.go @@ -0,0 +1,731 @@ +// +// 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 connector + +import ( + "context" + "encoding/base64" + "encoding/json" + "fmt" + "io" + "net/http" + "net/mail" + "net/url" + "os" + "strings" + "sync" + "time" + + "golang.org/x/oauth2" + "golang.org/x/oauth2/google" +) + +const ( + defaultGmailBatchSize = 32 + gmailItemsPerPage = 100 + gmailRequestTimeout = 60 * time.Second + gmailOAuthTokenURL = "https://oauth2.googleapis.com/token" +) + +var gmailScopes = []string{ + "https://www.googleapis.com/auth/gmail.readonly", + "https://www.googleapis.com/auth/admin.directory.user.readonly", + "https://www.googleapis.com/auth/admin.directory.group.readonly", +} + +// FIXME: IDK why everytime, gmail do sync, It will update all file's Metadata. +// FIXME: I think this need to be checked or fixed after all data syncer is done +// FIXME: Some file sync from gmail have no content, this need to be checked too + +// GmailConnector reads Gmail threads from a Workspace domain or one Gmail account. +type GmailConnector struct { + primaryAdminEmail string + credentials map[string]any + batchSize int + clientMu sync.Mutex + clients map[string]*http.Client + httpClientForUser func(ctx context.Context, userEmail string) (*http.Client, error) + listUsers func(ctx context.Context) ([]string, error) + listThreadPage func(ctx context.Context, userEmail, query, pageToken string, pageSize int) (gmailThreadListPage, error) + getThread func(ctx context.Context, userEmail, threadID string) (gmailThread, error) +} + +// NewGmailConnector creates a Gmail connector from Python-compatible config. +func NewGmailConnector(config map[string]any) (*GmailConnector, error) { + credentials, _ := config["credentials"].(map[string]any) + return &GmailConnector{ + primaryAdminEmail: strings.TrimSpace(stringConfig(credentials["google_primary_admin"])), + credentials: credentials, + batchSize: configInt(config["batch_size"], defaultGmailBatchSize), + clients: map[string]*http.Client{}, + }, nil +} + +// Validate validates Gmail connector settings and credentials. +func (c *GmailConnector) Validate(ctx context.Context) error { + if c == nil { + return fmt.Errorf("gmail connector is nil") + } + if c.primaryAdminEmail == "" { + return fmt.Errorf("Gmail connector is missing google_primary_admin") + } + if len(c.credentials) == 0 { + return fmt.Errorf("Gmail connector is missing credentials") + } + if c.batchSize <= 0 { + return fmt.Errorf("batch_size must be a positive integer") + } + _, err := c.getUserEmails(ctx) + return err +} + +// OpenSync opens one Gmail sync session. +func (c *GmailConnector) OpenSync(ctx context.Context, request SyncRequest) (SyncSession, error) { + users, err := c.getUserEmails(ctx) + if err != nil { + return nil, err + } + query := "" + if !request.FromBeginning && request.WindowStart != nil { + query = gmailTimeRangeQuery(request.WindowStart, request.WindowEnd) + } + return &gmailSyncSession{connector: c, users: users, batchSize: c.batchSize, query: query}, nil +} + +// OpenPrune opens one complete Gmail prune snapshot session. +func (c *GmailConnector) OpenPrune(ctx context.Context, request PruneRequest) (PruneSession, error) { + users, err := c.getUserEmails(ctx) + if err != nil { + return nil, err + } + return &gmailPruneSession{connector: c, users: users, batchSize: c.batchSize}, nil +} + +// getUserEmails returns Workspace users or falls back to the primary account for personal Gmail. +func (c *GmailConnector) getUserEmails(ctx context.Context) ([]string, error) { + if c.listUsers != nil { + return c.listUsers(ctx) + } + client, err := c.clientForUser(ctx, c.primaryAdminEmail) + if err != nil { + return nil, err + } + + domain := c.primaryAdminEmail + if _, after, ok := strings.Cut(c.primaryAdminEmail, "@"); ok { + domain = after + } + users := []string{} + pageToken := "" + for { + query := url.Values{ + "domain": {domain}, + "fields": {"nextPageToken,users(primaryEmail)"}, + "maxResults": {"500"}, + } + if pageToken != "" { + query.Set("pageToken", pageToken) + } + var page gmailUsersPage + err = c.getJSON(ctx, client, "https://admin.googleapis.com/admin/directory/v1/users?"+query.Encode(), &page) + if err != nil { + if httpErr, ok := err.(googleHTTPError); ok && httpErr.status == http.StatusNotFound { + return []string{c.primaryAdminEmail}, nil + } + return nil, err + } + for _, user := range page.Users { + if user.PrimaryEmail != "" { + users = append(users, user.PrimaryEmail) + } + } + if page.NextPageToken == "" { + break + } + pageToken = page.NextPageToken + } + if len(users) == 0 { + return []string{c.primaryAdminEmail}, nil + } + return users, nil +} + +// listThreads returns one Gmail thread list page. +func (c *GmailConnector) listThreads(ctx context.Context, userEmail, queryText, pageToken string, pageSize int) (gmailThreadListPage, error) { + if c.listThreadPage != nil { + return c.listThreadPage(ctx, userEmail, queryText, pageToken, pageSize) + } + client, err := c.clientForUser(ctx, userEmail) + if err != nil { + return gmailThreadListPage{}, err + } + query := url.Values{ + "fields": {"nextPageToken,threads(id)"}, + "maxResults": {fmt.Sprint(pageSize)}, + } + if queryText != "" { + query.Set("q", queryText) + } + if pageToken != "" { + query.Set("pageToken", pageToken) + } + var page gmailThreadListPage + err = c.getJSON(ctx, client, "https://gmail.googleapis.com/gmail/v1/users/"+url.PathEscape(userEmail)+"/threads?"+query.Encode(), &page) + return page, err +} + +// loadThread returns one full Gmail thread. +func (c *GmailConnector) loadThread(ctx context.Context, userEmail, threadID string) (gmailThread, error) { + if c.getThread != nil { + return c.getThread(ctx, userEmail, threadID) + } + client, err := c.clientForUser(ctx, userEmail) + if err != nil { + return gmailThread{}, err + } + query := url.Values{"fields": {"id,messages(id,labelIds,payload(headers,parts(body(data),mimeType)))"}} + var thread gmailThread + err = c.getJSON(ctx, client, "https://gmail.googleapis.com/gmail/v1/users/"+url.PathEscape(userEmail)+"/threads/"+url.PathEscape(threadID)+"?"+query.Encode(), &thread) + return thread, err +} + +// clientForUser builds an authenticated Google client for a user. +func (c *GmailConnector) clientForUser(ctx context.Context, userEmail string) (*http.Client, error) { + c.clientMu.Lock() + defer c.clientMu.Unlock() + if c.clients == nil { + c.clients = map[string]*http.Client{} + } + if client := c.clients[userEmail]; client != nil { + return client, nil + } + + var client *http.Client + var err error + if c.httpClientForUser != nil { + client, err = c.httpClientForUser(ctx, userEmail) + } else { + var tokenSource oauth2.TokenSource + tokenSource, err = c.tokenSource(ctx, userEmail) + if err == nil { + client = oauth2.NewClient(ctx, tokenSource) + } + } + if err != nil { + return nil, err + } + c.clients[userEmail] = client + return client, nil +} + +// tokenSource returns a Google OAuth token source for stored connector credentials. +func (c *GmailConnector) tokenSource(ctx context.Context, userEmail string) (oauth2.TokenSource, error) { + if value := stringConfig(c.credentials["google_service_account_key"]); value != "" { + config, err := google.JWTConfigFromJSON([]byte(value), gmailScopes...) + if err != nil { + return nil, err + } + config.Subject = userEmail + return config.TokenSource(ctx), nil + } + + tokenJSON := stringConfig(c.credentials["google_tokens"]) + if tokenJSON == "" { + return nil, fmt.Errorf("Gmail connector credentials must include google_tokens or google_service_account_key") + } + var tokenData gmailOAuthToken + if err := json.Unmarshal([]byte(tokenJSON), &tokenData); err != nil { + return nil, err + } + if tokenData.ClientID == "" { + tokenData.ClientID = os.Getenv("OAUTH_GOOGLE_DRIVE_CLIENT_ID") + } + if tokenData.ClientSecret == "" { + tokenData.ClientSecret = os.Getenv("OAUTH_GOOGLE_DRIVE_CLIENT_SECRET") + } + if tokenData.ClientID == "" || tokenData.ClientSecret == "" || tokenData.RefreshToken == "" { + return nil, fmt.Errorf("Gmail OAuth credentials are incomplete") + } + return (&oauth2.Config{ + ClientID: tokenData.ClientID, + ClientSecret: tokenData.ClientSecret, + Endpoint: oauth2.Endpoint{TokenURL: gmailOAuthTokenURL}, + Scopes: gmailScopes, + }).TokenSource(ctx, tokenData.token()), nil +} + +// getJSON fetches and decodes a Google JSON API response. +func (c *GmailConnector) getJSON(ctx context.Context, client *http.Client, apiURL string, out any) error { + ctx, cancel := context.WithTimeout(ctx, gmailRequestTimeout) + defer cancel() + + req, err := http.NewRequestWithContext(ctx, http.MethodGet, apiURL, nil) + if err != nil { + return err + } + + resp, err := client.Do(req) + if err != nil { + return err + } + defer resp.Body.Close() + if resp.StatusCode >= 400 { + body, _ := io.ReadAll(io.LimitReader(resp.Body, 4096)) + return googleHTTPError{status: resp.StatusCode, body: strings.TrimSpace(string(body))} + } + return json.NewDecoder(resp.Body).Decode(out) +} + +type gmailSyncSession struct { + connector *GmailConnector + users []string + userIndex int + pageToken string + batchSize int + query string + buffer []SourceDocument +} + +// NextBatch returns the next Gmail document batch. +func (s *gmailSyncSession) NextBatch(ctx context.Context) (SyncBatch, error) { + documents := make([]SourceDocument, 0, s.batchSize) + if len(s.buffer) > 0 { + n := min(s.batchSize, len(s.buffer)) + documents = append(documents, s.buffer[:n]...) + s.buffer = s.buffer[n:] + } + for len(documents) < s.batchSize { + if s.userIndex >= len(s.users) { + if len(documents) == 0 { + return SyncBatch{}, io.EOF + } + break + } + batch, err := s.nextDocumentPage(ctx) + if err != nil { + return SyncBatch{}, err + } + remaining := s.batchSize - len(documents) + if len(batch) > remaining { + documents = append(documents, batch[:remaining]...) + s.buffer = append(s.buffer, batch[remaining:]...) + break + } + documents = append(documents, batch...) + } + return SyncBatch{Documents: documents}, nil +} + +// Close closes the Gmail sync session. +func (s *gmailSyncSession) Close() error { + return nil +} + +// nextDocumentPage fetches one Gmail list page and expands threads. +func (s *gmailSyncSession) nextDocumentPage(ctx context.Context) ([]SourceDocument, error) { + userEmail := s.users[s.userIndex] + page, err := s.connector.listThreads(ctx, userEmail, s.query, s.pageToken, gmailItemsPerPage) + if err != nil { + if isGmailDisabled(err) { + s.advanceUser() + return nil, nil + } + return nil, err + } + documents := make([]SourceDocument, 0, len(page.Threads)) + for _, item := range page.Threads { + thread, err := s.connector.loadThread(ctx, userEmail, item.ID) + if err != nil { + if isGmailDisabled(err) || isGoogleForbiddenOrNotFound(err) { + continue + } + return nil, err + } + doc, ok := thread.toSourceDocument(userEmail) + if ok { + documents = append(documents, doc) + } + } + if page.NextPageToken == "" { + s.advanceUser() + } else { + s.pageToken = page.NextPageToken + } + return documents, nil +} + +// advanceUser moves a Gmail session to the next mailbox. +func (s *gmailSyncSession) advanceUser() { + s.userIndex++ + s.pageToken = "" +} + +type gmailPruneSession struct { + connector *GmailConnector + users []string + userIndex int + pageToken string + batchSize int + buffer []SlimDocument +} + +// NextBatch returns the next Gmail prune snapshot batch. +func (s *gmailPruneSession) NextBatch(ctx context.Context) (PruneBatch, error) { + documents := make([]SlimDocument, 0, s.batchSize) + if len(s.buffer) > 0 { + n := min(s.batchSize, len(s.buffer)) + documents = append(documents, s.buffer[:n]...) + s.buffer = s.buffer[n:] + } + for len(documents) < s.batchSize { + if s.userIndex >= len(s.users) { + if len(documents) == 0 { + return PruneBatch{}, io.EOF + } + break + } + batch, err := s.nextSlimPage(ctx) + if err != nil { + return PruneBatch{}, err + } + remaining := s.batchSize - len(documents) + if len(batch) > remaining { + documents = append(documents, batch[:remaining]...) + s.buffer = append(s.buffer, batch[remaining:]...) + break + } + documents = append(documents, batch...) + } + return PruneBatch{Documents: documents}, nil +} + +// Close closes the Gmail prune session. +func (s *gmailPruneSession) Close() error { + return nil +} + +// nextSlimPage fetches one Gmail thread ID page. +func (s *gmailPruneSession) nextSlimPage(ctx context.Context) ([]SlimDocument, error) { + userEmail := s.users[s.userIndex] + page, err := s.connector.listThreads(ctx, userEmail, "", s.pageToken, gmailItemsPerPage) + if err != nil { + if isGmailDisabled(err) { + s.advanceUser() + return nil, nil + } + return nil, err + } + documents := make([]SlimDocument, 0, len(page.Threads)) + for _, thread := range page.Threads { + if thread.ID != "" { + documents = append(documents, SlimDocument{SourceID: thread.ID}) + } + } + if page.NextPageToken == "" { + s.advanceUser() + } else { + s.pageToken = page.NextPageToken + } + return documents, nil +} + +// advanceUser moves a Gmail prune session to the next mailbox. +func (s *gmailPruneSession) advanceUser() { + s.userIndex++ + s.pageToken = "" +} + +type gmailUsersPage struct { + NextPageToken string `json:"nextPageToken"` + Users []struct { + PrimaryEmail string `json:"primaryEmail"` + } `json:"users"` +} + +type gmailThreadListPage struct { + NextPageToken string `json:"nextPageToken"` + Threads []struct { + ID string `json:"id"` + } `json:"threads"` +} + +type gmailThread struct { + ID string `json:"id"` + Messages []gmailMessage `json:"messages"` +} + +// toSourceDocument converts a Gmail thread into the syncer model. +func (t gmailThread) toSourceDocument(userEmail string) (SourceDocument, bool) { + if t.ID == "" || len(t.Messages) == 0 { + return SourceDocument{}, false + } + sections := make([]string, 0, len(t.Messages)) + semanticIdentifier := "" + updatedAt := time.Time{} + metadata := map[string]any{} + fromEmails := map[string]string{} + otherEmails := map[string]string{} + + for _, message := range t.Messages { + body, messageMetadata := message.toTextSection() + if body != "" { + sections = append(sections, body) + } + for key, value := range messageMetadata { + metadata[key] = value + if isGmailEmailHeader(key) { + email, name := parseGmailAddress(stringConfig(value)) + if email == "" { + continue + } + if key == "from" { + fromEmails[email] = name + } else { + otherEmails[email] = name + } + } + } + if semanticIdentifier == "" { + semanticIdentifier = sanitizeGmailName(stringConfig(messageMetadata["subject"])) + } + if value := stringConfig(messageMetadata["updated_at"]); value != "" { + if parsed, err := mail.ParseDate(value); err == nil && parsed.UTC().After(updatedAt) { + updatedAt = parsed.UTC() + } + } + } + if semanticIdentifier == "" { + semanticIdentifier = "(no subject)" + } + if updatedAt.IsZero() { + updatedAt = time.Now().UTC() + } + blob := []byte(strings.Join(sections, "\n\n")) + metadata["external_user_emails"] = []string{userEmail} + metadata["primary_owners"] = gmailOwnersMetadata(fromEmails) + metadata["secondary_owners"] = gmailOwnersMetadata(otherEmails) + return SourceDocument{ + SourceID: t.ID, + SemanticIdentifier: semanticIdentifier, + Extension: ".txt", + Blob: blob, + UpdatedAt: updatedAt, + SizeBytes: int64(len(blob)), + Metadata: metadata, + }, true +} + +type gmailMessage struct { + ID string `json:"id"` + LabelIDs []string `json:"labelIds"` + Payload gmailPayload `json:"payload"` +} + +// toTextSection converts a Gmail message to text and metadata. +func (m gmailMessage) toTextSection() (string, map[string]any) { + metadata := map[string]any{} + for _, header := range m.Payload.Headers { + name := strings.ToLower(header.Name) + if isGmailEmailHeader(name) { + metadata[name] = header.Value + } + if name == "subject" { + metadata["subject"] = header.Value + } + if name == "date" { + metadata["updated_at"] = header.Value + } + } + if len(m.LabelIDs) > 0 { + metadata["labels"] = m.LabelIDs + } + + var builder strings.Builder + builder.WriteString(gmailPayloadText(m.Payload)) + for _, key := range []string{"from", "to", "cc", "bcc", "subject", "labels"} { + if value, ok := metadata[key]; ok { + builder.WriteString(fmt.Sprintf("%s: %v\n", key, value)) + } + } + return builder.String(), metadata +} + +type gmailPayload struct { + Headers []gmailHeader `json:"headers"` + Parts []gmailPart `json:"parts"` + Body gmailBody `json:"body"` + MimeType string `json:"mimeType"` +} + +type gmailPart struct { + MimeType string `json:"mimeType"` + Body gmailBody `json:"body"` + Parts []gmailPart `json:"parts"` +} + +type gmailBody struct { + Data string `json:"data"` +} + +type gmailHeader struct { + Name string `json:"name"` + Value string `json:"value"` +} + +// gmailPayloadText extracts text/plain body content from a Gmail payload. +func gmailPayloadText(payload gmailPayload) string { + var builder strings.Builder + if payload.MimeType == "text/plain" && payload.Body.Data != "" { + builder.WriteString(decodeGmailBody(payload.Body.Data)) + } + for _, part := range payload.Parts { + builder.WriteString(gmailPartText(part)) + } + return builder.String() +} + +// gmailPartText extracts text/plain content recursively. +func gmailPartText(part gmailPart) string { + var builder strings.Builder + if part.MimeType == "text/plain" && part.Body.Data != "" { + builder.WriteString(decodeGmailBody(part.Body.Data)) + } + for _, child := range part.Parts { + builder.WriteString(gmailPartText(child)) + } + return builder.String() +} + +// decodeGmailBody decodes Gmail's URL-safe base64 body format. +func decodeGmailBody(value string) string { + value = strings.TrimSpace(value) + if value == "" { + return "" + } + if decoded, err := base64.RawURLEncoding.DecodeString(value); err == nil { + return string(decoded) + } + if decoded, err := base64.URLEncoding.DecodeString(value); err == nil { + return string(decoded) + } + return "" +} + +// gmailTimeRangeQuery builds Python-compatible Gmail time range query text. +func gmailTimeRangeQuery(windowStart *time.Time, windowEnd time.Time) string { + parts := []string{} + if windowStart != nil && !windowStart.IsZero() { + parts = append(parts, fmt.Sprintf("after:%d", windowStart.Unix()+1)) + } + if !windowEnd.IsZero() { + parts = append(parts, fmt.Sprintf("before:%d", windowEnd.Unix())) + } + return strings.Join(parts, " ") +} + +// isGmailEmailHeader reports whether a metadata key is an email header of interest. +func isGmailEmailHeader(name string) bool { + switch name { + case "cc", "bcc", "from", "to": + return true + default: + return false + } +} + +// parseGmailAddress extracts an email address and display name. +func parseGmailAddress(value string) (string, string) { + address, err := mail.ParseAddress(value) + if err != nil { + return strings.TrimSpace(value), "" + } + return address.Address, address.Name +} + +// gmailOwnersMetadata converts email owners to compact metadata. +func gmailOwnersMetadata(owners map[string]string) []map[string]string { + out := make([]map[string]string, 0, len(owners)) + for email, name := range owners { + item := map[string]string{"email": email} + if name != "" { + parts := strings.Fields(name) + if len(parts) > 1 { + item["first_name"] = strings.Join(parts[:len(parts)-1], " ") + item["last_name"] = parts[len(parts)-1] + } else { + item["last_name"] = name + } + } + out = append(out, item) + } + return out +} + +// sanitizeGmailName mirrors Python's sanitized filename intent. +func sanitizeGmailName(name string) string { + name = strings.TrimSpace(name) + if name == "" { + return "" + } + replacer := strings.NewReplacer("/", "_", "\\", "_", ":", "_", "*", "_", "?", "_", `"`, "_", "<", "_", ">", "_", "|", "_") + return replacer.Replace(name) +} + +// isGmailDisabled reports mailbox-not-provisioned Gmail failures. +func isGmailDisabled(err error) bool { + if httpErr, ok := err.(googleHTTPError); ok { + return httpErr.status == http.StatusBadRequest && (strings.Contains(httpErr.body, "Mail service not enabled") || strings.Contains(httpErr.body, "failedPrecondition")) + } + return false +} + +// isGoogleForbiddenOrNotFound reports item-level Google visibility failures. +func isGoogleForbiddenOrNotFound(err error) bool { + if httpErr, ok := err.(googleHTTPError); ok { + return httpErr.status == http.StatusForbidden || httpErr.status == http.StatusNotFound + } + return false +} + +type googleHTTPError struct { + status int + body string +} + +// Error returns the Google API error text. +func (e googleHTTPError) Error() string { + return fmt.Sprintf("Google API returned HTTP %d: %s", e.status, e.body) +} + +type gmailOAuthToken struct { + AccessToken string `json:"token"` + RefreshToken string `json:"refresh_token"` + ClientID string `json:"client_id"` + ClientSecret string `json:"client_secret"` + Expiry string `json:"expiry"` +} + +// token converts stored OAuth JSON to oauth2.Token. +func (t gmailOAuthToken) token() *oauth2.Token { + token := &oauth2.Token{AccessToken: t.AccessToken, RefreshToken: t.RefreshToken} + if t.Expiry != "" { + if parsed, err := time.Parse(time.RFC3339Nano, t.Expiry); err == nil { + token.Expiry = parsed + } + } + return token +} diff --git a/internal/syncer/connector/gmail_test.go b/internal/syncer/connector/gmail_test.go new file mode 100644 index 0000000000..f266dbb87d --- /dev/null +++ b/internal/syncer/connector/gmail_test.go @@ -0,0 +1,160 @@ +package connector + +import ( + "context" + "encoding/base64" + "errors" + "io" + "net/http" + "strings" + "testing" + "time" +) + +// TestGmailConnectorOpenSync verifies Gmail thread expansion and incremental query construction. +func TestGmailConnectorOpenSync(t *testing.T) { + connector := newFixtureGmailConnector() + start := mustTime(t, "2026-01-02T00:00:00Z") + end := mustTime(t, "2026-01-03T00:00:00Z") + + session, err := connector.OpenSync(context.Background(), SyncRequest{WindowStart: &start, WindowEnd: end}) + if err != nil { + t.Fatalf("OpenSync failed: %v", err) + } + batch, err := session.NextBatch(context.Background()) + if err != nil { + t.Fatalf("NextBatch failed: %v", err) + } + if connector.lastQuery != "after:1767312001 before:1767398400" { + t.Fatalf("query = %q", connector.lastQuery) + } + if len(batch.Documents) != 1 { + t.Fatalf("documents len = %d, want 1", len(batch.Documents)) + } + doc := batch.Documents[0] + if doc.SourceID != "thread-1" { + t.Fatalf("source id = %q", doc.SourceID) + } + if doc.SemanticIdentifier != "Hello_World" { + t.Fatalf("semantic identifier = %q", doc.SemanticIdentifier) + } + blob := string(doc.Blob) + if !strings.Contains(blob, "Body text") || + !strings.Contains(blob, "from: Alice Example ") || + !strings.Contains(blob, "to: Bob ") || + !strings.Contains(blob, "subject: Hello/World") { + t.Fatalf("blob = %q", string(doc.Blob)) + } + if !doc.UpdatedAt.Equal(mustTime(t, "2026-01-02T03:04:05Z")) { + t.Fatalf("updated at = %s", doc.UpdatedAt) + } + if doc.Metadata["external_user_emails"].([]string)[0] != "admin@example.com" { + t.Fatalf("external user metadata = %v", doc.Metadata["external_user_emails"]) + } + if _, err = session.NextBatch(context.Background()); !errors.Is(err, io.EOF) { + t.Fatalf("NextBatch EOF = %v", err) + } +} + +// TestGmailConnectorOpenPrune verifies Gmail prune emits thread IDs only. +func TestGmailConnectorOpenPrune(t *testing.T) { + connector := newFixtureGmailConnector() + session, err := connector.OpenPrune(context.Background(), PruneRequest{}) + if err != nil { + t.Fatalf("OpenPrune failed: %v", err) + } + batch, err := session.NextBatch(context.Background()) + if err != nil { + t.Fatalf("NextBatch failed: %v", err) + } + if len(batch.Documents) != 1 || batch.Documents[0].SourceID != "thread-1" { + t.Fatalf("unexpected prune batch: %+v", batch.Documents) + } +} + +// TestGmailConnectorClientForUserCachesByMailbox verifies clients are reused per user. +func TestGmailConnectorClientForUserCachesByMailbox(t *testing.T) { + connector := &GmailConnector{} + calls := map[string]int{} + connector.httpClientForUser = func(ctx context.Context, userEmail string) (*http.Client, error) { + calls[userEmail]++ + return &http.Client{}, nil + } + + first, err := connector.clientForUser(context.Background(), "a@example.com") + if err != nil { + t.Fatalf("first client: %v", err) + } + second, err := connector.clientForUser(context.Background(), "a@example.com") + if err != nil { + t.Fatalf("second client: %v", err) + } + other, err := connector.clientForUser(context.Background(), "b@example.com") + if err != nil { + t.Fatalf("other client: %v", err) + } + if first != second { + t.Fatalf("same mailbox did not reuse client") + } + if first == other { + t.Fatalf("different mailboxes shared client") + } + if calls["a@example.com"] != 1 || calls["b@example.com"] != 1 { + t.Fatalf("client creation calls = %+v", calls) + } +} + +type fixtureGmailConnector struct { + *GmailConnector + lastQuery string +} + +func newFixtureGmailConnector() *fixtureGmailConnector { + base, _ := NewGmailConnector(map[string]any{ + "batch_size": 1, + "credentials": map[string]any{ + "google_primary_admin": "admin@example.com", + "google_tokens": `{"client_id":"client","client_secret":"secret","refresh_token":"refresh"}`, + }, + }) + connector := &fixtureGmailConnector{GmailConnector: base} + base.listUsers = func(ctx context.Context) ([]string, error) { + return []string{"admin@example.com"}, nil + } + base.listThreadPage = func(ctx context.Context, userEmail, query, pageToken string, pageSize int) (gmailThreadListPage, error) { + connector.lastQuery = query + return gmailThreadListPage{Threads: []struct { + ID string `json:"id"` + }{{ID: "thread-1"}}}, nil + } + base.getThread = func(ctx context.Context, userEmail, threadID string) (gmailThread, error) { + return gmailThread{ + ID: threadID, + Messages: []gmailMessage{{ + ID: "msg-1", + Payload: gmailPayload{ + Headers: []gmailHeader{ + {Name: "From", Value: "Alice Example "}, + {Name: "To", Value: "Bob "}, + {Name: "Subject", Value: "Hello/World"}, + {Name: "Date", Value: "Fri, 02 Jan 2026 03:04:05 +0000"}, + }, + Parts: []gmailPart{{ + MimeType: "text/plain", + Body: gmailBody{Data: base64.RawURLEncoding.EncodeToString([]byte("Body text"))}, + }}, + }, + }}, + }, nil + } + return connector +} + +// TestGmailTimeRangeQuery verifies Python-compatible after and before seconds. +func TestGmailTimeRangeQuery(t *testing.T) { + start := time.Unix(100, 0).UTC() + end := time.Unix(200, 0).UTC() + if got := gmailTimeRangeQuery(&start, end); got != "after:101 before:200" { + t.Fatalf("query = %q", got) + } +} diff --git a/internal/syncer/prune_runner.go b/internal/syncer/prune_runner.go index c8435fc3ba..e99daa19fc 100644 --- a/internal/syncer/prune_runner.go +++ b/internal/syncer/prune_runner.go @@ -40,6 +40,7 @@ func (r *PruneRunner) Run(ctx context.Context, taskContext service.SyncTaskConte if r.pruneService == nil { return errors.New("prune service is not configured") } + // Get the run session session, err := connector.OpenPrune(ctx, syncerconnector.PruneRequest{ TaskID: taskContext.Task.ID, ConnectorID: taskContext.Connector.ID, @@ -50,30 +51,22 @@ func (r *PruneRunner) Run(ctx context.Context, taskContext service.SyncTaskConte } defer session.Close() - if err = r.pruneService.ClearSnapshot(ctx, taskContext.Task.ID); err != nil { - return err - } - clearSnapshot := func() { - _ = r.pruneService.ClearSnapshot(context.WithoutCancel(ctx), taskContext.Task.ID) - } + retain := map[string]struct{}{} for { batch, nextErr := session.NextBatch(ctx) if errors.Is(nextErr, io.EOF) { break } if nextErr != nil { - clearSnapshot() return nextErr } - if err = r.pruneService.AddSnapshotBatch(ctx, taskContext.Task.ID, batch.Documents); err != nil { - clearSnapshot() - return err + for _, doc := range batch.Documents { + service.AddRetainDocumentID(retain, taskContext.Knowledgebase.ID, taskContext.Connector.ID, doc.SourceID) } } - removed, err := r.pruneService.DeleteStale(ctx, taskContext) + removed, err := r.pruneService.DeleteStale(ctx, taskContext, retain) if err != nil { - clearSnapshot() return err } return r.taskService.CompletePrune(ctx, taskContext, removed) diff --git a/internal/syncer/syncer.go b/internal/syncer/syncer.go index 23cf67f88d..53afa04a20 100644 --- a/internal/syncer/syncer.go +++ b/internal/syncer/syncer.go @@ -22,11 +22,14 @@ import ( "sync" "time" + "ragflow/internal/common" "ragflow/internal/dao" "ragflow/internal/service" documentservice "ragflow/internal/service/document" syncerconnector "ragflow/internal/syncer/connector" "ragflow/internal/utility" + + "go.uber.org/zap" ) // Syncer owns the scheduler, task queue, and bounded worker pool. @@ -55,7 +58,7 @@ func NewSyncer(maxConcurrency int, pollInterval time.Duration) *Syncer { registerBuiltInConnectors(registry) documentService := documentservice.NewDocumentService() - pruneService := service.NewSyncPruneService(dao.NewSyncPruneSnapshotDAO(taskDAO.DB()), documentService, nil) + pruneService := service.NewSyncPruneService(documentService, nil) return New(config, taskDAO, registry, documentService, pruneService) } @@ -134,6 +137,22 @@ func (s *Syncer) Stop() { }) } +// logTemporarySyncTaskDuration records sync task wall time while concurrency tuning is in progress. +func logTemporarySyncTaskDuration(taskContext service.SyncTaskContext, startedAt time.Time) { + if taskContext.Task.TaskType != service.TaskTypeSync { + return + } + common.Info( + "sync task duration", + zap.String("temporary_code", "remove_after_sync_concurrency_optimization"), + zap.String("task_id", taskContext.Task.ID), + zap.String("connector_id", taskContext.Connector.ID), + zap.String("kb_id", taskContext.Knowledgebase.ID), + zap.String("source", taskContext.Connector.Source), + zap.Duration("elapsed", time.Since(startedAt)), + ) +} + // registerBuiltInConnectors registers datasource connectors available in the server binary. func registerBuiltInConnectors(registry *syncerconnector.Registry) { registry.Register("rss", func(ctx context.Context, taskContext any) (syncerconnector.Connector, error) { @@ -150,4 +169,11 @@ func registerBuiltInConnectors(registry *syncerconnector.Registry) { } return syncerconnector.NewGitHubConnector(map[string]any(row.Connector.Config)) }) + registry.Register("gmail", func(ctx context.Context, taskContext any) (syncerconnector.Connector, error) { + row, ok := taskContext.(dao.SyncTaskContext) + if !ok { + return nil, errors.New("gmail connector received an invalid task context") + } + return syncerconnector.NewGmailConnector(map[string]any(row.Connector.Config)) + }) } diff --git a/internal/syncer/syncer_test.go b/internal/syncer/syncer_test.go index 9c24e43d31..8e4077274e 100644 --- a/internal/syncer/syncer_test.go +++ b/internal/syncer/syncer_test.go @@ -130,7 +130,7 @@ func setupSyncerDB(t *testing.T) *gorm.DB { if err != nil { t.Fatalf("open sqlite: %v", err) } - if err = db.AutoMigrate(&entity.Connector{}, &entity.Connector2Kb{}, &entity.Knowledgebase{}, &entity.SyncLogs{}, &entity.Document{}, &dao.SyncPruneSnapshot{}); err != nil { + if err = db.AutoMigrate(&entity.Connector{}, &entity.Connector2Kb{}, &entity.Knowledgebase{}, &entity.SyncLogs{}, &entity.Document{}); err != nil { t.Fatalf("migrate sqlite: %v", err) } orig := dao.DB @@ -455,14 +455,14 @@ func TestRecoverStaleRunningTasks(t *testing.T) { } } -// TestPruneSnapshotFailureDoesNotDelete verifies incomplete snapshots never delete. -func TestPruneSnapshotFailureDoesNotDelete(t *testing.T) { +// TestPruneSourceFailureDoesNotDelete verifies incomplete source listings never delete. +func TestPruneSourceFailureDoesNotDelete(t *testing.T) { db := setupSyncerDB(t) insertTaskContext(t, db, "conn-1", "kb-1", "task-1", dao.TaskTypePrune) _ = db.Model(&entity.SyncLogs{}).Where("id = ?", "task-1").Update("status", dao.SyncStatusRunning).Error taskService := service.NewSyncTaskService(dao.NewSyncTaskDAO(db)) deleter := &fakeDeleter{} - pruneService := service.NewSyncPruneService(dao.NewSyncPruneSnapshotDAO(db), deleter, fakeStore{ids: map[string]struct{}{"stale": {}}}) + pruneService := service.NewSyncPruneService(deleter, fakeStore{ids: map[string]struct{}{"stale": {}}}) connector := &connectormock.Connector{PruneErrAt: 1, PruneBatches: []syncerconnector.PruneBatch{{Documents: []syncerconnector.SlimDocument{{SourceID: "keep"}}}}} queue := make(chan TaskEnvelope, 1) worker := NewTaskWorker(queue, taskService, newCoordinator(taskService, newTestRegistry(map[string]*connectormock.Connector{"conn-1": connector}), &fakeSink{}, pruneService, fakeStore{}, 1), NewConnectorLock()) @@ -470,13 +470,6 @@ func TestPruneSnapshotFailureDoesNotDelete(t *testing.T) { if len(deleter.deleted) != 0 { t.Fatalf("deleted %v despite incomplete snapshot", deleter.deleted) } - var count int64 - if err := db.Model(&dao.SyncPruneSnapshot{}).Where("task_id = ?", "task-1").Count(&count).Error; err != nil { - t.Fatalf("count snapshot: %v", err) - } - if count != 0 { - t.Fatalf("snapshot rows = %d, want cleared", count) - } } // TestPruneRetainIDsIncludeLegacyAndNew verifies PRUNE retains both ID schemes. diff --git a/internal/syncer/task_worker.go b/internal/syncer/task_worker.go index 230ee9bbd3..34b3931e77 100644 --- a/internal/syncer/task_worker.go +++ b/internal/syncer/task_worker.go @@ -21,6 +21,7 @@ import ( "fmt" "ragflow/internal/service" "sync" + "time" ) // TaskWorker consumes claimed task envelopes and runs coordinators. @@ -79,11 +80,15 @@ func (w *TaskWorker) handle(ctx context.Context, envelope TaskEnvelope) { } defer w.locker.Unlock(taskContext.Connector.ID) + startedAt := time.Now() if err = w.coordinator.Execute(ctx, taskContext); err != nil { + logTemporarySyncTaskDuration(taskContext, startedAt) if ctx.Err() != nil { _ = w.taskService.RescheduleClaimed(context.WithoutCancel(ctx), taskContext.Task.ID) return } _ = w.taskService.Fail(ctx, taskContext.Task.ID, taskContext.Connector.ID, fmt.Errorf("sync task failed: %w", err)) + return } + logTemporarySyncTaskDuration(taskContext, startedAt) }