WeKnora deployment source with the reviewed wiki identity, summary and recovery fixes

internal/application/repository/retriever/doris/streamload.go

@@ -0,0 +1,508 @@
+package doris
+
+import (
+	"bytes"
+	"context"
+	"encoding/base64"
+	"encoding/json"
+	"fmt"
+	"io"
+	"net/http"
+	"net/url"
+	"strings"
+
+	"github.com/Tencent/WeKnora/internal/logger"
+	secutils "github.com/Tencent/WeKnora/internal/utils"
+)
+
+// Stream Load 相关常量。
+const (
+	// 单批 Stream Load 的 JSON body 上限（保守值，远小于 Doris 默认 streaming_load_max_mb=10240）。
+	// 主要目的是控制单次 HTTP 请求的尾延迟，超过则自动拆批。
+	streamLoadMaxBatchBytes = 1 << 20 // 1 MiB
+
+	// HTTP 头 Authorization 用 Basic auth；Stream Load 也支持 token，
+	// 这里走最常用的用户名/密码方案与 MySQL 协议保持一致。
+	headerAuthorization = "Authorization"
+	headerExpect        = "Expect"
+	headerContentType   = "Content-Type"
+)
+
+// streamLoadResponse 是 Doris FE/BE 返回的 Stream Load 结果体。
+//
+// 关键字段：Status 应为 "Success" 或 "Publish Timeout"（后者表示数据已写入但发布事务超时，
+// 仍视为成功）。其它状态都视为失败。
+type streamLoadResponse struct {
+	TxnId                  int64  `json:"TxnId"`
+	Label                  string `json:"Label"`
+	Status                 string `json:"Status"`
+	Message                string `json:"Message"`
+	NumberTotalRows        int64  `json:"NumberTotalRows"`
+	NumberLoadedRows       int64  `json:"NumberLoadedRows"`
+	NumberFilteredRows     int64  `json:"NumberFilteredRows"`
+	NumberUnselectedRows   int64  `json:"NumberUnselectedRows"`
+	LoadBytes              int64  `json:"LoadBytes"`
+	LoadTimeMs             int64  `json:"LoadTimeMs"`
+	BeginTxnTimeMs         int64  `json:"BeginTxnTimeMs"`
+	StreamLoadPutTimeMs    int64  `json:"StreamLoadPutTimeMs"`
+	ReadDataTimeMs         int64  `json:"ReadDataTimeMs"`
+	WriteDataTimeMs        int64  `json:"WriteDataTimeMs"`
+	CommitAndPublishTimeMs int64  `json:"CommitAndPublishTimeMs"`
+	ErrorURL               string `json:"ErrorURL"`
+}
+
+// streamLoadURL 拼装某张表的 Stream Load HTTP 端点。
+func (r *dorisRepository) streamLoadURL(table string) string {
+	return fmt.Sprintf("%s/api/%s/%s/_stream_load",
+		r.feHTTPBase, url.PathEscape(r.database), url.PathEscape(table))
+}
+
+// newDorisStreamLoadHTTPClient protects both the initial FE request and every
+// FE -> BE redirect at connection time. Doris requires Basic auth to survive
+// the 307 redirect, so redirect targets on another host must be explicitly
+// trusted through the SSRF whitelist before credentials are forwarded.
+func newDorisStreamLoadHTTPClient() *http.Client {
+	cfg := secutils.DefaultSSRFSafeHTTPClientConfig()
+	// Stream Load calls carry their own context deadline. Preserve the previous
+	// client behaviour instead of imposing the generic 30-second HTTP timeout.
+	cfg.Timeout = 0
+
+	client := secutils.NewSSRFSafeHTTPClient(cfg)
+	ssrfCheckRedirect := client.CheckRedirect
+	client.CheckRedirect = func(req *http.Request, via []*http.Request) error {
+		// Keep the shared client's redirect limit, scheme/target validation and
+		// sensitive-header stripping as the authoritative SSRF behaviour.
+		if err := ssrfCheckRedirect(req, via); err != nil {
+			return err
+		}
+		if len(via) == 0 {
+			return nil
+		}
+
+		sourceHost := via[len(via)-1].URL.Hostname()
+		targetHost := req.URL.Hostname()
+		if !strings.EqualFold(sourceHost, targetHost) && !secutils.IsSSRFWhitelisted(targetHost) {
+			return fmt.Errorf(
+				"stream load redirect blocked: target host %q is not trusted to receive credentials",
+				targetHost,
+			)
+		}
+
+		// The shared client deliberately strips Authorization on cross-host
+		// redirects. Doris is the exceptional case where an explicitly trusted
+		// BE needs the same Basic credentials after the FE's 307 response.
+		req.Header.Set(headerAuthorization, via[0].Header.Get(headerAuthorization))
+		return nil
+	}
+	return client
+}
+
+// partialUpdateRows 把若干行通过 Stream Load 的 partial update 模式写回目标表。
+//
+// columns 是参与本次 partial update 的列（必须包含 UNIQUE KEY 列，即 "id"）。
+// rows 中每一项是一个与 columns 等长的字段值数组。
+//
+// 实现要点：
+//  1. 用 JSON 数组的 body 形式，header 加 strip_outer_array=true。
+//  2. 设置 partial_columns=true、merge_type=APPEND，触发 Doris 的 partial update 模式
+//     (Doris 4.1 + UNIQUE KEY MoW 表的标准玩法)。
+//  3. 按 streamLoadMaxBatchBytes 自动拆批，避免单次过大。
+//  4. 处理 307：Doris 的 FE 会 redirect 到 BE，net/http 默认会跟随；
+//     这里需要保证 GetBody 可重发（已通过 bytes.NewReader 构造 Body 满足）。
+func (r *dorisRepository) partialUpdateRows(ctx context.Context,
+	table string, columns []string, rows []map[string]any,
+) error {
+	if len(rows) == 0 {
+		return nil
+	}
+	for _, batch := range chunkRows(rows, streamLoadMaxBatchBytes) {
+		if err := r.streamLoadOnce(ctx, table, columns, batch); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+// streamLoadOnce 发出一次 Stream Load HTTP 请求。
+func (r *dorisRepository) streamLoadOnce(ctx context.Context,
+	table string, columns []string, rows []map[string]any,
+) error {
+	log := logger.GetLogger(ctx)
+	if len(rows) == 0 {
+		return nil
+	}
+
+	body, err := json.Marshal(rows)
+	if err != nil {
+		return fmt.Errorf("marshal stream load body: %w", err)
+	}
+
+	url := r.streamLoadURL(table)
+	// Validate at the final outbound boundary as well as when user-supplied
+	// vector-store configuration is created. This covers stored/env configs and
+	// keeps the tainted value from reaching http.Client.Do unchecked.
+	if err := secutils.ValidateURLForSSRF(url); err != nil {
+		return fmt.Errorf("stream load URL blocked by SSRF validation: %w", err)
+	}
+	req, err := http.NewRequestWithContext(ctx, http.MethodPut, url, bytes.NewReader(body))
+	if err != nil {
+		return fmt.Errorf("build stream load request: %w", err)
+	}
+
+	// GetBody 让 redirect 时可以重新读 body（FE -> BE 的 307 需要重发 PUT body）。
+	bodyCopy := append([]byte(nil), body...)
+	req.GetBody = func() (io.ReadCloser, error) {
+		return io.NopCloser(bytes.NewReader(bodyCopy)), nil
+	}
+	req.ContentLength = int64(len(body))
+
+	auth := base64.StdEncoding.EncodeToString([]byte(r.username + ":" + r.password))
+	req.Header.Set(headerAuthorization, "Basic "+auth)
+	req.Header.Set(headerExpect, "100-continue")
+	req.Header.Set(headerContentType, "application/json")
+	req.Header.Set("format", "json")
+	req.Header.Set("strip_outer_array", "true")
+	req.Header.Set("partial_columns", "true")
+	req.Header.Set("columns", strings.Join(columns, ","))
+	req.Header.Set("merge_type", "APPEND")
+
+	resp, err := r.httpClient.Do(req)
+	if err != nil {
+		return fmt.Errorf("stream load HTTP: %w", err)
+	}
+	defer resp.Body.Close()
+
+	respBody, err := io.ReadAll(resp.Body)
+	if err != nil {
+		return fmt.Errorf("read stream load response: %w", err)
+	}
+	if resp.StatusCode/100 != 2 {
+		return fmt.Errorf("stream load HTTP %d: %s", resp.StatusCode, string(respBody))
+	}
+
+	var result streamLoadResponse
+	if err := json.Unmarshal(respBody, &result); err != nil {
+		return fmt.Errorf("decode stream load response: %w (raw=%s)", err, string(respBody))
+	}
+
+	switch result.Status {
+	case "Success", "Publish Timeout":
+		log.Infof("[Doris] Stream load %s OK: rows=%d, loaded=%d, label=%s",
+			table, result.NumberTotalRows, result.NumberLoadedRows, result.Label)
+		return nil
+	default:
+		return fmt.Errorf("stream load failed: status=%s msg=%s err_url=%s",
+			result.Status, result.Message, result.ErrorURL)
+	}
+}
+
+// chunkRows 把行按累积 JSON 体大小切分，每段不超过 maxBytes。
+//
+// 注意：JSON 序列化的实际开销大约是 marshal 后的字节数，而单行 marshal
+// 加上逗号 + 数组括号约等于本估算。这里用粗略估计避免每段都 marshal。
+func chunkRows(rows []map[string]any, maxBytes int) [][]map[string]any {
+	if len(rows) == 0 {
+		return nil
+	}
+
+	var (
+		out    [][]map[string]any
+		curr   []map[string]any
+		size   int
+		header = 2 // "[" + "]"
+	)
+
+	for _, row := range rows {
+		raw, err := json.Marshal(row)
+		if err != nil {
+			// marshal 失败时把这一行单独成段，由上层 streamLoadOnce 再次 marshal 报错。
+			if len(curr) > 0 {
+				out = append(out, curr)
+			}
+			out = append(out, []map[string]any{row})
+			curr = nil
+			size = 0
+			continue
+		}
+		// 加上逗号位（除第一行之外）。
+		need := len(raw)
+		if len(curr) > 0 {
+			need++
+		}
+		if size+need+header > maxBytes && len(curr) > 0 {
+			out = append(out, curr)
+			curr = nil
+			size = 0
+		}
+		curr = append(curr, row)
+		size += need
+	}
+	if len(curr) > 0 {
+		out = append(out, curr)
+	}
+	return out
+}
+
+// ---------------------------------------------------------------------------
+// 业务面方法：BatchUpdateChunkEnabledStatus / BatchUpdateChunkTagID
+// ---------------------------------------------------------------------------
+
+// BatchUpdateChunkEnabledStatus 批量更新 chunk 的 is_enabled 字段。
+// legacy 模式走 Stream Load partial update；inner_product_duplicate 模式改为读整行后 replaceRows 写回。
+func (r *dorisRepository) BatchUpdateChunkEnabledStatus(ctx context.Context,
+	chunkStatusMap map[string]bool,
+) error {
+	if len(chunkStatusMap) == 0 {
+		return nil
+	}
+	compatMode, err := r.resolveCompatMode(ctx)
+	if err != nil {
+		return err
+	}
+	if !compatMode.usesRewriteChunkUpdates() {
+		return r.batchUpdateChunkEnabledStatusLegacy(ctx, chunkStatusMap)
+	}
+
+	chunkIDs := make([]string, 0, len(chunkStatusMap))
+	for id := range chunkStatusMap {
+		chunkIDs = append(chunkIDs, id)
+	}
+
+	return r.rewriteChunkRows(ctx, chunkIDs, func(row *DorisVectorEmbedding) bool {
+		enabled, ok := chunkStatusMap[row.ChunkID]
+		if !ok || row.IsEnabled == enabled {
+			return false
+		}
+		row.IsEnabled = enabled
+		return true
+	}, "rewrite is_enabled")
+}
+
+// BatchUpdateChunkTagID 批量更新 chunk 的 tag_id 字段。逻辑与 EnabledStatus 一致。
+func (r *dorisRepository) BatchUpdateChunkTagID(ctx context.Context,
+	chunkTagMap map[string]string,
+) error {
+	if len(chunkTagMap) == 0 {
+		return nil
+	}
+	compatMode, err := r.resolveCompatMode(ctx)
+	if err != nil {
+		return err
+	}
+	if !compatMode.usesRewriteChunkUpdates() {
+		return r.batchUpdateChunkTagIDLegacy(ctx, chunkTagMap)
+	}
+
+	chunkIDs := make([]string, 0, len(chunkTagMap))
+	for id := range chunkTagMap {
+		chunkIDs = append(chunkIDs, id)
+	}
+
+	return r.rewriteChunkRows(ctx, chunkIDs, func(row *DorisVectorEmbedding) bool {
+		tagID, ok := chunkTagMap[row.ChunkID]
+		if !ok || row.TagID == tagID {
+			return false
+		}
+		row.TagID = tagID
+		return true
+	}, "rewrite tag_id")
+}
+
+func (r *dorisRepository) rewriteChunkRows(ctx context.Context,
+	chunkIDs []string,
+	mutate func(*DorisVectorEmbedding) bool,
+	action string,
+) error {
+	if len(chunkIDs) == 0 {
+		return nil
+	}
+
+	tables, err := r.listEmbeddingTables(ctx)
+	if err != nil {
+		return fmt.Errorf("list tables: %w", err)
+	}
+
+	for _, table := range tables {
+		rows, err := r.loadRowsByChunkIDs(ctx, table, chunkIDs)
+		if err != nil {
+			return fmt.Errorf("load chunk rows from %s: %w", table, err)
+		}
+
+		updated := make([]*DorisVectorEmbedding, 0, len(rows))
+		for _, row := range rows {
+			if !mutate(row) {
+				continue
+			}
+			updated = append(updated, row)
+		}
+		if len(updated) == 0 {
+			continue
+		}
+
+		if err := r.replaceRows(ctx, table, updated); err != nil {
+			return fmt.Errorf("%s in %s: %w", action, table, err)
+		}
+	}
+	return nil
+}
+
+func (r *dorisRepository) batchUpdateChunkEnabledStatusLegacy(ctx context.Context,
+	chunkStatusMap map[string]bool,
+) error {
+	chunkIDs := make([]string, 0, len(chunkStatusMap))
+	for id := range chunkStatusMap {
+		chunkIDs = append(chunkIDs, id)
+	}
+
+	mapping, err := r.lookupChunkRowKeys(ctx, chunkIDs)
+	if err != nil {
+		return err
+	}
+
+	byTable := make(map[string][]map[string]any)
+	for chunkID, locations := range mapping {
+		enabled, ok := chunkStatusMap[chunkID]
+		if !ok {
+			continue
+		}
+		for _, loc := range locations {
+			byTable[loc.table] = append(byTable[loc.table], map[string]any{
+				fieldID:        loc.id,
+				fieldIsEnabled: enabled,
+			})
+		}
+	}
+	for table, rows := range byTable {
+		if err := r.partialUpdateRows(ctx, table, []string{fieldID, fieldIsEnabled}, rows); err != nil {
+			return fmt.Errorf("partial update is_enabled in %s: %w", table, err)
+		}
+	}
+	return nil
+}
+
+func (r *dorisRepository) batchUpdateChunkTagIDLegacy(ctx context.Context,
+	chunkTagMap map[string]string,
+) error {
+	chunkIDs := make([]string, 0, len(chunkTagMap))
+	for id := range chunkTagMap {
+		chunkIDs = append(chunkIDs, id)
+	}
+
+	mapping, err := r.lookupChunkRowKeys(ctx, chunkIDs)
+	if err != nil {
+		return err
+	}
+
+	byTable := make(map[string][]map[string]any)
+	for chunkID, locations := range mapping {
+		tagID, ok := chunkTagMap[chunkID]
+		if !ok {
+			continue
+		}
+		for _, loc := range locations {
+			byTable[loc.table] = append(byTable[loc.table], map[string]any{
+				fieldID:    loc.id,
+				fieldTagID: tagID,
+			})
+		}
+	}
+	for table, rows := range byTable {
+		if err := r.partialUpdateRows(ctx, table, []string{fieldID, fieldTagID}, rows); err != nil {
+			return fmt.Errorf("partial update tag_id in %s: %w", table, err)
+		}
+	}
+	return nil
+}
+
+func (r *dorisRepository) loadRowsByChunkIDs(ctx context.Context,
+	table string, chunkIDs []string,
+) ([]*DorisVectorEmbedding, error) {
+	if len(chunkIDs) == 0 {
+		return nil, nil
+	}
+
+	placeholders := make([]string, len(chunkIDs))
+	args := make([]any, len(chunkIDs))
+	for i, v := range chunkIDs {
+		placeholders[i] = "?"
+		args[i] = v
+	}
+
+	stmt := fmt.Sprintf(
+		"SELECT %s FROM `%s` WHERE %s IN (%s)",
+		strings.Join(columnsForCopy, ", "),
+		table,
+		fieldChunkID,
+		strings.Join(placeholders, ", "),
+	)
+	rows, err := r.db.QueryContext(ctx, stmt, args...)
+	if err != nil {
+		return nil, err
+	}
+	batch, err := scanCopyRows(rows)
+	_ = rows.Close()
+	if err != nil {
+		return nil, err
+	}
+	return batch, nil
+}
+
+// rowLocation 表示某行在哪个表里、主键 id 是什么。
+type rowLocation struct {
+	table string
+	id    string
+}
+
+// lookupChunkRowKeys 查询给定的 chunkIDs 在所有 <base>_<dim> 表中的物理位置：
+//   - key：chunk_id
+//   - value：[(table, id), ...]，因为同一 chunk 可能在多个维度的表里都有副本。
+//
+// 跨表查询使用 listEmbeddingTables 列出的所有匹配表；每张表执行一次
+// SELECT id, chunk_id FROM <table> WHERE chunk_id IN (?, ?, ...)。
+func (r *dorisRepository) lookupChunkRowKeys(ctx context.Context,
+	chunkIDs []string,
+) (map[string][]rowLocation, error) {
+	if len(chunkIDs) == 0 {
+		return nil, nil
+	}
+	tables, err := r.listEmbeddingTables(ctx)
+	if err != nil {
+		return nil, fmt.Errorf("list tables: %w", err)
+	}
+	if len(tables) == 0 {
+		return nil, nil
+	}
+
+	placeholders := make([]string, len(chunkIDs))
+	args := make([]any, len(chunkIDs))
+	for i, v := range chunkIDs {
+		placeholders[i] = "?"
+		args[i] = v
+	}
+
+	out := make(map[string][]rowLocation)
+	for _, table := range tables {
+		stmt := fmt.Sprintf(
+			"SELECT %s, %s FROM `%s` WHERE %s IN (%s)",
+			fieldID, fieldChunkID, table, fieldChunkID, strings.Join(placeholders, ", "),
+		)
+		rows, err := r.db.QueryContext(ctx, stmt, args...)
+		if err != nil {
+			return nil, fmt.Errorf("lookup chunk row keys in %s: %w", table, err)
+		}
+		for rows.Next() {
+			var id, chunkID string
+			if err := rows.Scan(&id, &chunkID); err != nil {
+				_ = rows.Close()
+				return nil, fmt.Errorf("scan row keys: %w", err)
+			}
+			out[chunkID] = append(out[chunkID], rowLocation{table: table, id: id})
+		}
+		if err := rows.Err(); err != nil {
+			_ = rows.Close()
+			return nil, err
+		}
+		_ = rows.Close()
+	}
+	return out, nil
+}