358 lines
14 KiB
Go
358 lines
14 KiB
Go
package repository
|
|
|
|
import (
|
|
"context"
|
|
"database/sql"
|
|
"errors"
|
|
"fmt"
|
|
"strings"
|
|
"time"
|
|
|
|
"billing/internal/collection"
|
|
)
|
|
|
|
const requestDetailProjectionMigration = "request-detail-index-v1"
|
|
|
|
func (r *SQLiteUsageRepository) ensureRequestDetailProjection(ctx context.Context) error {
|
|
var completed string
|
|
err := r.db.QueryRowContext(ctx, `SELECT completed_at FROM cpa_ext_migrations WHERE name=?`, requestDetailProjectionMigration).Scan(&completed)
|
|
if err == nil {
|
|
return nil
|
|
}
|
|
if !errors.Is(err, sql.ErrNoRows) {
|
|
return fmt.Errorf("检查请求明细索引迁移: %w", err)
|
|
}
|
|
|
|
tx, err := r.db.BeginTx(ctx, nil)
|
|
if err != nil {
|
|
return fmt.Errorf("开始请求明细索引迁移: %w", err)
|
|
}
|
|
defer func() { _ = tx.Rollback() }()
|
|
if _, err := tx.ExecContext(ctx, `DELETE FROM request_detail_index`); err != nil {
|
|
return fmt.Errorf("清理请求明细索引: %w", err)
|
|
}
|
|
if _, err := tx.ExecContext(ctx, usageProjectionInsertSQL+`
|
|
SELECT u.id, CASE WHEN r.request_id IS NULL THEN '' ELSE r.request_id END,
|
|
u.requested_at, COALESCE(NULLIF(u.managed_key_id, ''), r.managed_key_id, ''),
|
|
u.model,
|
|
CASE WHEN u.failed=1 THEN 'failed'
|
|
WHEN r.outcome IN ('failed','rejected','canceled') THEN r.outcome
|
|
ELSE 'succeeded' END,
|
|
u.auth_id, `+endpointKindSQL("COALESCE(NULLIF(u.endpoint, ''), r.endpoint, '')")+`,
|
|
COALESCE(NULLIF(u.request_id, ''), r.request_id, '')
|
|
FROM usage_records u
|
|
LEFT JOIN request_records r ON u.request_id <> '' AND r.request_id=u.request_id`); err != nil {
|
|
return fmt.Errorf("回填用量请求明细索引: %w", err)
|
|
}
|
|
if _, err := tx.ExecContext(ctx, lifecycleProjectionInsertSQL+`
|
|
SELECT NULL, r.request_id, r.requested_at, r.managed_key_id,
|
|
COALESCE(NULLIF(r.model, ''), r.requested_model),
|
|
CASE WHEN r.outcome IN ('failed','rejected','canceled') THEN r.outcome ELSE 'succeeded' END,
|
|
'', `+endpointKindSQL("r.endpoint")+`, r.request_id
|
|
FROM request_records r
|
|
WHERE NOT EXISTS (SELECT 1 FROM usage_records u WHERE u.request_id=r.request_id)`); err != nil {
|
|
return fmt.Errorf("回填终态请求明细索引: %w", err)
|
|
}
|
|
if err := tx.Commit(); err != nil {
|
|
return fmt.Errorf("提交请求明细索引迁移: %w", err)
|
|
}
|
|
|
|
rows, err := r.db.QueryContext(ctx, `SELECT requested_at FROM usage_records WHERE request_id='' ORDER BY requested_at`)
|
|
if err != nil {
|
|
return fmt.Errorf("扫描孤立用量时间: %w", err)
|
|
}
|
|
var orphanTimes []time.Time
|
|
for rows.Next() {
|
|
var raw string
|
|
if err := rows.Scan(&raw); err != nil {
|
|
_ = rows.Close()
|
|
return fmt.Errorf("读取孤立用量时间: %w", err)
|
|
}
|
|
parsed, parseErr := time.Parse(time.RFC3339Nano, raw)
|
|
if parseErr != nil {
|
|
_ = rows.Close()
|
|
return fmt.Errorf("解析孤立用量时间: %w", parseErr)
|
|
}
|
|
orphanTimes = append(orphanTimes, parsed)
|
|
}
|
|
if err := rows.Close(); err != nil {
|
|
return fmt.Errorf("关闭孤立用量扫描: %w", err)
|
|
}
|
|
var last time.Time
|
|
for _, candidate := range orphanTimes {
|
|
if !last.IsZero() && candidate.Sub(last) < orphanLifecycleMatchWindow {
|
|
continue
|
|
}
|
|
if err := r.reconcileOrphanProjection(ctx, candidate); err != nil {
|
|
return err
|
|
}
|
|
last = candidate
|
|
}
|
|
if _, err := r.db.ExecContext(ctx, `INSERT OR REPLACE INTO cpa_ext_migrations(name, completed_at) VALUES(?, ?)`, requestDetailProjectionMigration, formatTime(time.Now().UTC())); err != nil {
|
|
return fmt.Errorf("完成请求明细索引迁移: %w", err)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
const usageProjectionInsertSQL = `
|
|
INSERT OR IGNORE INTO request_detail_index
|
|
(usage_id, lifecycle_request_id, requested_at, managed_key_id, model, result, auth_id, endpoint_kind, request_id) `
|
|
|
|
const lifecycleProjectionInsertSQL = `
|
|
INSERT OR IGNORE INTO request_detail_index
|
|
(usage_id, lifecycle_request_id, requested_at, managed_key_id, model, result, auth_id, endpoint_kind, request_id) `
|
|
|
|
func endpointKindSQL(expression string) string {
|
|
return `CASE
|
|
WHEN LOWER(TRIM(` + expression + `)) LIKE '%/responses/compact' THEN 'compact'
|
|
WHEN LOWER(TRIM(` + expression + `)) LIKE '%/chat/completions' THEN 'chat'
|
|
WHEN LOWER(TRIM(` + expression + `)) LIKE '%/responses' THEN 'responses'
|
|
ELSE LOWER(TRIM(` + expression + `)) END`
|
|
}
|
|
|
|
func endpointKind(value string) string {
|
|
path := strings.ToLower(strings.TrimSpace(value))
|
|
if fields := strings.Fields(path); len(fields) > 1 {
|
|
path = fields[len(fields)-1]
|
|
}
|
|
if index := strings.IndexByte(path, '?'); index >= 0 {
|
|
path = path[:index]
|
|
}
|
|
path = strings.TrimRight(path, "/")
|
|
switch {
|
|
case strings.HasSuffix(path, "/responses/compact"):
|
|
return "compact"
|
|
case strings.HasSuffix(path, "/chat/completions"):
|
|
return "chat"
|
|
case strings.HasSuffix(path, "/responses"):
|
|
return "responses"
|
|
default:
|
|
return path
|
|
}
|
|
}
|
|
|
|
func usageResult(failed bool, outcome string) string {
|
|
switch strings.ToLower(strings.TrimSpace(outcome)) {
|
|
case collection.UsageResultCanceled:
|
|
return collection.UsageResultCanceled
|
|
case collection.UsageResultRejected:
|
|
return collection.UsageResultRejected
|
|
case collection.UsageResultFailed:
|
|
return collection.UsageResultFailed
|
|
}
|
|
if failed {
|
|
return collection.UsageResultFailed
|
|
}
|
|
return collection.UsageResultSucceeded
|
|
}
|
|
|
|
func insertUsageProjection(ctx context.Context, tx *sql.Tx, usageID int64) error {
|
|
_, err := tx.ExecContext(ctx, usageProjectionInsertSQL+`
|
|
SELECT u.id, CASE WHEN r.request_id IS NULL THEN '' ELSE r.request_id END,
|
|
u.requested_at, COALESCE(NULLIF(u.managed_key_id, ''), r.managed_key_id, ''),
|
|
u.model,
|
|
CASE WHEN u.failed=1 THEN 'failed'
|
|
WHEN r.outcome IN ('failed','rejected','canceled') THEN r.outcome
|
|
ELSE 'succeeded' END,
|
|
u.auth_id, `+endpointKindSQL("COALESCE(NULLIF(u.endpoint, ''), r.endpoint, '')")+`,
|
|
COALESCE(NULLIF(u.request_id, ''), r.request_id, '')
|
|
FROM usage_records u
|
|
LEFT JOIN request_records r ON u.request_id <> '' AND r.request_id=u.request_id
|
|
WHERE u.id=?`, usageID)
|
|
if err != nil {
|
|
return fmt.Errorf("写入请求明细索引: %w", err)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func syncRequestProjection(ctx context.Context, tx *sql.Tx, requestID string) (bool, error) {
|
|
_, err := tx.ExecContext(ctx, `
|
|
UPDATE request_detail_index
|
|
SET lifecycle_request_id=?, request_id=?,
|
|
managed_key_id=COALESCE(NULLIF((SELECT managed_key_id FROM usage_records WHERE id=request_detail_index.usage_id), ''),
|
|
(SELECT managed_key_id FROM request_records WHERE request_id=?), ''),
|
|
result=CASE
|
|
WHEN (SELECT failed FROM usage_records WHERE id=request_detail_index.usage_id)=1 THEN 'failed'
|
|
WHEN (SELECT outcome FROM request_records WHERE request_id=?) IN ('failed','rejected','canceled')
|
|
THEN (SELECT outcome FROM request_records WHERE request_id=?)
|
|
ELSE 'succeeded' END
|
|
WHERE usage_id IN (SELECT id FROM usage_records WHERE request_id=?)`, requestID, requestID, requestID, requestID, requestID, requestID)
|
|
if err != nil {
|
|
return false, fmt.Errorf("关联请求终态索引: %w", err)
|
|
}
|
|
var linked bool
|
|
if err := tx.QueryRowContext(ctx, `
|
|
SELECT EXISTS (
|
|
SELECT 1 FROM request_detail_index
|
|
WHERE usage_id IN (SELECT id FROM usage_records WHERE request_id=?)
|
|
)`, requestID).Scan(&linked); err != nil {
|
|
return false, fmt.Errorf("确认请求终态索引关联: %w", err)
|
|
}
|
|
if linked {
|
|
if _, err := tx.ExecContext(ctx, `UPDATE request_detail_index SET endpoint_kind=`+endpointKindSQL(`COALESCE(NULLIF((SELECT endpoint FROM usage_records WHERE id=request_detail_index.usage_id), ''), (SELECT endpoint FROM request_records WHERE request_id=?), '')`)+` WHERE usage_id IN (SELECT id FROM usage_records WHERE request_id=?)`, requestID, requestID, requestID, requestID, requestID); err != nil {
|
|
return false, fmt.Errorf("更新请求端点索引: %w", err)
|
|
}
|
|
if _, err := tx.ExecContext(ctx, `DELETE FROM request_detail_index WHERE usage_id IS NULL AND lifecycle_request_id=?`, requestID); err != nil {
|
|
return false, fmt.Errorf("删除重复请求终态索引: %w", err)
|
|
}
|
|
return true, nil
|
|
}
|
|
_, err = tx.ExecContext(ctx, lifecycleProjectionInsertSQL+`
|
|
SELECT NULL, r.request_id, r.requested_at, r.managed_key_id,
|
|
COALESCE(NULLIF(r.model, ''), r.requested_model),
|
|
CASE WHEN r.outcome IN ('failed','rejected','canceled') THEN r.outcome ELSE 'succeeded' END,
|
|
'', `+endpointKindSQL("r.endpoint")+`, r.request_id
|
|
FROM request_records r WHERE r.request_id=?`, requestID)
|
|
if err != nil {
|
|
return false, fmt.Errorf("写入终态请求明细索引: %w", err)
|
|
}
|
|
return false, nil
|
|
}
|
|
|
|
type orphanUsageProjection struct {
|
|
ID int64
|
|
ManagedKeyID string
|
|
RequestedAt time.Time
|
|
Model string
|
|
Failed bool
|
|
AuthID string
|
|
Endpoint string
|
|
}
|
|
|
|
type orphanLifecycleProjection struct {
|
|
RequestID string
|
|
ManagedKeyID string
|
|
RequestedAt time.Time
|
|
Model string
|
|
Outcome string
|
|
Endpoint string
|
|
}
|
|
|
|
func (r *SQLiteUsageRepository) reconcileOrphanProjection(ctx context.Context, at time.Time) error {
|
|
window := time.Second
|
|
from, to := formatTime(at.Add(-window)), formatTime(at.Add(window))
|
|
usageRows, err := r.db.QueryContext(ctx, `
|
|
SELECT id, managed_key_id, requested_at, model, failed, auth_id, endpoint
|
|
FROM usage_records WHERE request_id='' AND requested_at BETWEEN ? AND ?`, from, to)
|
|
if err != nil {
|
|
return fmt.Errorf("查询孤立用量候选: %w", err)
|
|
}
|
|
var usages []orphanUsageProjection
|
|
for usageRows.Next() {
|
|
var item orphanUsageProjection
|
|
var requestedAt string
|
|
if err := usageRows.Scan(&item.ID, &item.ManagedKeyID, &requestedAt, &item.Model, &item.Failed, &item.AuthID, &item.Endpoint); err != nil {
|
|
_ = usageRows.Close()
|
|
return fmt.Errorf("读取孤立用量候选: %w", err)
|
|
}
|
|
item.RequestedAt, err = time.Parse(time.RFC3339Nano, requestedAt)
|
|
if err != nil {
|
|
_ = usageRows.Close()
|
|
return fmt.Errorf("解析孤立用量候选时间: %w", err)
|
|
}
|
|
usages = append(usages, item)
|
|
}
|
|
_ = usageRows.Close()
|
|
|
|
lifecycleRows, err := r.db.QueryContext(ctx, `
|
|
SELECT r.request_id, r.managed_key_id, r.requested_at,
|
|
COALESCE(NULLIF(r.model, ''), r.requested_model), r.outcome, r.endpoint
|
|
FROM request_records r
|
|
WHERE r.requested_at BETWEEN ? AND ?
|
|
AND NOT EXISTS (SELECT 1 FROM usage_records u WHERE u.request_id=r.request_id)`, from, to)
|
|
if err != nil {
|
|
return fmt.Errorf("查询孤立终态候选: %w", err)
|
|
}
|
|
var lifecycles []orphanLifecycleProjection
|
|
for lifecycleRows.Next() {
|
|
var item orphanLifecycleProjection
|
|
var requestedAt string
|
|
if err := lifecycleRows.Scan(&item.RequestID, &item.ManagedKeyID, &requestedAt, &item.Model, &item.Outcome, &item.Endpoint); err != nil {
|
|
_ = lifecycleRows.Close()
|
|
return fmt.Errorf("读取孤立终态候选: %w", err)
|
|
}
|
|
item.RequestedAt, err = time.Parse(time.RFC3339Nano, requestedAt)
|
|
if err != nil {
|
|
_ = lifecycleRows.Close()
|
|
return fmt.Errorf("解析孤立终态候选时间: %w", err)
|
|
}
|
|
lifecycles = append(lifecycles, item)
|
|
}
|
|
_ = lifecycleRows.Close()
|
|
|
|
tx, err := r.db.BeginTx(ctx, nil)
|
|
if err != nil {
|
|
return fmt.Errorf("开始重建孤立请求索引: %w", err)
|
|
}
|
|
defer func() { _ = tx.Rollback() }()
|
|
for _, usage := range usages {
|
|
if _, err := tx.ExecContext(ctx, `UPDATE request_detail_index SET lifecycle_request_id='', request_id='', managed_key_id=?, model=?, result=?, auth_id=?, endpoint_kind=? WHERE usage_id=?`, usage.ManagedKeyID, usage.Model, usageResult(usage.Failed, ""), usage.AuthID, endpointKind(usage.Endpoint), usage.ID); err != nil {
|
|
return fmt.Errorf("重置孤立用量索引: %w", err)
|
|
}
|
|
}
|
|
for _, lifecycle := range lifecycles {
|
|
if _, err := tx.ExecContext(ctx, lifecycleProjectionInsertSQL+`VALUES(NULL, ?, ?, ?, ?, ?, '', ?, ?)`, lifecycle.RequestID, formatTime(lifecycle.RequestedAt), lifecycle.ManagedKeyID, lifecycle.Model, usageResult(false, lifecycle.Outcome), endpointKind(lifecycle.Endpoint), lifecycle.RequestID); err != nil {
|
|
return fmt.Errorf("恢复孤立终态索引: %w", err)
|
|
}
|
|
}
|
|
|
|
usageMatches := make(map[int]lifecycleMatch)
|
|
requestMatches := make(map[int]lifecycleMatch)
|
|
for usageIndex, usage := range usages {
|
|
for requestIndex, lifecycle := range lifecycles {
|
|
if !sameOrphanProjection(usage, lifecycle) {
|
|
continue
|
|
}
|
|
distance := usage.RequestedAt.Sub(lifecycle.RequestedAt)
|
|
if distance < 0 {
|
|
distance = -distance
|
|
}
|
|
if distance > orphanLifecycleMatchWindow {
|
|
continue
|
|
}
|
|
if current, ok := usageMatches[usageIndex]; ok {
|
|
usageMatches[usageIndex] = betterLifecycleMatch(current, requestIndex, distance)
|
|
} else {
|
|
usageMatches[usageIndex] = lifecycleMatch{index: requestIndex, distance: distance}
|
|
}
|
|
if current, ok := requestMatches[requestIndex]; ok {
|
|
requestMatches[requestIndex] = betterLifecycleMatch(current, usageIndex, distance)
|
|
} else {
|
|
requestMatches[requestIndex] = lifecycleMatch{index: usageIndex, distance: distance}
|
|
}
|
|
}
|
|
}
|
|
for usageIndex, requestMatch := range usageMatches {
|
|
usageMatch, ok := requestMatches[requestMatch.index]
|
|
if requestMatch.ambiguous || !ok || usageMatch.ambiguous || usageMatch.index != usageIndex {
|
|
continue
|
|
}
|
|
usage, lifecycle := usages[usageIndex], lifecycles[requestMatch.index]
|
|
managedKeyID := usage.ManagedKeyID
|
|
if managedKeyID == "" {
|
|
managedKeyID = lifecycle.ManagedKeyID
|
|
}
|
|
endpoint := usage.Endpoint
|
|
if strings.TrimSpace(endpoint) == "" {
|
|
endpoint = lifecycle.Endpoint
|
|
}
|
|
if _, err := tx.ExecContext(ctx, `UPDATE request_detail_index SET lifecycle_request_id=?, request_id=?, managed_key_id=?, result=?, endpoint_kind=? WHERE usage_id=?`, lifecycle.RequestID, lifecycle.RequestID, managedKeyID, usageResult(usage.Failed, lifecycle.Outcome), endpointKind(endpoint), usage.ID); err != nil {
|
|
return fmt.Errorf("合并孤立请求索引: %w", err)
|
|
}
|
|
if _, err := tx.ExecContext(ctx, `DELETE FROM request_detail_index WHERE usage_id IS NULL AND lifecycle_request_id=?`, lifecycle.RequestID); err != nil {
|
|
return fmt.Errorf("删除已合并终态索引: %w", err)
|
|
}
|
|
}
|
|
if err := tx.Commit(); err != nil {
|
|
return fmt.Errorf("提交孤立请求索引: %w", err)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func sameOrphanProjection(usage orphanUsageProjection, lifecycle orphanLifecycleProjection) bool {
|
|
if usage.ManagedKeyID != "" && lifecycle.ManagedKeyID != "" {
|
|
return usage.ManagedKeyID == lifecycle.ManagedKeyID
|
|
}
|
|
return strings.TrimSpace(usage.Model) != "" && strings.EqualFold(strings.TrimSpace(usage.Model), strings.TrimSpace(lifecycle.Model))
|
|
}
|